mirror of
https://github.com/vladmandic/automatic
synced 2026-09-20 01:31:13 +02:00
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:
@@ -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}')
|
||||
|
||||
Reference in New Issue
Block a user