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