redesign live preview and bring full quality toggle to original backend

CHANGELOG.md

@@ -45,6 +45,11 @@ or even free speedups and quality improvements (regardless of which workflows yo
       for example: *"Extra: sampler: Euler a, width: 480, height: 640, steps: 30, cfg scale: 10, clip skip: 2"*
   - **VAE**  
     - VAEs are now also listed as part of extra networks  
+    - Image preview methods have been redesigned: simple, approximate, taesd, full  
+      please set desired preview method in settings  
+    - both original and diffusers backend now support "full quality" setting  
+      if you desired model or platform does not support FP16 and/or you have a low-end hardware and cannot use FP32  
+      you can disable "full quality" in advanced params and it will likely reduce decode errors (infamous black images)  
   - **LoRA**  
     - LoRAs are now automatically filtered based on compatibility with currently loaded model  
       note that if lora type cannot be auto-determined, it will be left in the list  
@@ -152,6 +157,8 @@ or even free speedups and quality improvements (regardless of which workflows yo
     - default updated to *0.0.23*  
     - note that latest xformers are still not compatible with cuda 12.1  
       recommended to use torch 2.1.0 with cuda 11.8  
+      if you attempt to use xformers with cuda 12.1, it will force a full xformers rebuild on install  
+      which can take a very long time and may/may-not work  
     - added cmd param `--use-xformers` to force usage of exformers  
   - **GC**:  
     - custom garbage collect threshold to reduce vram memory usage, thanks @Disty0  

modules/loader.py

@@ -36,4 +36,4 @@ errors.install([gradio])
 
 import diffusers # pylint: disable=W0611,C0411
 timer.startup.record("diffusers")
-errors.log.debug(f'Loaded packages: torch={torch.__version__} diffusers={diffusers.__version__} gradio={gradio.__version__}')
+errors.log.debug(f'Loaded packages: torch={getattr(torch, "__long_version__", torch.__version__)} diffusers={diffusers.__version__} gradio={gradio.__version__}')

modules/processing.py

@@ -35,6 +35,7 @@ import modules.sd_samplers_common
 import modules.sd_models
 import modules.sd_vae
 import modules.sd_vae_approx
+import modules.taesd.sd_vae_taesd
 import modules.generation_parameters_copypaste
 from modules.sd_hijack_hypertile import context_hypertile_vae, context_hypertile_unet, hypertile_set
 
@@ -433,15 +434,25 @@ def create_random_tensors(shape, seeds, subseeds=None, subseed_strength=0.0, see
     return x
 
 
-def decode_first_stage(model, x):
+def decode_first_stage(model, x, full_quality=True):
     with devices.autocast(disable = x.dtype==devices.dtype_vae):
-        if hasattr(model, 'decode_first_stage'):
-            x = model.decode_first_stage(x)
-        elif hasattr(model, 'vae'):
-            x = model.vae(x)
-        else:
-            shared.log.warning('Cannot decode first stage')
-    return x
+        try:
+            if full_quality:
+                if hasattr(model, 'decode_first_stage'):
+                    x_sample = model.decode_first_stage(x)
+                elif hasattr(model, 'vae'):
+                    x_sample = model.vae(x)
+                else:
+                    x_sample = x
+                    shared.log.error('Decode VAE unknown model')
+            else:
+                x_sample = torch.zeros((len(x), 3, x.shape[2] * 8, x.shape[3] * 8), dtype=devices.dtype_vae, device=devices.device)
+                for i in range(len(x_sample)):
+                    x_sample[i] = (modules.taesd.sd_vae_taesd.decode(x[i]) * 2.0) - 1.0
+        except Exception as e:
+            x_sample = x
+            shared.log.error(f'Decode VAE: {e}')
+    return x_sample
 
 
 def get_fixed_seed(seed):
@@ -755,8 +766,6 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
         t0 = time.time()
         with devices.autocast():
             p.init(p.all_prompts, p.all_seeds, p.all_subseeds)
-            if shared.opts.live_previews_enable and shared.opts.show_progress_type == "Approximate NN" and shared.backend == shared.Backend.ORIGINAL:
-                modules.sd_vae_approx.model()
         if shared.state.job_count == -1:
             shared.state.job_count = p.n_iter
         extra_network_data = None
@@ -801,7 +810,7 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
                         comments[comment] = 1
                 with devices.without_autocast() if devices.unet_needs_upcast else devices.autocast():
                     samples_ddim = p.sample(conditioning=c, unconditional_conditioning=uc, seeds=p.seeds, subseeds=p.subseeds, subseed_strength=p.subseed_strength, prompts=p.prompts)
-                x_samples_ddim = [decode_first_stage(p.sd_model, samples_ddim[i:i+1].to(dtype=devices.dtype_vae))[0].cpu() for i in range(samples_ddim.size(0))]
+                x_samples_ddim = [decode_first_stage(p.sd_model, samples_ddim[i:i+1].to(dtype=devices.dtype_vae), p.full_quality)[0].cpu() for i in range(samples_ddim.size(0))]
                 try:
                     for x in x_samples_ddim:
                         devices.test_for_nans(x, "vae")
@@ -811,7 +820,7 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
                         devices.dtype_vae = torch.bfloat16
                         vae_file, vae_source = modules.sd_vae.resolve_vae(p.sd_model.sd_model_checkpoint)
                         modules.sd_vae.load_vae(p.sd_model, vae_file, vae_source)
-                        x_samples_ddim = [decode_first_stage(p.sd_model, samples_ddim[i:i+1].to(dtype=devices.dtype_vae))[0].cpu() for i in range(samples_ddim.size(0))]
+                        x_samples_ddim = [decode_first_stage(p.sd_model, samples_ddim[i:i+1].to(dtype=devices.dtype_vae), p.full_quality)[0].cpu() for i in range(samples_ddim.size(0))]
                         for x in x_samples_ddim:
                             devices.test_for_nans(x, "vae")
                     else:
@@ -1067,7 +1076,7 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
             for i in range(samples.shape[0]):
                 save_intermediate(samples, i)
             if latent_scale_mode is None or self.hr_force: # non-latent upscaling
-                decoded_samples = decode_first_stage(self.sd_model, samples.to(dtype=devices.dtype_vae))
+                decoded_samples = decode_first_stage(self.sd_model, samples.to(dtype=devices.dtype_vae), self.full_quality)
                 lowres_samples = torch.clamp((decoded_samples + 1.0) / 2.0, min=0.0, max=1.0)
                 batch_images = []
                 for _i, x_sample in enumerate(lowres_samples):
@@ -1093,7 +1102,7 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
             else:
                 samples = torch.nn.functional.interpolate(samples, size=(target_height // 8, target_width // 8), mode=latent_scale_mode["mode"], antialias=latent_scale_mode["antialias"])
                 if getattr(self, "inpainting_mask_weight", shared.opts.inpainting_mask_weight) < 1.0:
-                    image_conditioning = self.img2img_image_conditioning(decode_first_stage(self.sd_model, samples.to(dtype=devices.dtype_vae)), samples)
+                    image_conditioning = self.img2img_image_conditioning(decode_first_stage(self.sd_model, samples.to(dtype=devices.dtype_vae), self.full_quality), samples)
                 else:
                     image_conditioning = self.txt2img_image_conditioning(samples.to(dtype=devices.dtype_vae))
                 if self.latent_sampler == "PLMS":

modules/sd_samplers_common.py

@@ -1,13 +1,21 @@
 from collections import namedtuple
-import numpy as np
 import torch
+import torchvision.transforms as T
 from PIL import Image
 from modules import devices, processing, images, sd_vae_approx, sd_samplers, shared
 import modules.taesd.sd_vae_taesd as sd_vae_taesd
 
 
 SamplerData = namedtuple('SamplerData', ['name', 'constructor', 'aliases', 'options'])
-approximation_indexes = {"Full VAE": 0, "Approximate NN": 1, "Approximate simple": 2, "TAESD": 3}
+approximation_indexes = { "Simple": 0, "Approximate": 1, "TAESD": 2, "Full VAE": 3 }
+warned = False
+
+
+def warn_once(message):
+    global warned # pylint: disable=global-statement
+    if not warned:
+        shared.log.warning(message)
+        warned = True
 
 
 def setup_img2img_steps(p, steps=None):
@@ -23,26 +31,33 @@ def setup_img2img_steps(p, steps=None):
 
 
 def single_sample_to_image(sample, approximation=None):
+
     if approximation is None:
-        approximation = approximation_indexes.get(shared.opts.show_progress_type, 0)
-    if approximation == 0:
-        x_sample = processing.decode_first_stage(shared.sd_model, sample.unsqueeze(0))[0] * 0.5 + 0.5
-    elif approximation == 1:
-        x_sample = sd_vae_approx.model()(sample.to(devices.device, devices.dtype).unsqueeze(0))[0].detach() * 0.5 + 0.5
+        approximation = approximation_indexes.get(shared.opts.show_progress_type, None)
+        if approximation is None:
+            warn_once('Unknown decode type, please reset preview method')
+            approximation = 0
+
+    if approximation == 0: # Simple
+        x_sample = sd_vae_approx.cheap_approximation(sample) * 0.5 + 0.5
+    elif approximation == 1: # Approximate
+        x_sample = sd_vae_approx.nn_approximation(sample) * 0.5 + 0.5
         if shared.sd_model_type == "sdxl":
             x_sample = x_sample[[2,1,0],:,:] # BGR to RGB
-    elif approximation == 2:
-        x_sample = sd_vae_approx.cheap_approximation(sample) * 0.5 + 0.5
-    elif approximation == 3:
-        # x_sample = sample * 1.5
-        # x_sample = sd_vae_taesd.model()(x_sample.to(devices.device, devices.dtype).unsqueeze(0))[0].detach()
+    elif approximation == 2: # TAESD
         x_sample = sd_vae_taesd.decode(sample)
+    elif approximation == 3: # Full VAE
+        x_sample = processing.decode_first_stage(shared.sd_model, sample.unsqueeze(0))[0] * 0.5 + 0.5
     else:
-        shared.log.warning(f"Unknown image decode type: {approximation}")
+        warn_once(f"Unknown latent decode type: {approximation}")
         return Image.new(mode="RGB", size=(512, 512))
-    x_sample = torch.clamp(255 * x_sample, min=0.0, max=255).cpu()
-    x_sample = np.moveaxis(x_sample.numpy(), 0, 2).astype(np.uint8)
-    image = Image.fromarray(x_sample)
+
+    try:
+        transform = T.ToPILImage()
+        image = transform(x_sample)
+    except Exception as e:
+        warn_once(f'Transform tensor to image: {e}')
+        image = Image.new(mode="RGB", size=(512, 512))
     return image
 
 
@@ -58,7 +73,7 @@ def images_tensor_to_samples(image, approximation=None, model=None):
     '''image[0, 1] -> latent'''
     if approximation is None:
         approximation = approximation_indexes.get(shared.opts.show_progress_type, 0)
-    if approximation == 3:
+    if approximation == 2:
         image = image.to(devices.device, devices.dtype)
         x_latent = sd_vae_taesd.encode(image)
     else:
@@ -68,10 +83,8 @@ 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:
-            x_latent = torch.stack([
-                model.get_first_stage_encoding(model.encode_first_stage(torch.unsqueeze(img, 0)))[0]
-                for img in image
-            ])
+            image_latents = [model.get_first_stage_encoding(model.encode_first_stage(torch.unsqueeze(img, 0)))[0] for img in image]
+            x_latent = torch.stack(image_latents)
         else:
             x_latent = model.get_first_stage_encoding(model.encode_first_stage(image))
     return x_latent

modules/sd_vae_approx.py

@@ -5,6 +5,8 @@ from modules import devices, paths, shared
 
 
 sd_vae_approx_model = None
+simple_weights = None
+simple_bias = None
 
 
 class VAEApprox(nn.Module):
@@ -32,38 +34,49 @@ class VAEApprox(nn.Module):
         return x
 
 
-def model():
+def nn_approximation(sample): # Approximate NN
     global sd_vae_approx_model # pylint: disable=global-statement
     if sd_vae_approx_model is None:
-        from modules.shared import log
         model_path = os.path.join(paths.models_path, "VAE-approx", "model.pt")
         sd_vae_approx_model = VAEApprox()
         if not os.path.exists(model_path):
             model_path = os.path.join(paths.script_path, "models", "VAE-approx", "model.pt")
-        sd_vae_approx_model.load_state_dict(torch.load(model_path, map_location='cpu' if devices.device.type != 'cuda' else None))
+        approx_weights = torch.load(model_path, map_location='cpu' if devices.device.type != 'cuda' else None)
+        sd_vae_approx_model.load_state_dict(approx_weights)
         sd_vae_approx_model.eval()
         sd_vae_approx_model.to(devices.device, devices.dtype)
-        log.info(f"Loaded VAE-approx: model={model_path}")
-    return sd_vae_approx_model
-
-
-def cheap_approximation(sample):
-    # https://discuss.huggingface.co/t/decoding-latents-to-rgb-without-upscaling/23204/2
-    if shared.sd_model_type == "sdxl":
-        weight = torch.tensor([
-            [0.4543,-0.2868, 0.1566,-0.4748],
-            [0.5008, 0.0952, 0.2155,-0.3268],
-            [0.5294, 0.1625,-0.0624,-0.3793]
-        ]).reshape(3, 4, 1, 1).to(sample.device)
-        bias = torch.tensor([0.1375, 0.0144, -0.0675]).to(sample.device)
-    else:
-        weight = torch.tensor([
-            [0.298, 0.187,-0.158,-0.184],
-            [0.207, 0.286, 0.189,-0.271],
-            [0.208, 0.173, 0.264,-0.473],
-        ]).reshape(3, 4, 1, 1).to(sample.device)
-        bias = None
+        shared.log.debug(f'Loaded VAE decode approximate: model="{model_path}"')
     try:
-        return nn.functional.conv2d(sample, weight, bias) # pylint: disable=not-callable
-    except Exception:
+        in_sample = sample.to(devices.device, devices.dtype).unsqueeze(0)
+        x_sample = sd_vae_approx_model(in_sample)
+        x_sample = x_sample[0]
+        return x_sample
+    except Exception as e:
+        shared.log.error(f'Decode approximate: {e}')
+        return sample
+
+
+def cheap_approximation(sample): # Approximate simple
+    # https://discuss.huggingface.co/t/decoding-latents-to-rgb-without-upscaling/23204/2
+    global simple_weights, simple_bias # pylint: disable=global-statement
+    if simple_weights is None or simple_bias is None:
+        if shared.sd_model_type == "sdxl":
+            simple_weights = torch.tensor([
+                [0.4543,-0.2868, 0.1566,-0.4748],
+                [0.5008, 0.0952, 0.2155,-0.3268],
+                [0.5294, 0.1625,-0.0624,-0.3793]
+            ]).reshape(3, 4, 1, 1).to(sample.device)
+            simple_bias = torch.tensor([0.1375, 0.0144, -0.0675]).to(sample.device)
+        else:
+            simple_weights = torch.tensor([
+                [0.298, 0.187,-0.158,-0.184],
+                [0.207, 0.286, 0.189,-0.271],
+                [0.208, 0.173, 0.264,-0.473],
+            ]).reshape(3, 4, 1, 1).to(sample.device)
+            simple_bias = None
+    try:
+        x_sample = nn.functional.conv2d(sample, simple_weights, simple_bias) # pylint: disable=not-callable
+        return x_sample
+    except Exception as e:
+        shared.log.error(f'Decode simple: {e}')
         return sample

modules/shared.py

@@ -600,7 +600,7 @@ options_templates.update(options_section(('live-preview', "Live Previews"), {
     "notification_audio_enable": OptionInfo(False, "Play a sound when images are finished generating"),
     "notification_audio_path": OptionInfo("html/notification.mp3","Path to notification sound", component_args=hide_dirs, folder=True),
     "show_progress_every_n_steps": OptionInfo(1, "Live preview display period", gr.Slider, {"minimum": -1, "maximum": 32, "step": 1}),
-    "show_progress_type": OptionInfo("Approximate NN", "Live preview method", gr.Radio, {"choices": ["Full VAE", "Approximate NN", "Approximate simple", "TAESD"]}),
+    "show_progress_type": OptionInfo("Approximate", "Live preview method", gr.Radio, {"choices": ["Simple", "Approximate", "TAESD", "Full VAE"]}),
     "live_preview_content": OptionInfo("Combined", "Live preview subject", gr.Radio, {"choices": ["Combined", "Prompt", "Negative prompt"], "visible": False}),
     "live_preview_refresh_period": OptionInfo(500, "Progress update period", gr.Slider, {"minimum": 0, "maximum": 5000, "step": 25}),
     "logmonitor_show": OptionInfo(True, "Show log view"),

modules/taesd/sd_vae_taesd.py

@@ -62,6 +62,7 @@ def decode(latents):
     image = vae.decoder(enc).clamp(0, 1).detach()
     return image[0]
 
+
 def encode(image):
     from modules import shared
     model_class = shared.sd_model_type