return live preview defaults to how they were

only download TAESD model when it's needed
return calculations in single_sample_to_image to just if/elif/elif blocks
keep taesd model in its own directory

modules/sd_samplers_common.py

@@ -22,28 +22,29 @@ def setup_img2img_steps(p, steps=None):
     return steps, t_enc
 
 
-approximation_indexes = {"Full": 0, "Tiny AE": 1, "Approx NN": 2, "Approx cheap": 3}
+approximation_indexes = {"Full": 0, "Approx NN": 1, "Approx cheap": 2, "TAESD": 3}
 
 
 def single_sample_to_image(sample, approximation=None):
-    if approximation is None or approximation not in approximation_indexes.keys():
-        approximation = approximation_indexes.get(opts.show_progress_type, 1)
 
-    if approximation == 1:
-        x_sample = sd_vae_taesd.decode()(sample.to(devices.device, devices.dtype).unsqueeze(0))[0].detach()
-        x_sample = sd_vae_taesd.TAESD.unscale_latents(x_sample)
-        x_sample = torch.clamp((x_sample * 0.25) + 0.5, 0, 1)
-    else:
-        if approximation == 3:
+    if approximation is None:
+        approximation = approximation_indexes.get(opts.show_progress_type, 0)
+
+    if approximation == 2:
         x_sample = sd_vae_approx.cheap_approximation(sample)
-        elif approximation == 2:
+    elif approximation == 1:
         x_sample = sd_vae_approx.model()(sample.to(devices.device, devices.dtype).unsqueeze(0))[0].detach()
+    elif approximation == 3:
+        x_sample = sd_vae_taesd.model()(sample.to(devices.device, devices.dtype).unsqueeze(0))[0].detach()
+        x_sample = sd_vae_taesd.TAESD.unscale_latents(x_sample)  # returns value in [-2, 2]
+        x_sample = x_sample * 0.5
     else:
         x_sample = processing.decode_first_stage(shared.sd_model, sample.unsqueeze(0))[0]
-        x_sample = torch.clamp((x_sample + 1.0) / 2.0, min=0.0, max=1.0)
 
+    x_sample = torch.clamp((x_sample + 1.0) / 2.0, min=0.0, max=1.0)
     x_sample = 255. * np.moveaxis(x_sample.cpu().numpy(), 0, 2)
     x_sample = x_sample.astype(np.uint8)
+
     return Image.fromarray(x_sample)
 
 

modules/sd_vae_taesd.py

@@ -61,16 +61,28 @@ class TAESD(nn.Module):
         return x.sub(TAESD.latent_shift).mul(2 * TAESD.latent_magnitude)
 
 
-def decode():
+def download_model(model_path):
+    model_url = 'https://github.com/madebyollin/taesd/raw/main/taesd_decoder.pth'
+
+    if not os.path.exists(model_path):
+        os.makedirs(os.path.dirname(model_path), exist_ok=True)
+
+        print(f'Downloading TAESD decoder to: {model_path}')
+        torch.hub.download_url_to_file(model_url, model_path)
+
+
+def model():
     global sd_vae_taesd
 
     if sd_vae_taesd is None:
-        model_path = os.path.join(paths_internal.models_path, "VAE-approx", "taesd_decoder.pth")
+        model_path = os.path.join(paths_internal.models_path, "VAE-taesd", "taesd_decoder.pth")
+        download_model(model_path)
+
         if os.path.exists(model_path):
             sd_vae_taesd = TAESD(model_path)
             sd_vae_taesd.eval()
             sd_vae_taesd.to(devices.device, devices.dtype)
         else:
-            raise FileNotFoundError('Tiny AE model not found')
+            raise FileNotFoundError('TAESD model not found')
 
     return sd_vae_taesd.decoder

modules/shared.py

@@ -425,7 +425,7 @@ options_templates.update(options_section(('ui', "Live previews"), {
     "live_previews_enable": OptionInfo(True, "Show live previews of the created image"),
     "show_progress_grid": OptionInfo(True, "Show previews of all images generated in a batch as a grid"),
     "show_progress_every_n_steps": OptionInfo(10, "Show new live preview image every N sampling steps. Set to -1 to show after completion of batch.", gr.Slider, {"minimum": -1, "maximum": 32, "step": 1}),
-    "show_progress_type": OptionInfo("Tiny AE", "Image creation progress preview mode", gr.Radio, {"choices": ["Full", "Tiny AE", "Approx NN", "Approx cheap"]}),
+    "show_progress_type": OptionInfo("Approx NN", "Image creation progress preview mode", gr.Radio, {"choices": ["Full", "Approx NN", "Approx cheap", "TAESD"]}),
     "live_preview_content": OptionInfo("Prompt", "Live preview subject", gr.Radio, {"choices": ["Combined", "Prompt", "Negative prompt"]}),
     "live_preview_refresh_period": OptionInfo(1000, "Progressbar/preview update period, in milliseconds")
 }))

webui.py

@@ -144,21 +144,10 @@ Use --skip-version-check commandline argument to disable this check.
             """.strip())
 
 
-def check_taesd():
-    from modules.paths_internal import models_path
-
-    model_url = 'https://github.com/madebyollin/taesd/raw/main/taesd_decoder.pth'
-    model_path = os.path.join(models_path, "VAE-approx", "taesd_decoder.pth")
-    if not os.path.exists(model_path):
-        print('From taesd repo download decoder model')
-        torch.hub.download_url_to_file(model_url, model_path)
-
-
 def initialize():
     fix_asyncio_event_loop_policy()
 
     check_versions()
-    check_taesd()
 
     extensions.list_extensions()
     localization.list_localizations(cmd_opts.localizations_dir)