Merge pull request #4883 from vladmandic/fix/taesd-preview-size

fix(preview): keep taesd preview size constant across decode layers

modules/vae/sd_vae_taesd.py

@@ -173,6 +173,23 @@ def load_model(model_type = 'decoder', variant = None, vae_file: str | None = No
     return None, variant
 
 
+def restore_preview_size(image, vae):
+    # TAESD (image) and TAEHV (video) drop spatial upsample blocks when taesd_layers < 3, shrinking output 2x/4x.
+    # Rescale spatial dims so preview size stays constant. Other taes (TAEM1, Hybrid) ignore taesd_layers, so skip them.
+    from modules.taesd.taesd import TAESD
+    from modules.taesd.taehv import TAEHV
+    layers = shared.opts.taesd_layers
+    if layers >= 3 or not isinstance(vae, (TAESD, TAEHV)) or not isinstance(image, torch.Tensor) or image.ndim < 3 or image.shape[-3] != 3:
+        return image
+    try:
+        frames = image.reshape(-1, *image.shape[-3:]) # flatten any leading dims to a batch of CHW frames
+        frames = torch.nn.functional.interpolate(frames, scale_factor=float(2 ** (3 - layers)), mode='bilinear', align_corners=False)
+        image = frames.reshape(*image.shape[:-2], frames.shape[-2], frames.shape[-1])
+    except Exception:
+        pass
+    return image
+
+
 def decode(latents):
     global first_run # pylint: disable=global-statement
     with lock:
@@ -202,6 +219,7 @@ def decode(latents):
                 else:
                     image = vae.decode(tensor, return_dict=False)[0]
                     image = (image / 2.0 + 0.5).clamp(0, 1).detach()
+                image = restore_preview_size(image, vae)
                 t1 = time.time()
                 if (t1 - t0) > 3.0 and not first_run:
                     log.warning(f'Decode: type="taesd" variant="{variant}" long decode time={t1 - t0:.2f}')