improve live preview

modules/sd_samplers_common.py

@@ -14,7 +14,7 @@ warned = False
 def warn_once(message):
     global warned # pylint: disable=global-statement
     if not warned:
-        shared.log.warning(message)
+        shared.log.warning(f'VAE: {message}')
         warned = True
 
 
@@ -34,9 +34,8 @@ def single_sample_to_image(sample, approximation=None):
     if approximation is None:
         approximation = approximation_indexes.get(shared.opts.show_progress_type, None)
         if approximation is None:
-            warn_once('Unknown decode type, please reset preview method')
+            warn_once('Unknown decode type')
             approximation = 0
-
     # normal sample is [4,64,64]
     if sample.dtype == torch.bfloat16:
         sample = sample.to(torch.float16)
@@ -58,14 +57,13 @@ def single_sample_to_image(sample, approximation=None):
     else:
         warn_once(f"Unknown latent decode type: {approximation}")
         return Image.new(mode="RGB", size=(512, 512))
-
     try:
         if x_sample.dtype == torch.bfloat16:
             x_sample.to(torch.float16)
         transform = T.ToPILImage()
         image = transform(x_sample)
     except Exception as e:
-        warn_once(f'Live preview: {e}')
+        warn_once(f'live preview: {e}')
         image = Image.new(mode="RGB", size=(512, 512))
     return image
 

modules/sd_vae_approx.py

@@ -43,7 +43,7 @@ def nn_approximation(sample): # Approximate NN
         sd_vae_approx_model.load_state_dict(approx_weights)
         sd_vae_approx_model.eval()
         sd_vae_approx_model.to(devices.device, sample.dtype)
-        shared.log.debug(f'Load VAE decode approximate: model="{model_path}"')
+        shared.log.debug(f'VAE load: type=approximate model={model_path}')
     try:
         in_sample = sample.to(devices.device).unsqueeze(0)
         sd_vae_approx_model.to(devices.device, devices.dtype)
@@ -51,7 +51,7 @@ def nn_approximation(sample): # Approximate NN
         x_sample = x_sample[0].detach().cpu()
         return x_sample
     except Exception as e:
-        shared.log.error(f'Decode approximate: {e}')
+        shared.log.error(f'VAE decode approximate: {e}')
         return sample
 
 
@@ -77,5 +77,5 @@ def cheap_approximation(sample): # Approximate simple
         x_sample = nn.functional.conv2d(sample, weights, bias) # pylint: disable=not-callable
         return x_sample
     except Exception as e:
-        shared.log.error(f'Decode simple: {e}')
+        shared.log.error(f'VAE decode simple: {e}')
         return sample

modules/taesd/sd_vae_taesd.py

@@ -15,10 +15,10 @@ def download_model(model_path):
     model_name = os.path.basename(model_path)
     model_url = f'https://github.com/madebyollin/taesd/raw/main/{model_name}'
     if not os.path.exists(model_path):
-        os.makedirs(os.path.dirname(model_path), exist_ok=True)
-        from modules.shared import log
-        log.info(f'Downloading TAESD decoder: {model_path}')
         import torch
+        from modules.shared import log
+        os.makedirs(os.path.dirname(model_path), exist_ok=True)
+        log.info(f'Downloading TAESD decoder: {model_path}')
         torch.hub.download_url_to_file(model_url, model_path)
 
 
@@ -56,20 +56,25 @@ def decode(latents):
         download_model(model_path)
         if os.path.exists(model_path):
             taesd_models[f'{model_class}-decoder'] = TAESD(decoder_path=model_path, encoder_path=None)
+            shared.log.debug(f'VAE load: type=taesd model={model_path}')
             vae = taesd_models[f'{model_class}-decoder']
             vae.to(devices.device, devices.dtype_vae)
     latents.to(devices.device, devices.dtype_vae)
-    if len(latents.shape) == 3:
-        latents = latents.unsqueeze(0)
-        image = vae.decoder(latents).clamp(0, 1).detach()
-        image = 2.0 * image - 1.0 # typical normalized range except for preview which runs denormalization
-        return image[0]
-    elif len(latents.shape) == 4:
-        image = vae.decoder(latents).clamp(0, 1).detach()
-        image = 2.0 * image - 1.0 # typical normalized range except for preview which runs denormalization
-        return image
-    else:
-        shared.log.error(f'TAESD decode unsupported latent type: {latents.shape}')
+    try:
+        if len(latents.shape) == 3:
+            latents = latents.unsqueeze(0)
+            image = vae.decoder(latents).clamp(0, 1).detach()
+            image = 2.0 * image - 1.0 # typical normalized range except for preview which runs denormalization
+            return image[0]
+        elif len(latents.shape) == 4:
+            image = vae.decoder(latents).clamp(0, 1).detach()
+            image = 2.0 * image - 1.0 # typical normalized range except for preview which runs denormalization
+            return image
+        else:
+            shared.log.error(f'TAESD decode unsupported latent type: {latents.shape}')
+            return latents
+    except Exception as e:
+        shared.log.error(f'VAE decode taesd: {e}')
         return latents
 
 
@@ -86,6 +91,7 @@ def encode(image):
         model_path = os.path.join(paths.models_path, "TAESD", f"tae{model_class}_encoder.pth")
         download_model(model_path)
         if os.path.exists(model_path):
+            shared.log.debug(f'VAE load: type=taesd model={model_path}')
             taesd_models[f'{model_class}-encoder'] = TAESD(encoder_path=model_path, decoder_path=None)
             vae = taesd_models[f'{model_class}-encoder']
             vae.to(devices.device, devices.dtype_vae)