From 3fe1d090e4f0b5fe9aca88cda37acf616b8732df Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Mon, 6 Oct 2025 19:13:29 -0400 Subject: [PATCH] add configurable layers to taehv Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 9 +++++---- modules/sd_vae_taesd.py | 4 +++- modules/taesd/taehv.py | 34 +++++++++++++++++++++++++++------- 3 files changed, 35 insertions(+), 12 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index b2d4023e3..c4d78920c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -62,10 +62,10 @@ - networks ability to filter lora by base model version - add interrogate button to input images - **SDNQ** - - add `SVDQuant` quantization method support - - make sdnq scales compatible with balanced offload - - add int8 matmul support for RDNA2 GPUs via triton - - improve int8 mamtul performance on Intel GPUs + - add `SVDQuant` quantization method support + - make sdnq scales compatible with balanced offload + - add int8 matmul support for RDNA2 GPUs via triton + - improve int8 mamtul performance on Intel GPUs - **Other** - server will note when restart is recommended due to package updates - **interrrupt** will now show last known preview image @@ -90,6 +90,7 @@ can be combined with sdp, enabling may improve stability when used on iGPU or shared memory systems - **nunchaku** update to `1.0.1` and enhance installer - **xyz-grid** add guidance section + - **preview** implement configurable layers for WAN, Qwen, HV - **Video** - use shared T5 text encoder for video models when possible - unified video save code across all video models diff --git a/modules/sd_vae_taesd.py b/modules/sd_vae_taesd.py index aed995f1b..510215855 100644 --- a/modules/sd_vae_taesd.py +++ b/modules/sd_vae_taesd.py @@ -79,6 +79,7 @@ def get_model(model_type = 'decoder', variant = None): uri += '/tae' + model_cls + '_' + model_type + '.pth' try: torch.hub.download_url_to_file(uri, fn) + shared.log.print() # new line shared.log.info(f'Decode: type="taesd" variant="{variant}": uri="{uri}" fn="{fn}" download') except Exception as e: warn_once(f'download uri={uri} {e}', variant=variant) @@ -86,7 +87,8 @@ def get_model(model_type = 'decoder', variant = None): prev_cls = model_cls prev_type = model_type prev_model = variant - shared.log.debug(f'Decode: type="taesd" variant="{variant}" fn="{fn}" load') + shared.log.print() # new line + shared.log.debug(f'Decode: type="taesd" variant="{variant}" fn="{fn}" layers={shared.opts.taesd_layers} load') vae = None if 'TAE HunyuanVideo' in variant: from modules.taesd.taehv import TAEHV diff --git a/modules/taesd/taehv.py b/modules/taesd/taehv.py index 0f1828e45..bfeced7a7 100644 --- a/modules/taesd/taehv.py +++ b/modules/taesd/taehv.py @@ -159,6 +159,7 @@ class TAEHV(nn.Module): decoder_space_upscale: whether spatial upsampling is enabled for each block. upsampling can be disabled for a cheaper preview. """ super().__init__() + from modules import shared self.encoder = nn.Sequential( conv(TAEHV.image_channels, 64), nn.ReLU(inplace=True), TPool(64, 2), conv(64, 64, stride=2, bias=False), MemBlock(64, 64), MemBlock(64, 64), MemBlock(64, 64), @@ -168,13 +169,32 @@ class TAEHV(nn.Module): ) n_f = [256, 128, 64, 64] self.frames_to_trim = 2**sum(decoder_time_upscale) - 1 - self.decoder = nn.Sequential( - Clamp(), conv(TAEHV.latent_channels, n_f[0]), nn.ReLU(inplace=True), - MemBlock(n_f[0], n_f[0]), MemBlock(n_f[0], n_f[0]), MemBlock(n_f[0], n_f[0]), nn.Upsample(scale_factor=2 if decoder_space_upscale[0] else 1), TGrow(n_f[0], 1), conv(n_f[0], n_f[1], bias=False), - MemBlock(n_f[1], n_f[1]), MemBlock(n_f[1], n_f[1]), MemBlock(n_f[1], n_f[1]), nn.Upsample(scale_factor=2 if decoder_space_upscale[1] else 1), TGrow(n_f[1], 2 if decoder_time_upscale[0] else 1), conv(n_f[1], n_f[2], bias=False), - MemBlock(n_f[2], n_f[2]), MemBlock(n_f[2], n_f[2]), MemBlock(n_f[2], n_f[2]), nn.Upsample(scale_factor=2 if decoder_space_upscale[2] else 1), TGrow(n_f[2], 2 if decoder_time_upscale[1] else 1), conv(n_f[2], n_f[3], bias=False), - nn.ReLU(inplace=True), conv(n_f[3], TAEHV.image_channels), - ) + + if shared.opts.taesd_layers == 1: + self.decoder = nn.Sequential( + Clamp(), conv(TAEHV.latent_channels, n_f[0]), nn.ReLU(inplace=True), + MemBlock(n_f[0], n_f[0]), MemBlock(n_f[0], n_f[0]), MemBlock(n_f[0], n_f[0]), nn.Upsample(scale_factor=2 if decoder_space_upscale[0] else 1), TGrow(n_f[0], 1), conv(n_f[0], n_f[1], bias=False), + MemBlock(n_f[1], n_f[1]), MemBlock(n_f[1], n_f[1]), MemBlock(n_f[1], n_f[1]), nn.Identity(), TGrow(n_f[1], 2 if decoder_time_upscale[0] else 1), conv(n_f[1], n_f[2], bias=False), + MemBlock(n_f[2], n_f[2]), MemBlock(n_f[2], n_f[2]), MemBlock(n_f[2], n_f[2]), nn.Identity(), TGrow(n_f[2], 2 if decoder_time_upscale[1] else 1), conv(n_f[2], n_f[3], bias=False), + nn.ReLU(inplace=True), conv(n_f[3], TAEHV.image_channels), + ) + elif shared.opts.taesd_layers == 2: + self.decoder = nn.Sequential( + Clamp(), conv(TAEHV.latent_channels, n_f[0]), nn.ReLU(inplace=True), + MemBlock(n_f[0], n_f[0]), MemBlock(n_f[0], n_f[0]), MemBlock(n_f[0], n_f[0]), nn.Upsample(scale_factor=2 if decoder_space_upscale[0] else 1), TGrow(n_f[0], 1), conv(n_f[0], n_f[1], bias=False), + MemBlock(n_f[1], n_f[1]), MemBlock(n_f[1], n_f[1]), MemBlock(n_f[1], n_f[1]), nn.Upsample(scale_factor=2 if decoder_space_upscale[1] else 1), TGrow(n_f[1], 2 if decoder_time_upscale[0] else 1), conv(n_f[1], n_f[2], bias=False), + MemBlock(n_f[2], n_f[2]), MemBlock(n_f[2], n_f[2]), MemBlock(n_f[2], n_f[2]), nn.Identity(), TGrow(n_f[2], 2 if decoder_time_upscale[1] else 1), conv(n_f[2], n_f[3], bias=False), + nn.ReLU(inplace=True), conv(n_f[3], TAEHV.image_channels), + ) + else: + self.decoder = nn.Sequential( + Clamp(), conv(TAEHV.latent_channels, n_f[0]), nn.ReLU(inplace=True), + MemBlock(n_f[0], n_f[0]), MemBlock(n_f[0], n_f[0]), MemBlock(n_f[0], n_f[0]), nn.Upsample(scale_factor=2 if decoder_space_upscale[0] else 1), TGrow(n_f[0], 1), conv(n_f[0], n_f[1], bias=False), + MemBlock(n_f[1], n_f[1]), MemBlock(n_f[1], n_f[1]), MemBlock(n_f[1], n_f[1]), nn.Upsample(scale_factor=2 if decoder_space_upscale[1] else 1), TGrow(n_f[1], 2 if decoder_time_upscale[0] else 1), conv(n_f[1], n_f[2], bias=False), + MemBlock(n_f[2], n_f[2]), MemBlock(n_f[2], n_f[2]), MemBlock(n_f[2], n_f[2]), nn.Upsample(scale_factor=2 if decoder_space_upscale[2] else 1), TGrow(n_f[2], 2 if decoder_time_upscale[1] else 1), conv(n_f[2], n_f[3], bias=False), + nn.ReLU(inplace=True), conv(n_f[3], TAEHV.image_channels), + ) + if checkpoint_path is not None: self.load_state_dict(self.patch_tgrow_layers(torch.load(checkpoint_path, map_location="cpu", weights_only=True)))