From 93104ea456f4dc014707d3a46961eca85d468d79 Mon Sep 17 00:00:00 2001 From: Haoming Date: Mon, 1 Sep 2025 15:21:46 +0800 Subject: [PATCH] vae --- modules/sd_samplers_common.py | 39 ++++++++++++++++------------------- modules/sd_vae_approx.py | 31 +++++++++++++++++++++++----- modules/sd_vae_taesd.py | 2 ++ 3 files changed, 46 insertions(+), 26 deletions(-) diff --git a/modules/sd_samplers_common.py b/modules/sd_samplers_common.py index a8f19832..94987297 100644 --- a/modules/sd_samplers_common.py +++ b/modules/sd_samplers_common.py @@ -42,27 +42,24 @@ def samples_to_images_tensor(sample, approximation=None, model=None): if approximation is None or (shared.state.interrupted and opts.live_preview_fast_interrupt): approximation = approximation_indexes.get(opts.show_progress_type, 0) if approximation == 0: - approximation = 1 + approximation = 2 + + if approximation == 1: + if (mdl := sd_vae_approx.model()) is not None: + x_sample = mdl(sample.to(devices.device, devices.dtype)).detach() + else: + approximation = 2 + elif approximation == 3: + if (mdl := sd_vae_taesd.decoder_model()) is not None: + x_sample = mdl(sample.to(devices.device, devices.dtype)).detach() + x_sample = x_sample * 2 - 1 + else: + approximation = 2 if approximation == 2: - x_sample = sd_vae_approx.cheap_approximation(sample) - elif approximation == 1: - m = sd_vae_approx.model() - if m is None: - x_sample = sd_vae_approx.cheap_approximation(sample) - else: - x_sample = m(sample.to(devices.device, devices.dtype)).detach() - elif approximation == 3: - m = sd_vae_taesd.decoder_model() - if m is None: - x_sample = sd_vae_approx.cheap_approximation(sample) - else: - x_sample = m(sample.to(devices.device, devices.dtype)).detach() - x_sample = x_sample * 2 - 1 + x_sample = sd_vae_approx.cheap_approximation(sample).detach() else: - if model is None: - model = shared.sd_model - x_sample = model.decode_first_stage(sample) + x_sample = (model or shared.sd_model).decode_first_stage(sample) return x_sample @@ -107,7 +104,7 @@ def images_tensor_to_samples(image, approximation=None, model=None): image = image.to(shared.device, dtype=devices.dtype_vae) image = image * 2 - 1 - if len(image) > 1: + if len(image) > 1 and not model.is_wan: x_latent = torch.stack([ model.get_first_stage_encoding( model.encode_first_stage(torch.unsqueeze(img, 0)) @@ -193,13 +190,13 @@ def apply_refiner(cfg_denoiser, x): with sd_models.SkipWritingToConfig(): fp_checkpoint = getattr(shared.opts, 'sd_model_checkpoint') - checkpoint_changed = main_entry.checkpoint_change(refiner_checkpoint_info.short_title, save=False, refresh=False) + checkpoint_changed = main_entry.checkpoint_change(refiner_checkpoint_info.short_title, preset=None, save=False, refresh=False) if checkpoint_changed: try: main_entry.refresh_model_loading_parameters() sd_models.forge_model_reload() finally: - main_entry.checkpoint_change(fp_checkpoint, save=False, refresh=True) + main_entry.checkpoint_change(fp_checkpoint, preset=None, save=False, refresh=True) if not cfg_denoiser.p.disable_extra_networks: extra_networks.activate(cfg_denoiser.p, cfg_denoiser.p.extra_network_data) diff --git a/modules/sd_vae_approx.py b/modules/sd_vae_approx.py index 013f164e..144530bf 100644 --- a/modules/sd_vae_approx.py +++ b/modules/sd_vae_approx.py @@ -1,7 +1,9 @@ import os +from functools import cache import torch from torch import nn + from modules import devices, paths, shared sd_vae_approx_models = {} @@ -24,7 +26,7 @@ class VAEApprox(nn.Module): x = nn.functional.interpolate(x, (x.shape[2] * 2, x.shape[3] * 2)) x = nn.functional.pad(x, (extra, extra, extra, extra)) - for layer in [self.conv1, self.conv2, self.conv3, self.conv4, self.conv5, self.conv6, self.conv7, self.conv8, ]: + for layer in [self.conv1, self.conv2, self.conv3, self.conv4, self.conv5, self.conv6, self.conv7, self.conv8]: x = layer(x) x = nn.functional.leaky_relu(x, 0.1) @@ -35,7 +37,7 @@ def download_model(model_path, model_url): if not os.path.exists(model_path): os.makedirs(os.path.dirname(model_path), exist_ok=True) - print(f'Downloading VAEApprox model to: {model_path}') + print(f"Downloading VAEApprox model to: {model_path}") torch.hub.download_url_to_file(model_url, model_path) @@ -57,10 +59,10 @@ def model(): if not os.path.exists(model_path): model_path = os.path.join(paths.models_path, "VAE-approx", model_name) - download_model(model_path, 'https://github.com/AUTOMATIC1111/stable-diffusion-webui/releases/download/v1.0.0-pre/' + model_name) + download_model(model_path, "https://github.com/AUTOMATIC1111/stable-diffusion-webui/releases/download/v1.0.0-pre/" + model_name) loaded_model = VAEApprox(latent_channels=shared.sd_model.forge_objects.vae.latent_channels) - loaded_model.load_state_dict(torch.load(model_path, map_location='cpu' if devices.device.type != 'cuda' else None)) + loaded_model.load_state_dict(torch.load(model_path, map_location="cpu" if devices.device.type != "cuda" else None)) loaded_model.eval() loaded_model.to(devices.device, devices.dtype) sd_vae_approx_models[model_name] = loaded_model @@ -68,5 +70,24 @@ def model(): return loaded_model +@cache +def latent_format(latent_format, dtype, device): + factors = torch.tensor(latent_format.latent_rgb_factors, dtype=dtype, device=device).transpose(0, 1) + + if getattr(latent_format, "latent_rgb_factors_bias", None) is None: + return factors, None + + bias = torch.tensor(latent_format.latent_rgb_factors_bias, dtype=dtype, device=device) + return factors, bias + + def cheap_approximation(sample): - return torch.einsum("...lxy,lr -> ...rxy", sample, torch.tensor(shared.sd_model.model_config.latent_format.latent_rgb_factors).to(sample.device)) + factors, bias = latent_format(shared.sd_model.model_config.latent_format, sample.dtype, sample.device) + + if sample.ndim == 5: + sample = sample[0, :, 0] + else: + sample = sample[0] + + latent_image = torch.nn.functional.linear(sample.movedim(0, -1), factors, bias=bias) + return latent_image.permute(2, 0, 1).unsqueeze(0) diff --git a/modules/sd_vae_taesd.py b/modules/sd_vae_taesd.py index c1394cec..87605954 100644 --- a/modules/sd_vae_taesd.py +++ b/modules/sd_vae_taesd.py @@ -102,6 +102,8 @@ def download_model(model_path, model_url): def decoder_model(): if not shared.sd_model.is_webui_legacy_model(): + if shared.sd_model.is_wan: + return None model_name = "taef1_decoder.pth" elif shared.sd_model.is_sdxl: model_name = "taesdxl_decoder.pth"