From f8bd53a96b394680c9562cd0c76e098852d057e3 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Sun, 31 May 2026 04:43:10 +0100 Subject: [PATCH] fix(preview): keep taesd preview size constant across decode layers taesd_layers < 3 drops spatial upsample blocks in the TAESD/TAEHV decoders, shrinking preview output 2x/4x. Both UIs size the live preview from the image's intrinsic pixel dimensions (modern via object-fit: scale-down, standard via max(naturalWidth, 512px)), so lower layer counts rendered the preview physically small. Rescale the decoded preview spatially by 2^(3-layers) in sd_vae_taesd.decode. Gated to TAESD and TAEHV, the only decoders that honor taesd_layers; TAEM1 and Hybrid VAEs decode at full size and are left untouched. Rank-agnostic so it covers both image (CHW) and video (TCHW) previews, including single-frame video models used for txt2img. --- modules/vae/sd_vae_taesd.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/modules/vae/sd_vae_taesd.py b/modules/vae/sd_vae_taesd.py index 46a18fd9d..cbae08815 100644 --- a/modules/vae/sd_vae_taesd.py +++ b/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}')