From 682da607a969600dcfa0ca9c031c250f73286b80 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Mon, 29 Jan 2024 08:53:48 -0500 Subject: [PATCH] improve live preview --- modules/sd_samplers_common.py | 8 +++----- modules/sd_vae_approx.py | 6 +++--- modules/taesd/sd_vae_taesd.py | 34 ++++++++++++++++++++-------------- 3 files changed, 26 insertions(+), 22 deletions(-) diff --git a/modules/sd_samplers_common.py b/modules/sd_samplers_common.py index 86c090324..08571a063 100644 --- a/modules/sd_samplers_common.py +++ b/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 diff --git a/modules/sd_vae_approx.py b/modules/sd_vae_approx.py index 0d8caf6d9..e66b78011 100644 --- a/modules/sd_vae_approx.py +++ b/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 diff --git a/modules/taesd/sd_vae_taesd.py b/modules/taesd/sd_vae_taesd.py index 5b38ed259..47330d4b8 100644 --- a/modules/taesd/sd_vae_taesd.py +++ b/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)