From 4aa115d5c167b03df5ad0fd2e966f5e4bc99233e Mon Sep 17 00:00:00 2001 From: Sakura-Luna <53183413+Sakura-Luna@users.noreply.github.com> Date: Sun, 2 Apr 2023 17:28:44 +0800 Subject: [PATCH 1/4] Add bf16 support. --- modules/processing.py | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/modules/processing.py b/modules/processing.py index 6d9c6a8de..ce0dbbabf 100644 --- a/modules/processing.py +++ b/modules/processing.py @@ -653,8 +653,20 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed: samples_ddim = p.sample(conditioning=c, unconditional_conditioning=uc, seeds=seeds, subseeds=subseeds, subseed_strength=p.subseed_strength, prompts=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))] - for x in x_samples_ddim: - devices.test_for_nans(x, "vae") + try: + for x in x_samples_ddim: + devices.test_for_nans(x, "vae") + except devices.NansException as e: + if not shared.cmd_opts.no_half and not shared.cmd_opts.no_half_vae and torch.cuda.get_device_capability()[0] >= 8: + print('\nA tensor with all NaNs was produced in VAE, try converting to bf16.') + devices.dtype_vae = torch.bfloat16 + vae_file, vae_source = sd_vae.resolve_vae(p.sd_model.sd_model_checkpoint) + 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))] + for x in x_samples_ddim: + devices.test_for_nans(x, "vae") + else: + raise e x_samples_ddim = torch.stack(x_samples_ddim).float() x_samples_ddim = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0) From 157c25f123bfa24aa5d72d700016de5642fcb715 Mon Sep 17 00:00:00 2001 From: Sakura-Luna <53183413+Sakura-Luna@users.noreply.github.com> Date: Sun, 2 Apr 2023 17:41:30 +0800 Subject: [PATCH 2/4] Restore type --- modules/sd_vae.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/modules/sd_vae.py b/modules/sd_vae.py index 9b00f76e9..707d1fb2c 100644 --- a/modules/sd_vae.py +++ b/modules/sd_vae.py @@ -183,6 +183,8 @@ unspecified = object() def reload_vae_weights(sd_model=None, vae_file=unspecified): from modules import lowvram, devices, sd_hijack + if devices.dtype_vae == torch.bfloat16: + devices.dtype_vae = torch.float16 if not sd_model: sd_model = shared.sd_model From d19d227138c6f0448849ca7b19ac8a8e876f249c Mon Sep 17 00:00:00 2001 From: Sakura-Luna <53183413+Sakura-Luna@users.noreply.github.com> Date: Thu, 6 Apr 2023 19:52:18 +0800 Subject: [PATCH 3/4] Add startup parameters and version check --- modules/cmd_args.py | 1 + modules/processing.py | 2 +- modules/sd_vae.py | 2 +- webui.py | 10 ++++++++++ 4 files changed, 13 insertions(+), 2 deletions(-) diff --git a/modules/cmd_args.py b/modules/cmd_args.py index 81c0b82a3..547e8dc89 100644 --- a/modules/cmd_args.py +++ b/modules/cmd_args.py @@ -101,3 +101,4 @@ parser.add_argument("--no-gradio-queue", action='store_true', help="Disables gra parser.add_argument("--skip-version-check", action='store_true', help="Do not check versions of torch and xformers") parser.add_argument("--no-hashing", action='store_true', help="disable sha256 hashing of checkpoints to help loading performance", default=False) parser.add_argument("--no-download-sd-model", action='store_true', help="don't download SD1.5 model even if no model is found in --ckpt-dir", default=False) +parser.add_argument("--rollback-vae", action='store_true', help="trying to roll back vae when produced nan image, need to enable nan check", default=False) diff --git a/modules/processing.py b/modules/processing.py index ce0dbbabf..98402aa57 100644 --- a/modules/processing.py +++ b/modules/processing.py @@ -657,7 +657,7 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed: for x in x_samples_ddim: devices.test_for_nans(x, "vae") except devices.NansException as e: - if not shared.cmd_opts.no_half and not shared.cmd_opts.no_half_vae and torch.cuda.get_device_capability()[0] >= 8: + if not shared.cmd_opts.no_half and not shared.cmd_opts.no_half_vae and shared.cmd_opts.rollback_vae: print('\nA tensor with all NaNs was produced in VAE, try converting to bf16.') devices.dtype_vae = torch.bfloat16 vae_file, vae_source = sd_vae.resolve_vae(p.sd_model.sd_model_checkpoint) diff --git a/modules/sd_vae.py b/modules/sd_vae.py index 707d1fb2c..ee3902a4b 100644 --- a/modules/sd_vae.py +++ b/modules/sd_vae.py @@ -183,7 +183,7 @@ unspecified = object() def reload_vae_weights(sd_model=None, vae_file=unspecified): from modules import lowvram, devices, sd_hijack - if devices.dtype_vae == torch.bfloat16: + if shared.cmd_opts.rollback_vae and devices.dtype_vae == torch.bfloat16: devices.dtype_vae = torch.float16 if not sd_model: sd_model = shared.sd_model diff --git a/webui.py b/webui.py index b570895fb..2f8a3e9fc 100644 --- a/webui.py +++ b/webui.py @@ -97,9 +97,19 @@ To reinstall the desired version, run with commandline flag --reinstall-xformers Use --skip-version-check commandline argument to disable this check. """.strip()) +def check_rollback_vae(): + if shared.cmd_opts.rollback_vae: + if version.parse(torch.__version__) < version.parse('2.1'): + print("If your PyTorch version is lower than PyTorch 2.1, Rollback VAE will not work.") + shared.cmd_opts.rollback_vae = False + elif 0 < torch.cuda.get_device_capability()[0] < 8: + print('Rollback VAE will not work because your device does not support it.') + shared.cmd_opts.rollback_vae = False + def initialize(): check_versions() + check_rollback_vae() extensions.list_extensions() localization.list_localizations(cmd_opts.localizations_dir) From 942c7d6158a160bf675ef8d0ce2630318edb827c Mon Sep 17 00:00:00 2001 From: Sakura-Luna <53183413+Sakura-Luna@users.noreply.github.com> Date: Sat, 8 Apr 2023 23:50:22 +0800 Subject: [PATCH 4/4] Bug fix --- modules/sd_vae.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/modules/sd_vae.py b/modules/sd_vae.py index ee3902a4b..8cff15916 100644 --- a/modules/sd_vae.py +++ b/modules/sd_vae.py @@ -183,8 +183,6 @@ unspecified = object() def reload_vae_weights(sd_model=None, vae_file=unspecified): from modules import lowvram, devices, sd_hijack - if shared.cmd_opts.rollback_vae and devices.dtype_vae == torch.bfloat16: - devices.dtype_vae = torch.float16 if not sd_model: sd_model = shared.sd_model @@ -205,6 +203,8 @@ def reload_vae_weights(sd_model=None, vae_file=unspecified): sd_model.to(devices.cpu) sd_hijack.model_hijack.undo_hijack(sd_model) + if shared.cmd_opts.rollback_vae and devices.dtype_vae == torch.bfloat16: + devices.dtype_vae = torch.float16 load_vae(sd_model, vae_file, vae_source)