From bb912c35f9c2a4d1deb078e35e4d275642a0e815 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Mon, 16 Oct 2023 14:07:00 -0400 Subject: [PATCH] redesign live preview and bring full quality toggle to original backend --- CHANGELOG.md | 7 +++ extensions-builtin/sd-extension-chainner | 2 +- modules/loader.py | 2 +- modules/processing.py | 37 ++++++++------ modules/sd_samplers_common.py | 55 +++++++++++++-------- modules/sd_vae_approx.py | 63 ++++++++++++++---------- modules/shared.py | 2 +- modules/taesd/sd_vae_taesd.py | 1 + 8 files changed, 106 insertions(+), 63 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 857177ea5..cfc13e2d6 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/extensions-builtin/sd-extension-chainner b/extensions-builtin/sd-extension-chainner index aa3050d96..656293e7b 160000 --- a/extensions-builtin/sd-extension-chainner +++ b/extensions-builtin/sd-extension-chainner @@ -1 +1 @@ -Subproject commit aa3050d96245430728b682c6b3a5b26a0c475673 +Subproject commit 656293e7bf4c2c462f7904dd2b5b03a244989a39 diff --git a/modules/loader.py b/modules/loader.py index 9733b735b..c0cc8dea3 100644 --- a/modules/loader.py +++ b/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__}') diff --git a/modules/processing.py b/modules/processing.py index 6705db2a1..a59680660 100644 --- a/modules/processing.py +++ b/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": diff --git a/modules/sd_samplers_common.py b/modules/sd_samplers_common.py index 71a0b2840..75198b832 100644 --- a/modules/sd_samplers_common.py +++ b/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 diff --git a/modules/sd_vae_approx.py b/modules/sd_vae_approx.py index a12d103ce..f5c577458 100644 --- a/modules/sd_vae_approx.py +++ b/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 diff --git a/modules/shared.py b/modules/shared.py index ad4db47e9..9dc7e7720 100644 --- a/modules/shared.py +++ b/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"), diff --git a/modules/taesd/sd_vae_taesd.py b/modules/taesd/sd_vae_taesd.py index 35d71a776..ff50b4156 100644 --- a/modules/taesd/sd_vae_taesd.py +++ b/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