From 372ab48991e121bfad99ea843f41d83ecf98a19f Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 13 Jun 2026 16:48:55 +0000 Subject: [PATCH] Real per-step progress bars for local GPU inference MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 | 33 ++++ backend/app/services/local_diffusion.py | 177 +++++++++++++----- frontend/src/js/libs/progress_overlay.js | 38 ++++ .../src/js/modules/generate/text_to_image.js | 8 +- frontend/src/js/tools/selection_actions.js | 8 +- 5 files changed, 214 insertions(+), 50 deletions(-) diff --git a/backend/app/routers/ai_tools.py b/backend/app/routers/ai_tools.py index be39254..dba991e 100644 --- a/backend/app/routers/ai_tools.py +++ b/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.""" diff --git a/backend/app/services/local_diffusion.py b/backend/app/services/local_diffusion.py index b873027..be2c1a1 100644 --- a/backend/app/services/local_diffusion.py +++ b/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( - 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) + try: + return pipe( + prompt=prompt, + negative_prompt=neg, + image=img_r, + mask_image=mask_r, + num_inference_steps=steps, + 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: - return pipe( - prompt=prompt, - width=w, height=h, - num_inference_steps=4, # FLUX.1-schnell is a 4-step model - guidance_scale=0.0, # fully CFG-distilled - 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=int(params.get("steps", 30)), - guidance_scale=float(params.get("cfg_scale", 7.5)), - generator=gen, - ).images[0] + try: + 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, + callback_on_step_end=step_cb, + callback_on_step_end_tensor_inputs=["latents"], + ).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, + 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: - result = pipe( - prompt=prompt, - image=img_r, - strength=strength, - num_inference_steps=4, - 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=int(params.get("steps", 30)), - guidance_scale=float(params.get("cfg_scale", 7.5)), - ).images[0] + try: + if is_flux: + result = pipe( + prompt=prompt, image=img_r, strength=strength, + num_inference_steps=steps, guidance_scale=0.0, + callback_on_step_end=step_cb, + callback_on_step_end_tensor_inputs=["latents"], + ).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)), + callback_on_step_end=step_cb, + 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 diff --git a/frontend/src/js/libs/progress_overlay.js b/frontend/src/js/libs/progress_overlay.js index da10cb0..7792e4c 100644 --- a/frontend/src/js/libs/progress_overlay.js +++ b/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(); diff --git a/frontend/src/js/modules/generate/text_to_image.js b/frontend/src/js/modules/generate/text_to_image.js index 86f4853..82b5aa4 100644 --- a/frontend/src/js/modules/generate/text_to_image.js +++ b/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; diff --git a/frontend/src/js/tools/selection_actions.js b/frontend/src/js/tools/selection_actions.js index 8cbec02..caaa404 100644 --- a/frontend/src/js/tools/selection_actions.js +++ b/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); }