Add startup parameters and version check

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)

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)

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

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)