Real per-step progress bars for local GPU inference

Backend:
- local_diffusion.py: add _make_step_cb() that writes step/total_steps/
  progress into _states on every diffusers callback_on_step_end; wired into
  txt2img, inpaint, img2img with TypeError fallback for older diffusers
- ai_tools.py: GET /api/generate/progress SSE endpoint — streams _states
  as JSON array every 200ms so clients get live denoising step counts

Frontend:
- progress_overlay.js: add connectProgressSSE(pipeType, baseUrl) /
  disconnectProgressSSE() — opens EventSource, maps step/total_steps
  to bar percentage (0→85% during denoising, 85→100 for decode/place)
- text_to_image.js: connect SSE before POST, disconnect on done/error
- selection_actions.js: connect SSE for AI edit / asymmetry operations

Result: for local GPU, progress bar shows "Step 12 / 30" with exact fill;
for remote providers and upscale (no step callbacks), shimmer animates.

https://claude.ai/code/session_01WVDg7amsy1TTtxvpku7bcM

backend/app/routers/ai_tools.py

@@ -4,10 +4,12 @@ All endpoints are under /api prefix.
 """
 
 from fastapi import APIRouter, HTTPException
+from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
 from typing import Optional
 import base64
 import asyncio
+import json
 
 from app.services.local_inpaint import (
     lama_inpaint, opencv_inpaint, lama_available, gpu_available, rembg_available,
@@ -191,6 +193,37 @@ async def inpaint_remote(req: InpaintRemoteRequest):
         raise HTTPException(status_code=500, detail=str(e))
 
 
+@router.get("/generate/progress")
+async def generation_progress_stream():
+    """
+    SSE stream of local GPU pipeline inference progress.
+    Events are JSON arrays of pipeline state objects, emitted every 200 ms.
+    Each object: {pipeline, state, step, total_steps, progress, message, model_id, …}
+    Clients open this with EventSource before firing a generation POST,
+    then close it when the POST resolves.
+    """
+    from app.services.local_diffusion import get_all_model_states
+
+    async def event_gen():
+        try:
+            while True:
+                states = get_all_model_states()
+                yield f"data: {json.dumps(states)}\n\n"
+                await asyncio.sleep(0.2)
+        except asyncio.CancelledError:
+            pass
+
+    return StreamingResponse(
+        event_gen(),
+        media_type="text/event-stream",
+        headers={
+            "Cache-Control": "no-cache",
+            "X-Accel-Buffering": "no",
+            "Connection": "keep-alive",
+        },
+    )
+
+
 @router.post("/generate/txt2img")
 async def txt2img(req: Txt2ImgRequest):
     """Text-to-image via configured remote provider."""

backend/app/services/local_diffusion.py

@@ -48,6 +48,24 @@ def get_all_model_states() -> list[dict]:
         return list(_states.values())
 
 
+def _make_step_cb(pipe_type: str, total_steps: int):
+    """
+    Returns a diffusers callback_on_step_end that writes per-step progress
+    into _states so the SSE /api/generate/progress endpoint can stream it.
+    Called from a thread executor — _set_state is thread-safe.
+    """
+    def cb(pipe, step_index: int, timestep, callback_kwargs: dict) -> dict:
+        done = step_index + 1
+        _set_state(pipe_type,
+                   state="running",
+                   step=done,
+                   total_steps=total_steps,
+                   progress=round(done / total_steps * 85, 1),
+                   message=f"Step {done} / {total_steps}")
+        return callback_kwargs
+    return cb
+
+
 # ── LRU pipeline cache ────────────────────────────────────────────────────────
 
 class _PipelineCache:
@@ -295,18 +313,35 @@ class LocalDiffusionProvider(RemoteAIProvider):
         steps = int(params.get("steps", 30))
         cfg   = float(params.get("cfg_scale", 7.5))
         neg   = params.get("negative_prompt", "") or None
+        step_cb = _make_step_cb("inpaint", steps)
+
+        _set_state("inpaint", state="running", step=0, total_steps=steps, progress=0, message="Starting…")
 
         def _run():
-            return pipe(
+            try:
-                prompt=prompt,
+                return pipe(
-                negative_prompt=neg,
+                    prompt=prompt,
-                image=img_r,
+                    negative_prompt=neg,
-                mask_image=mask_r,
+                    image=img_r,
-                num_inference_steps=steps,
+                    mask_image=mask_r,
-                guidance_scale=cfg,
+                    num_inference_steps=steps,
-            ).images[0].resize(orig, Image.LANCZOS)
+                    guidance_scale=cfg,
+                    callback_on_step_end=step_cb,
+                    callback_on_step_end_tensor_inputs=["latents"],
+                ).images[0].resize(orig, Image.LANCZOS)
+            except TypeError:
+                return pipe(
+                    prompt=prompt,
+                    negative_prompt=neg,
+                    image=img_r,
+                    mask_image=mask_r,
+                    num_inference_steps=steps,
+                    guidance_scale=cfg,
+                ).images[0].resize(orig, Image.LANCZOS)
 
-        return _to_png(await asyncio.get_event_loop().run_in_executor(None, _run))
+        result = _to_png(await asyncio.get_event_loop().run_in_executor(None, _run))
+        _set_state("inpaint", state="ready", step=None, total_steps=None, progress=100, message="Ready")
+        return result
 
     async def txt2img(self, prompt: str, width: int, height: int, params: dict) -> bytes:
         pipe = await self._get_pipeline("txt2img")
@@ -316,34 +351,61 @@ class LocalDiffusionProvider(RemoteAIProvider):
         w = min(width,  max_dim) // 8 * 8
         h = min(height, max_dim) // 8 * 8
         seed = int(params.get("seed", 0))
-
         is_flux = spec.family == "flux"
+        steps = 4 if is_flux else int(params.get("steps", 30))
+        step_cb = _make_step_cb("txt2img", steps)
+
+        _set_state("txt2img", state="running", step=0, total_steps=steps, progress=0, message="Starting…")
 
         def _run():
             import torch
             device = self._info.backend
             gen = torch.Generator(device=device).manual_seed(seed) if seed else None
 
-            if is_flux:
+            try:
-                return pipe(
+                if is_flux:
-                    prompt=prompt,
+                    return pipe(
-                    width=w, height=h,
+                        prompt=prompt,
-                    num_inference_steps=4,   # FLUX.1-schnell is a 4-step model
+                        width=w, height=h,
-                    guidance_scale=0.0,      # fully CFG-distilled
+                        num_inference_steps=steps,
-                    max_sequence_length=256,
+                        guidance_scale=0.0,
-                    generator=gen,
+                        max_sequence_length=256,
-                ).images[0]
+                        generator=gen,
-            else:
+                        callback_on_step_end=step_cb,
-                return pipe(
+                        callback_on_step_end_tensor_inputs=["latents"],
-                    prompt=prompt,
+                    ).images[0]
-                    negative_prompt=params.get("negative_prompt", "") or None,
+                else:
-                    width=w, height=h,
+                    return pipe(
-                    num_inference_steps=int(params.get("steps", 30)),
+                        prompt=prompt,
-                    guidance_scale=float(params.get("cfg_scale", 7.5)),
+                        negative_prompt=params.get("negative_prompt", "") or None,
-                    generator=gen,
+                        width=w, height=h,
-                ).images[0]
+                        num_inference_steps=steps,
+                        guidance_scale=float(params.get("cfg_scale", 7.5)),
+                        generator=gen,
+                        callback_on_step_end=step_cb,
+                        callback_on_step_end_tensor_inputs=["latents"],
+                    ).images[0]
+            except TypeError:
+                # Older diffusers without callback_on_step_end
+                if is_flux:
+                    return pipe(
+                        prompt=prompt, width=w, height=h,
+                        num_inference_steps=steps, guidance_scale=0.0,
+                        max_sequence_length=256, generator=gen,
+                    ).images[0]
+                else:
+                    return pipe(
+                        prompt=prompt,
+                        negative_prompt=params.get("negative_prompt", "") or None,
+                        width=w, height=h,
+                        num_inference_steps=steps,
+                        guidance_scale=float(params.get("cfg_scale", 7.5)),
+                        generator=gen,
+                    ).images[0]
 
-        return _to_png(await asyncio.get_event_loop().run_in_executor(None, _run))
+        result = _to_png(await asyncio.get_event_loop().run_in_executor(None, _run))
+        _set_state("txt2img", state="ready", step=None, total_steps=None, progress=100, message="Ready")
+        return result
 
     async def img2img(self, image_bytes: bytes, prompt: str, strength: float, params: dict) -> bytes:
         pipe = await self._get_pipeline("img2img")
@@ -353,28 +415,49 @@ class LocalDiffusionProvider(RemoteAIProvider):
         orig = img.size
         img_r = _resize_square(img, spec.native_res)
         is_flux = spec.family == "flux"
+        steps = 4 if is_flux else int(params.get("steps", 30))
+        step_cb = _make_step_cb("img2img", steps)
+
+        _set_state("img2img", state="running", step=0, total_steps=steps, progress=0, message="Starting…")
 
         def _run():
-            if is_flux:
+            try:
-                result = pipe(
+                if is_flux:
-                    prompt=prompt,
+                    result = pipe(
-                    image=img_r,
+                        prompt=prompt, image=img_r, strength=strength,
-                    strength=strength,
+                        num_inference_steps=steps, guidance_scale=0.0,
-                    num_inference_steps=4,
+                        callback_on_step_end=step_cb,
-                    guidance_scale=0.0,
+                        callback_on_step_end_tensor_inputs=["latents"],
-                ).images[0]
+                    ).images[0]
-            else:
+                else:
-                result = pipe(
+                    result = pipe(
-                    prompt=prompt,
+                        prompt=prompt,
-                    negative_prompt=params.get("negative_prompt", "") or None,
+                        negative_prompt=params.get("negative_prompt", "") or None,
-                    image=img_r,
+                        image=img_r, strength=strength,
-                    strength=strength,
+                        num_inference_steps=steps,
-                    num_inference_steps=int(params.get("steps", 30)),
+                        guidance_scale=float(params.get("cfg_scale", 7.5)),
-                    guidance_scale=float(params.get("cfg_scale", 7.5)),
+                        callback_on_step_end=step_cb,
-                ).images[0]
+                        callback_on_step_end_tensor_inputs=["latents"],
+                    ).images[0]
+            except TypeError:
+                if is_flux:
+                    result = pipe(
+                        prompt=prompt, image=img_r, strength=strength,
+                        num_inference_steps=steps, guidance_scale=0.0,
+                    ).images[0]
+                else:
+                    result = pipe(
+                        prompt=prompt,
+                        negative_prompt=params.get("negative_prompt", "") or None,
+                        image=img_r, strength=strength,
+                        num_inference_steps=steps,
+                        guidance_scale=float(params.get("cfg_scale", 7.5)),
+                    ).images[0]
             return result.resize(orig, Image.LANCZOS)
 
-        return _to_png(await asyncio.get_event_loop().run_in_executor(None, _run))
+        result = _to_png(await asyncio.get_event_loop().run_in_executor(None, _run))
+        _set_state("img2img", state="ready", step=None, total_steps=None, progress=100, message="Ready")
+        return result
 
     async def outpaint(self, image_bytes: bytes, direction: str, size: int, prompt: str) -> bytes:
         from PIL import ImageDraw

frontend/src/js/libs/progress_overlay.js

@@ -19,6 +19,44 @@ var _shimmerAnim = null;
 var _fakeTimer = null;
 var _currentPct = 0;
 
+// ── SSE progress connection ───────────────────────────────────────────────────
+
+var _sse = null;
+
+/**
+ * Open an EventSource to /api/generate/progress and drive the bar with real
+ * denoising step counts from the local GPU pipeline.
+ *
+ * @param {string} pipeType  - 'txt2img' | 'inpaint' | 'img2img'
+ * @param {string} baseUrl   - window.API_BASE_URL or ''
+ */
+export function connectProgressSSE(pipeType, baseUrl) {
+    disconnectProgressSSE();
+    try {
+        var url = (baseUrl || '') + '/api/generate/progress';
+        _sse = new EventSource(url);
+        _sse.onmessage = (e) => {
+            try {
+                var states = JSON.parse(e.data);
+                var s = Array.isArray(states)
+                    ? states.find(st => st.pipeline === pipeType)
+                    : null;
+                if (s && s.state === 'running' && s.total_steps) {
+                    var pct = Math.round(s.step / s.total_steps * 85);
+                    updateProgress(pct, s.message || `Step ${s.step} / ${s.total_steps}`);
+                }
+            } catch { /* malformed event — ignore */ }
+        };
+        _sse.onerror = () => disconnectProgressSSE();
+    } catch { /* SSE not supported */ }
+}
+
+export function disconnectProgressSSE() {
+    if (_sse) { _sse.close(); _sse = null; }
+}
+
+// ── Progress overlay ──────────────────────────────────────────────────────────
+
 export function showProgress(message, estimatedSeconds) {
     hideProgress();
 

frontend/src/js/modules/generate/text_to_image.js

@@ -12,7 +12,7 @@ import Dialog_class from './../../libs/popup.js';
 import alertify from './../../../../node_modules/alertifyjs/build/alertify.min.js';
 import apiService from './../../services/api.js';
 import { getCapabilities } from './../../api/capabilities.js';
-import { showProgress, updateProgress, hideProgress } from './../../libs/progress_overlay.js';
+import { showProgress, updateProgress, hideProgress, connectProgressSSE, disconnectProgressSSE } from './../../libs/progress_overlay.js';
 
 var instance = null;
 
@@ -142,7 +142,8 @@ class Generate_text_to_image_class {
         if (this.isProcessing) return;
         this.isProcessing = true;
 
-        showProgress('Generating image… this may take a minute on local GPU', estSec || 60);
+        connectProgressSSE('txt2img', window.API_BASE_URL || '');
+        showProgress('Generating image…', estSec || 60);
 
         try {
             var result = await apiService.textToImage(params.prompt, {
@@ -185,11 +186,13 @@ class Generate_text_to_image_class {
                         ])
                     );
                 }
+                disconnectProgressSSE();
                 hideProgress();
                 alertify.success('Image generated!');
                 this.isProcessing = false;
             };
             img.onerror = () => {
+                disconnectProgressSSE();
                 hideProgress();
                 alertify.error('Failed to load generated image.');
                 this.isProcessing = false;
@@ -197,6 +200,7 @@ class Generate_text_to_image_class {
             img.src = 'data:image/png;base64,' + result.result;
 
         } catch (err) {
+            disconnectProgressSSE();
             hideProgress();
             alertify.error('Generation failed: ' + (err.message || err));
             this.isProcessing = false;

frontend/src/js/tools/selection_actions.js

@@ -18,7 +18,7 @@ import app from './../app.js';
 import config from './../config.js';
 import Base_layers_class from './../core/base-layers.js';
 import alertify from './../../../node_modules/alertifyjs/build/alertify.min.js';
-import { showProgress, hideProgress } from './../libs/progress_overlay.js';
+import { showProgress, updateProgress, hideProgress, connectProgressSSE, disconnectProgressSSE } from './../libs/progress_overlay.js';
 
 const BASE = window.API_BASE_URL || '';
 
@@ -196,6 +196,7 @@ export class SelectionActions {
     async _makeAsymmetric() {
         if (!this._check()) return;
         this.hide();
+        connectProgressSSE('inpaint', window.API_BASE_URL || '');
         showProgress('AI is adding natural asymmetry…', 60);
         try {
             var res = await _post('/api/image/ai-edit-region', {
@@ -208,9 +209,11 @@ export class SelectionActions {
             });
             this.tool.updateLayerWithResult(res.result);
             this.tool.clearSelection();
+            disconnectProgressSSE();
             hideProgress();
             alertify.success('Made less symmetrical!');
         } catch (e) {
+            disconnectProgressSSE();
             hideProgress();
             alertify.error('AI edit failed: ' + e.message);
         }
@@ -219,6 +222,7 @@ export class SelectionActions {
     async _aiEditRegion(instruction) {
         if (!this._check()) return;
         this.hide();
+        connectProgressSSE('inpaint', window.API_BASE_URL || '');
         showProgress('AI is editing the region…', 60);
         try {
             var res = await _post('/api/image/ai-edit-region', {
@@ -230,9 +234,11 @@ export class SelectionActions {
             });
             this.tool.updateLayerWithResult(res.result);
             this.tool.clearSelection();
+            disconnectProgressSSE();
             hideProgress();
             alertify.success('Done!');
         } catch (e) {
+            disconnectProgressSSE();
             hideProgress();
             alertify.error('AI edit failed: ' + e.message);
         }