Adds per step based live preview as a option

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)

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)

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

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

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"),