Made live preview and progress tracker per step based with new UI setting

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

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)

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")}'

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

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<any>(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<any>(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<any>(gradioApp().querySelectorAll('.gallery_footer'));
-    for (const footer of footers) footer.style.display = 'flex'; // restore all footers
-    const galleries = Array.from<any>(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<any>(footers)) footer.style.display = 'flex';
+    const galleries = gradioApp().querySelectorAll('.gallery_main');
+    for (const gallery of Array.from<any>(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;