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.
This commit is contained in:
CalamitousFelicitousness
2026-05-31 04:43:10 +01:00
parent a0531c70fd
commit f8bd53a96b
+18
View File
@@ -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}')