diff --git a/modules/processing_callbacks.py b/modules/processing_callbacks.py index 0df281c47..bdd1474c9 100644 --- a/modules/processing_callbacks.py +++ b/modules/processing_callbacks.py @@ -73,6 +73,18 @@ def diffusers_callback(pipe, step: int = 0, timestep: int = 0, kwargs: dict | No shared.state.current_latent = None shared.state.step() # increase step shared.state.preview_job = -1 # indicate that preview image has changed + # Step-based progress and live preview push + from modules import progress + total = max(shared.state.sampling_steps, 1) + prog = round(min(1, shared.state.sampling_step / total), 2) + progress.push_progress(shared.state.sampling_step, total, prog) + if shared.opts.show_progress_every_n_steps > 0 and step % shared.opts.show_progress_every_n_steps == 0: + try: + shared.state.current_image_sampling_step = shared.state.sampling_step + shared.state.assign_current_image(image) + progress.push_live_preview(shared.state.current_image, shared.state.sampling_step, total, prog) + except Exception: + pass debug_callback(f'Callback: step={step} timestep={timestep} image={image if image is not None else None} kwargs={list(kwargs)}') return kwargs @@ -229,6 +241,17 @@ def diffusers_callback(pipe, step: int = 0, timestep: int = 0, kwargs: dict | No # errors.display(e, 'Callback') if shared.cmd_opts.profile and shared.profiler is not None: shared.profiler.step() + # Step-based progress and live preview push + from modules import progress + total = max(shared.state.sampling_steps, 1) + prog = round(min(1, shared.state.sampling_step / total), 2) + progress.push_progress(shared.state.sampling_step, total, prog) + if shared.opts.show_progress_every_n_steps > 0 and step % shared.opts.show_progress_every_n_steps == 0: + try: + if shared.state.set_current_image() and shared.state.current_image is not None: + progress.push_live_preview(shared.state.current_image, shared.state.sampling_step, total, prog) + except Exception: + pass t1 = time.time() timer.process.add('callback', t1 - t0) return kwargs diff --git a/modules/progress.py b/modules/progress.py index 2fcd4dc22..46a94f798 100644 --- a/modules/progress.py +++ b/modules/progress.py @@ -2,7 +2,9 @@ import base64 import os import io import time +import asyncio from pydantic import BaseModel, Field # pylint: disable=no-name-in-module +from starlette.websockets import WebSocket, WebSocketState import modules.shared as shared from modules.logger import log @@ -16,6 +18,99 @@ debug = os.environ.get('SD_PREVIEW_DEBUG', None) is not None debug_log = log.trace if debug else lambda *args, **kwargs: None +class PreviewManager: + def __init__(self): + self.active: list[WebSocket] = [] + self._loop = None + + async def connect(self, ws: WebSocket): + await ws.accept() + self._loop = asyncio.get_running_loop() + self.active.append(ws) + debug_log(f'Preview WS connect: client={ws.client.host} total={len(self.active)}') + + def disconnect(self, ws: WebSocket): + try: + self.active.remove(ws) + except ValueError: + pass + debug_log(f'Preview WS disconnect: client={ws.client.host} total={len(self.active)}') + + async def broadcast(self, data: dict): + disconnected = [] + for ws in list(self.active): + try: + if ws.client_state == WebSocketState.CONNECTED: + await ws.send_json(data) + else: + disconnected.append(ws) + except Exception: + disconnected.append(ws) + for ws in disconnected: + try: + self.active.remove(ws) + except ValueError: + pass + + def push(self, data: dict): + if self._loop is None or not self.active: + return + try: + asyncio.run_coroutine_threadsafe(self.broadcast(data), self._loop) + except Exception: + pass + + +preview_manager = PreviewManager() + + +def push_progress(step: int, steps: int, progress: float, active: bool = True): + try: + data = { + 'type': 'progress', + 'step': step, + 'steps': steps, + 'progress': progress, + 'active': active, + 'job': shared.state.job, + 'paused': shared.state.paused, + } + preview_manager.push(data) + except Exception as e: + debug_log(f'Progress push error: {e}') + + +def push_live_preview(image, step: int, steps: int, progress: float): + try: + buffered = io.BytesIO() + image.save(buffered, format='jpeg', quality=60) + b64 = base64.b64encode(buffered.getvalue()).decode('ascii') + data = { + 'type': 'preview', + 'live_preview': f'data:image/jpeg;base64,{b64}', + 'id_live_preview': shared.state.id_live_preview, + 'step': step, + 'steps': steps, + 'progress': progress, + 'job': shared.state.job, + } + preview_manager.push(data) + except Exception as e: + debug_log(f'Preview push error: {e}') + + +def push_complete(): + try: + data = { + 'type': 'complete', + 'active': False, + 'job': shared.state.job, + } + preview_manager.push(data) + except Exception as e: + debug_log(f'Complete push error: {e}') + + def start_task(id_task): global current_task # pylint: disable=global-statement current_task = id_task @@ -35,6 +130,7 @@ def finish_task(id_task): finished_tasks.append(id_task) if len(finished_tasks) > 16: finished_tasks.pop(0) + push_complete() def add_task_to_queue(id_job): @@ -94,15 +190,10 @@ def api_progress(req: ProgressRequest): debug_log(f'Progress: job="{shared.state.job}" active={active} progress={step}/{steps}/{progress} image={shared.state.current_image_sampling_step} request={id_live_preview} last={shared.state.id_live_preview} job={shared.state.preview_job} elapsed={elapsed:.3f}') - if active and (req.id_live_preview != -1): - have_image = shared.state.set_current_image() - if have_image and shared.state.current_image is not None: - buffered = io.BytesIO() - shared.state.current_image.save(buffered, format='jpeg', quality=60) - b64 = base64.b64encode(buffered.getvalue()) - live_preview = f'data:image/jpeg;base64,{b64.decode("ascii")}' - else: - live_preview = None + if active: + shared.state.set_current_image() + if shared.state.current_image is not None and req.id_live_preview != shared.state.id_live_preview: + live_preview = shared.state.live_preview_b64 id_live_preview = shared.state.id_live_preview @@ -129,3 +220,17 @@ def api_progress(req: ProgressRequest): def setup_progress_api(): shared.api.add_api_route("/internal/progress", api_progress, methods=["POST"], response_model=InternalProgressResponse) + + @shared.api.app.websocket("/ws/preview") + async def ws_preview(ws: WebSocket): + await preview_manager.connect(ws) + try: + while True: + data = await ws.receive_text() + if data == 'ping': + await ws.send_text('pong') + elif data == 'end': + break + except Exception: + pass + preview_manager.disconnect(ws) diff --git a/modules/shared_state.py b/modules/shared_state.py index ca017d9e9..1f07c6200 100644 --- a/modules/shared_state.py +++ b/modules/shared_state.py @@ -40,6 +40,7 @@ class State: current_image = None current_image_sampling_step = 0 id_live_preview = 0 + live_preview_b64 = None textinfo = None prediction_type = "epsilon" api = False @@ -271,13 +272,18 @@ class State: def do_set_current_image(self): from modules import shared, images, sd_samplers_common - if self.disable_preview or (self.preview_job == self.job_no): + if self.disable_preview: + return False + + # Skip if we already decoded a preview for the current sampling step + if self.current_image_sampling_step == self.sampling_step and self.current_image is not None: return False if (shared.opts.show_progress_type == "None") and (shared.history.last_image is not None): + self.preview_job = self.job_no last_image = images.image_grid(shared.history.last_image) self.assign_current_image(last_image) - self.preview_job = -1 + self.current_image_sampling_step = self.sampling_step return True if self.current_latent is not None: @@ -296,22 +302,25 @@ class State: pass # ignore sigma errors image = sd_samplers_common.samples_to_image_grid(sample) self.assign_current_image(image) - self.preview_job = -1 return True except Exception as e: - self.preview_job = -1 log.error(f'State image: last={self.id_live_preview} step={self.sampling_step} {e}') display(e, 'State image') return False elif self.current_image is not None: self.preview_job = self.job_no + self.current_image_sampling_step = self.sampling_step self.assign_current_image(self.current_image) - self.preview_job = -1 return True else: pass return False def assign_current_image(self, image): + import base64 + import io self.current_image = image self.id_live_preview += 1 + buffered = io.BytesIO() + image.save(buffered, format='jpeg', quality=60) + self.live_preview_b64 = f'data:image/jpeg;base64,{base64.b64encode(buffered.getvalue()).decode("ascii")}' diff --git a/modules/ui_definitions.py b/modules/ui_definitions.py index ec497f9e9..7e7fac4d2 100644 --- a/modules/ui_definitions.py +++ b/modules/ui_definitions.py @@ -565,9 +565,9 @@ def create_settings(cmd_opts): # --- Live Previews --- options_templates.update(options_section(('live-preview', "Live Previews"), { - "show_progress_every_n_steps": OptionInfo(1, "Live preview display period", gr.Slider, {"minimum": 0, "maximum": 20, "step": 1, "visible": False}), + "show_progress_every_n_steps": OptionInfo(1, "Show live preview every N steps", gr.Slider, {"minimum": 1, "maximum": 20, "step": 1}, info="1 = every step, 5 = every 5th step, etc."), "show_progress_type": OptionInfo("TAESD", "Live preview method", gr.Dropdown, {"choices": ["None", "Simple", "Approximate", "TAESD", "Full"]}), - "live_preview_refresh_period": OptionInfo(500, "Progress update period", gr.Slider, {"minimum": 0, "maximum": 5000, "step": 25}), + "live_preview_refresh_period": OptionInfo(500, "Progress status update period", gr.Slider, {"minimum": 100, "maximum": 5000, "step": 25}, info="How often to poll for progress text updates in ms"), "taesd_variant": OptionInfo(shared_items.sd_taesd_items()[0], "TAESD variant", gr.Dropdown, {"choices": shared_items.sd_taesd_items()}), "taesd_layers": OptionInfo(3, "TAESD decode layers", gr.Slider, {"minimum": 1, "maximum": 3, "step": 1}), "live_preview_downscale": OptionInfo(True, "Downscale high resolution live previews"), diff --git a/ui/progressBar.ts b/ui/progressBar.ts index 831822ead..33a6645c6 100644 --- a/ui/progressBar.ts +++ b/ui/progressBar.ts @@ -13,21 +13,20 @@ export function setRefreshInterval() { document.addEventListener('visibilitychange', () => { if (document.hidden) refreshInterval = Math.max(2500, window.opts.live_preview_refresh_period || 1000); else refreshInterval = window.opts.live_preview_refresh_period || 1000; - // log('refreshInterval', document.visibilityState, refreshInterval); }); } -function pad2(x) { +function pad2(x: number) { return x < 10 ? `0${x}` : x; } -function formatTime(secs) { +function formatTime(secs: number) { if (secs > 3600) return `${pad2(Math.floor(secs / 60 / 60))}:${pad2(Math.floor(secs / 60) % 60)}:${pad2(Math.floor(secs) % 60)}`; if (secs > 60) return `${pad2(Math.floor(secs / 60))}:${pad2(Math.floor(secs) % 60)}`; return `${Math.floor(secs)}s`; } -export function checkPaused(state) { +export function checkPaused(state?: boolean) { lastState.paused = state ? !state : !lastState.paused; const t_el = document.getElementById('txt2img_pause'); const i_el = document.getElementById('img2img_pause'); @@ -84,26 +83,33 @@ export function randomId() { return `task(${Math.random().toString(36).slice(2, 7)}${Math.random().toString(36).slice(2, 7)}${Math.random().toString(36).slice(2, 7)})`; } -// starts sending progress requests to "/internal/progress" uri, creating progressbar above progressbarContainer element and preview inside gallery element -// Cleans up all created stuff when the task is over and calls atEnd. calls onProgress every time there is a progress update -export function requestProgress(id_task = 'undefined', progressEl = null, galleryEl = null, atEnd = null, onProgress = null, once = false) { +function getWebSocketUrl(): string { + const proto = location.protocol === 'https:' ? 'wss:' : 'ws:'; + return `${proto}//${location.host}/ws/preview`; +} + +export function requestProgress(id_task = 'undefined', progressEl: HTMLElement | null = null, galleryEl: HTMLElement | null = null, atEnd: (() => void) | null = null, onProgress: ((res: any) => void) | null = null, once = false) { if (id_task) localStorage.setItem('task', id_task); let hasStarted = false; let dateStart = Date.now(); let prevProgress: any = null; - const parentGallery = galleryEl ? galleryEl.parentNode : null; + const parentGallery: HTMLElement | null = galleryEl ? galleryEl.parentNode as HTMLElement : null; let livePreview: HTMLElement | undefined; let img: HTMLImageElement; + let ws: WebSocket | null = null; + let wsReconnect: number | undefined; + let pollingTimer: number | undefined; + let id_live_preview = 0; const initLivePreview = () => { if (!parentGallery) return; const footers = Array.from(gradioApp().querySelectorAll('.gallery_footer')); for (const footer of footers) { - if (footer.id !== 'gallery_footer') footer.style.display = 'none'; // remove all footers + if (footer.id !== 'gallery_footer') footer.style.display = 'none'; } const galleries = Array.from(gradioApp().querySelectorAll('.gallery_main')); for (const gallery of galleries) { - if (gallery.id !== 'gallery_gallery') gallery.style.display = 'none'; // remove all footers + if (gallery.id !== 'gallery_gallery') gallery.style.display = 'none'; } livePreview = document.createElement('div'); @@ -124,50 +130,114 @@ export function requestProgress(id_task = 'undefined', progressEl = null, galler debug('taskEnd:', id_task); localStorage.removeItem('task'); setProgress(); - const footers = Array.from(gradioApp().querySelectorAll('.gallery_footer')); - for (const footer of footers) footer.style.display = 'flex'; // restore all footers - const galleries = Array.from(gradioApp().querySelectorAll('.gallery_main')); - for (const gallery of galleries) gallery.style.display = 'flex'; // remove all galleries - try { - if (parentGallery && livePreview) { - if (ok) { - const previewImg = gradioApp().querySelector('#livePreviewImage'); - const galleryImg = gradioApp().querySelector('#control_gallery img'); - if (previewImg?.src && galleryImg) galleryImg.src = previewImg.src; // copy preview to gallery if everything is ok - } - parentGallery.removeChild(livePreview); - parentGallery.style.minHeight = 'unset'; - parentGallery.style.maxHeight = 'unset'; - parentGallery.style.overflow = 'unset'; + if (ws) { + try { ws.send('end'); } catch { /* ignore */ } + try { ws.close(); } catch { /* ignore */ } + ws = null; + } + if (wsReconnect) { + clearTimeout(wsReconnect); + wsReconnect = undefined; + } + if (pollingTimer) { + clearTimeout(pollingTimer); + pollingTimer = undefined; + } + const footers = gradioApp().querySelectorAll('.gallery_footer'); + for (const footer of Array.from(footers)) footer.style.display = 'flex'; + const galleries = gradioApp().querySelectorAll('.gallery_main'); + for (const gallery of Array.from(galleries)) gallery.style.display = 'flex'; + if (parentGallery && livePreview && livePreview.parentNode) { + if (ok) { + const previewImg = gradioApp().querySelector('#livePreviewImage') as HTMLImageElement; + const galleryImg = gradioApp().querySelector('#control_gallery img') as HTMLImageElement; + if (previewImg?.src && galleryImg) galleryImg.src = previewImg.src; } - } catch { /* ignore */ } + parentGallery.removeChild(livePreview); + parentGallery.style.minHeight = 'unset'; + parentGallery.style.maxHeight = 'unset'; + parentGallery.style.overflow = 'unset'; + } checkPaused(true); sendNotification(); if (atEnd) atEnd(); }; - const startLivePreview = (taskId: string, id_live_preview: number) => { - if (window.opts.live_preview_refresh_period === 0) return; - const request_id = document.hidden ? -1 : id_live_preview; + const onWsMessage = (event: MessageEvent) => { + try { + const data = JSON.parse(event.data); + if (data.type === 'preview' && data.live_preview) { + if (!livePreview) initLivePreview(); + if (livePreview && galleryEl && img && img.src !== data.live_preview) { + img.src = data.live_preview; + id_live_preview = data.id_live_preview; + lastState = { ...lastState, step: data.step, steps: data.steps, progress: data.progress, job: data.job }; + setProgress(lastState); + if (onProgress) onProgress(lastState); + dateStart = Date.now(); + prevProgress = data.progress; + } + } else if (data.type === 'progress') { + lastState = { ...lastState, step: data.step, steps: data.steps, progress: data.progress, active: data.active, paused: data.paused, job: data.job }; + setProgress(lastState); + if (data.progress !== prevProgress) { + dateStart = Date.now(); + prevProgress = data.progress; + } + if (onProgress) onProgress(lastState); + } else if (data.type === 'complete') { + removeLivePreview(true); + } + } catch { /* ignore */ } + }; - const onProgressHandler = (res) => { - if (res?.debug) debug('progress:', { start: dateStart, id: request_id, res }); + const connectWebSocket = () => { + if (ws) return; + try { + ws = new WebSocket(getWebSocketUrl()); + ws.onmessage = onWsMessage; + ws.onopen = () => { + debug('ws', 'connected'); + if (wsReconnect) { + clearTimeout(wsReconnect); + wsReconnect = undefined; + } + }; + ws.onclose = () => { + debug('ws', 'disconnected'); + ws = null; + if (pollingTimer !== undefined) { + wsReconnect = window.setTimeout(connectWebSocket, 3000); + } + }; + ws.onerror = () => { + debug('ws', 'error'); + }; + } catch { + wsReconnect = window.setTimeout(connectWebSocket, 3000); + } + }; + + const pollProgress = () => { + if (window.opts.live_preview_refresh_period === 0) return; + const onProgressHandler = (res: any) => { + if (res?.debug) debug('progress:', { start: dateStart, res }); lastState = res; const elapsedFromStart = (Date.now() - dateStart) / 1000; hasStarted = hasStarted || res.active; if (res.completed || (!res.active && (hasStarted || once))) { debug('progress', { end: res, reason: res.completed ? 'completed' : 'inactive' }); - if (!res.paused) removeLivePreview(true); // only abort if not paused + if (!res.paused) removeLivePreview(true); return; } if (elapsedFromStart > progressTimeout && !res.queued && res.progress === prevProgress) { - debug('progress', { end: res, reason: 'progressSimeout' }); - if (!res.paused) removeLivePreview(false); // only abort if not paused + debug('progress', { end: res, reason: 'progressTimeout' }); + if (!res.paused) removeLivePreview(false); return; } if (elapsedFromStart > startTimeout && !res.queued && !res.active) { debug('progress', { end: res, reason: 'startTimeout' }); - if (!res.paused) removeLivePreview(false); // only abort if not paused + if (!res.paused) removeLivePreview(false); return; } if (res.progress !== prevProgress) { @@ -175,24 +245,21 @@ export function requestProgress(id_task = 'undefined', progressEl = null, galler prevProgress = res.progress; } setProgress(res); - if (res.live_preview && !livePreview) initLivePreview(); - if (res.live_preview && galleryEl) { - if (img.src !== res.live_preview) img.src = res.live_preview; - id_live_preview = res.id_live_preview; - } if (onProgress) onProgress(res); - setTimeout(() => startLivePreview(id_task, id_live_preview), window.opts.live_preview_refresh_period || 500); + pollingTimer = window.setTimeout(pollProgress, window.opts.live_preview_refresh_period || 500); }; - const onProgressErrorHandler = (err) => { + const onProgressErrorHandler = (err: any) => { error('progress', { error: err }); removeLivePreview(false); }; - xhrPost('./internal/progress', { id_task, id_live_preview: request_id }, onProgressHandler, onProgressErrorHandler, false, 30000); + xhrPost('./internal/progress', { id_task, id_live_preview: -1 }, onProgressHandler, onProgressErrorHandler, false, 30000); }; + debug('progress', { start: dateStart }); - startLivePreview(id_task, 0); + connectWebSocket(); + pollProgress(); } window.checkPaused = checkPaused;