diff --git a/modules/api/api.py b/modules/api/api.py index 86c67ae84..a5f828f66 100644 --- a/modules/api/api.py +++ b/modules/api/api.py @@ -1,7 +1,7 @@ import os from threading import Lock from secrets import compare_digest -from fastapi import FastAPI, APIRouter, Depends, Request +from fastapi import FastAPI, APIRouter, Depends, Request, WebSocket from fastapi.security import HTTPBasic, HTTPBasicCredentials from fastapi.exceptions import HTTPException from modules import errors, shared, paths @@ -133,6 +133,16 @@ class Api: from modules.api import gallery gallery.register_api(self) + # preview websocket api + from modules.api.preview import ws_preview + @self.app.websocket("/sdapi/v1/preview") + async def _ws_preview(websocket: WebSocket): + await ws_preview(websocket) + if shared.opts.subpath is not None and len(shared.opts.subpath) > 0: + @self.app.websocket(f"{shared.opts.subpath}/sdapi/v1/preview") + async def _ws_preview_subpath(websocket: WebSocket): + await ws_preview(websocket) + # nudenet api from modules.api import nudenet nudenet.register_api(self) diff --git a/modules/api/preview.py b/modules/api/preview.py new file mode 100644 index 000000000..9a55abbf2 --- /dev/null +++ b/modules/api/preview.py @@ -0,0 +1,131 @@ +import io +import asyncio +from fastapi import WebSocket, WebSocketDisconnect +from modules.logger import log + + +class PreviewManager: + """Manages WebSocket connections for step-based live preview streaming. + + Each denoising step, diffusers_callback calls push_step() from the + generation thread. Image encoding happens there (non‑blocking for + TAESD, the default preview method) and the binary JPEG + JSON progress + are dispatched to the event loop via run_coroutine_threadsafe(). + """ + + def __init__(self): + self.connections: set[WebSocket] = set() + self._client_hidden: bool = False + self._loop: asyncio.AbstractEventLoop | None = None + + # ------------------------------------------------------------------ + # Lifecycle (called from the event-loop thread) + # ------------------------------------------------------------------ + + async def connect(self, ws: WebSocket): + await ws.accept() + self.connections.add(ws) + if self._loop is None: + self._loop = asyncio.get_event_loop() + log.debug(f'Preview WS connect: client={ws.client.host} total={len(self.connections)}') + + async def disconnect(self, ws: WebSocket): + self.connections.discard(ws) + log.debug(f'Preview WS disconnect: client={ws.client.host} total={len(self.connections)}') + + async def handle_message(self, data: dict): + """Receive visibility-change messages from the client.""" + if data.get('type') == 'visibility': + self._client_hidden = not data.get('visible', True) + + # ------------------------------------------------------------------ + # Push from generation thread (thread‑safe) + # ------------------------------------------------------------------ + + def push_step(self, step: int, steps: int, state): + """Called from the generation thread inside diffusers_callback.""" + if not self.connections or self._client_hidden or self._loop is None: + return + + # --- lightweight metadata (no GPU ops) --- + progress = round(step / steps, 2) if steps > 0 else 0 + data = { + "active": True, + "step": step, + "steps": steps, + "progress": progress, + "job": state.job, + "paused": state.paused, + } + + # --- encode preview image --- + img_bytes = None + try: + if state.set_current_image() and state.current_image is not None: + buf = io.BytesIO() + state.current_image.save(buf, format='jpeg', quality=60) + img_bytes = buf.getvalue() + except Exception: + pass + + data["has_image"] = img_bytes is not None + + # --- dispatch to event loop --- + asyncio.run_coroutine_threadsafe( + self._broadcast(data, img_bytes), + self._loop + ) + + def push_complete(self): + """Called when a generation job finishes.""" + if not self.connections or self._loop is None: + return + data = { + "active": False, + "step": 0, + "steps": 0, + "progress": 1.0, + "job": "", + "paused": False, + "completed": True, + "has_image": False, + } + asyncio.run_coroutine_threadsafe( + self._broadcast(data, None), + self._loop + ) + + # ------------------------------------------------------------------ + # Async broadcast (runs on event loop) + # ------------------------------------------------------------------ + + async def _broadcast(self, data: dict, img_bytes: bytes | None): + dead: set[WebSocket] = set() + for ws in self.connections: + try: + await ws.send_json(data) + if img_bytes: + await ws.send_bytes(img_bytes) + except Exception: + dead.add(ws) + self.connections -= dead + + +preview_manager = PreviewManager() + + +async def ws_preview(websocket: WebSocket): + await preview_manager.connect(websocket) + try: + while True: + raw = await websocket.receive_text() + try: + import json + msg = json.loads(raw) + await preview_manager.handle_message(msg) + except Exception: + pass + except WebSocketDisconnect: + pass + finally: + await preview_manager.disconnect(websocket) diff --git a/modules/processing.py b/modules/processing.py index 8d4d6bab4..6b38831d7 100644 --- a/modules/processing.py +++ b/modules/processing.py @@ -576,5 +576,8 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed: if shared.cmd_opts.lowvram or shared.cmd_opts.medvram: devices.torch_gc(force=True, reason='final') + if shared.opts.live_preview_transport == "WebSocket": + from modules.api.preview import preview_manager + preview_manager.push_complete() shared.state.end(jobid) return results diff --git a/modules/processing_callbacks.py b/modules/processing_callbacks.py index 0df281c47..4aab82584 100644 --- a/modules/processing_callbacks.py +++ b/modules/processing_callbacks.py @@ -231,4 +231,13 @@ def diffusers_callback(pipe, step: int = 0, timestep: int = 0, kwargs: dict | No shared.profiler.step() t1 = time.time() timer.process.add('callback', t1 - t0) + if shared.opts.live_preview_transport == "WebSocket": + from modules.api.preview import preview_manager + interval = max(1, shared.opts.live_preview_ws_interval) + if shared.state.sampling_step % interval == 0: + preview_manager.push_step( + shared.state.sampling_step, + shared.state.sampling_steps, + shared.state + ) return kwargs diff --git a/modules/ui_definitions.py b/modules/ui_definitions.py index ec497f9e9..abac5d06b 100644 --- a/modules/ui_definitions.py +++ b/modules/ui_definitions.py @@ -567,7 +567,9 @@ def create_settings(cmd_opts): 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_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 update period of polling (ms)", gr.Slider, {"minimum": 0, "maximum": 5000, "step": 25}), + "live_preview_transport": OptionInfo("Polling", "Preview transport", gr.Dropdown, {"choices": ["Polling", "WebSocket"]}), + "live_preview_ws_interval": OptionInfo(1, "WebSocket preview step interval", gr.Slider, {"minimum": 1, "maximum": 20, "step": 1}), "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..0988591ef 100644 --- a/ui/progressBar.ts +++ b/ui/progressBar.ts @@ -146,53 +146,135 @@ export function requestProgress(id_task = 'undefined', progressEl = null, galler if (atEnd) atEnd(); }; - const startLivePreview = (taskId: string, id_live_preview: number) => { + const onProgressHandler = (res) => { + 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); + return; + } + if (elapsedFromStart > progressTimeout && !res.queued && res.progress === prevProgress) { + 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); + return; + } + if (res.progress !== prevProgress) { + dateStart = Date.now(); + 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; + } + if (onProgress) onProgress(res); + }; + + const onProgressErrorHandler = (err) => { + error('progress', { error: err }); + removeLivePreview(false); + }; + + const startHttpPolling = (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 onProgressHandler = (res) => { - if (res?.debug) debug('progress:', { start: dateStart, id: request_id, 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 - return; + const wrappedHandler = (res) => { + onProgressHandler(res); + if (res.completed || (!res.active && (hasStarted || once))) return; + if (!res.paused) { + setTimeout(() => startHttpPolling(taskId, res.id_live_preview || 0), window.opts.live_preview_refresh_period || 500); } - if (elapsedFromStart > progressTimeout && !res.queued && res.progress === prevProgress) { - debug('progress', { end: res, reason: 'progressSimeout' }); - if (!res.paused) removeLivePreview(false); // only abort if not paused - return; - } - if (elapsedFromStart > startTimeout && !res.queued && !res.active) { - debug('progress', { end: res, reason: 'startTimeout' }); - if (!res.paused) removeLivePreview(false); // only abort if not paused - return; - } - if (res.progress !== prevProgress) { - dateStart = Date.now(); - 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); }; - - const onProgressErrorHandler = (err) => { - 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: taskId, id_live_preview: request_id }, wrappedHandler, onProgressErrorHandler, false, 30000); }; + + const startLivePreviewWebSocket = () => { + const proto = location.protocol === 'https:' ? 'wss:' : 'ws:'; + const ws = new WebSocket(`${proto}//${location.host}/sdapi/v1/preview`); + let wsFailed = false; + let revokeUrl: string | undefined; + + const sendVisibility = () => { + if (ws.readyState === WebSocket.OPEN) { + ws.send(JSON.stringify({ type: 'visibility', visible: !document.hidden })); + } + }; + + const cleanup = (ok: boolean) => { + document.removeEventListener('visibilitychange', sendVisibility); + if (revokeUrl) { URL.revokeObjectURL(revokeUrl); revokeUrl = undefined; } + ws.close(); + if (!ok && !wsFailed) { + wsFailed = true; + debug('progress', 'ws fallback to polling'); + startHttpPolling(id_task, 0); + } + }; + + ws.onopen = () => { + sendVisibility(); + debug('progress', 'ws connected'); + }; + + ws.onmessage = (event) => { + if (event.data instanceof Blob) { + if (!livePreview) initLivePreview(); + const url = URL.createObjectURL(event.data); + if (revokeUrl) URL.revokeObjectURL(revokeUrl); + revokeUrl = url; + img.src = url; + if (onProgress) onProgress(lastState); + } else if (typeof event.data === 'string') { + try { + const res = JSON.parse(event.data); + if (res.completed) { + onProgressHandler(res); + cleanup(true); + removeLivePreview(true); + return; + } + onProgressHandler(res); + } catch (e) { + error('progress ws', { error: e }); + } + } + }; + + ws.onerror = () => { cleanup(false); }; + + ws.onclose = (event) => { + if (event.code !== 1000) cleanup(false); + }; + + document.addEventListener('visibilitychange', sendVisibility); + + const fallbackTimer = setTimeout(() => { + if (ws.readyState !== WebSocket.OPEN) { + wsFailed = true; + debug('progress', 'ws timeout, fallback to polling'); + ws.close(); + startHttpPolling(id_task, 0); + } + }, 1000); + + ws.addEventListener('open', () => clearTimeout(fallbackTimer)); + }; + debug('progress', { start: dateStart }); - startLivePreview(id_task, 0); + const transport = window.opts.live_preview_transport || 'Polling'; + if (transport === 'WebSocket') { + startLivePreviewWebSocket(); + } else { + startHttpPolling(id_task, 0); + } } window.checkPaused = checkPaused;