vae

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)

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)

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"