From fad90152708e716ece3a9528e9b7a09edd24b0da Mon Sep 17 00:00:00 2001
From: Vladimir Mandic <mandic00@live.com>
Date: Sat, 9 May 2026 17:09:25 +0200
Subject: [PATCH] fix vae preview flash previous image

Co-authored-by: Copilot <copilot@github.com>
Signed-off-by: Vladimir Mandic <mandic00@live.com>
---
 CHANGELOG.md                |  1 +
 javascript/progressBar.js   | 13 +++++++++----
 modules/api/server.py       |  4 ++--
 modules/vae/sd_vae_taesd.py |  5 ++++-
 4 files changed, 16 insertions(+), 7 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 61a36663e..68dde055e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -86,6 +86,7 @@ For full details, see [ChangeLog](https://github.com/vladmandic/automatic/blob/m
   - avoid `callback` duplicate registrations
   - pipeline task change causing loss of info on loaded `lora`
   - `detailer` handle `lora` internally
+  - vae preview flashes previous image
 
 ## Update for 2026-04-28
 
diff --git a/javascript/progressBar.js b/javascript/progressBar.js
index 866a0deb5..1ad591759 100644
--- a/javascript/progressBar.js
+++ b/javascript/progressBar.js
@@ -116,7 +116,7 @@ function requestProgress(id_task = 'undefined', progressEl = null, galleryEl = n
     };
   };
 
-  const done = () => {
+  const done = (ok = false) => {
     debug('taskEnd:', id_task);
     localStorage.removeItem('task');
     setProgress();
@@ -126,6 +126,11 @@ function requestProgress(id_task = 'undefined', progressEl = null, galleryEl = n
     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';
@@ -148,17 +153,17 @@ function requestProgress(id_task = 'undefined', progressEl = null, galleryEl = n
       hasStarted |= res.active;
       if (res.completed || (!res.active && (hasStarted || once))) {
         debug('progress', { end: res, reason: res.completed ? 'completed' : 'inactive' });
-        if (!res.paused) done(); // only abort if not paused
+        if (!res.paused) done(true); // only abort if not paused
         return;
       }
       if (elapsedFromStart > progressTimeout && !res.queued && res.progress === prevProgress) {
         debug('progress', { end: res, reason: 'progressSimeout' });
-        if (!res.paused) done(); // only abort if not paused
+        if (!res.paused) done(false); // only abort if not paused
         return;
       }
       if (elapsedFromStart > startTimeout && !res.queued && !res.active) {
         debug('progress', { end: res, reason: 'startTimeout' });
-        if (!res.paused) done(); // only abort if not paused
+        if (!res.paused) done(false); // only abort if not paused
         return;
       }
       if (res.progress !== prevProgress) {
diff --git a/modules/api/server.py b/modules/api/server.py
index 7daf91af5..db1e75f8d 100644
--- a/modules/api/server.py
+++ b/modules/api/server.py
@@ -116,8 +116,8 @@ def get_progress(req: models.ReqProgress = Depends()):
     progress = min((current / total) if current > 0 and total > 0 else 0, 1)
     time_since_start = time.time() - shared.state.time_start
     eta_relative = (time_since_start / progress) - time_since_start if progress > 0 else 0
-    # log.critical(f'get_progress: batch {batch_x}/{batch_y} step {step_x}/{step_y} current {current}/{total} time={time_since_start} eta={eta_relative}')
-    # log.critical(shared.state)
+    # log.trace(f'get_progress: batch {batch_x}/{batch_y} step {step_x}/{step_y} current {current}/{total} time={time_since_start} eta={eta_relative}')
+    # log.trace(shared.state)
     res = models.ResProgress(id=shared.state.id, progress=round(progress, 2), eta_relative=round(eta_relative, 2), current_image=current_image, textinfo=shared.state.textinfo, state=shared.state.dict(), )
     return res
 
diff --git a/modules/vae/sd_vae_taesd.py b/modules/vae/sd_vae_taesd.py
index 4050f9844..aedbf799e 100644
--- a/modules/vae/sd_vae_taesd.py
+++ b/modules/vae/sd_vae_taesd.py
@@ -59,9 +59,12 @@ def warn_once(msg, variant=None):
 def get_model(model_cls, variant=None):
     if variant is not None:
         pass
-    if model_cls in {'sd', 'sdxl', 'ldm', 'pixartalpha'}:
+    if model_cls in {'sd'}:
         model_cls = 'sd'
         variant = shared.opts.taesd_variant
+    elif model_cls in {'sdxl', 'ldm', 'pixartalpha'}:
+        model_cls = 'sdxl'
+        variant = shared.opts.taesd_variant
     elif model_cls in {'pixartsigma', 'hunyuandit', 'omnigen', 'auraflow'}:
         model_cls = 'sdxl'
         variant = shared.opts.taesd_variant