root/PaintPlus
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Real per-step progress bars for local GPU inference

Browse Source 
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
This commit is contained in:
Claude
2026-06-13 16:48:55 +00:00
parent ee94282753
commit 372ab48991
5 changed files with 214 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/app/routers/ai_tools.py
+33
View File
@@ -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
+130 -47
View File
@@ -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
frontend/src/js/libs/progress_overlay.js
+38
View File
@@ -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
+6 -2
View File
@@ -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
+7 -1
View File
@@ -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);
}