Merge pull request #438 from Sakura-Luna/pr9295

Add bf16 support for VAE

modules/cmd_args.py

@@ -44,6 +44,7 @@ parser.add_argument("--no-hashing", action='store_true', help="Disable sha256 ha
 parser.add_argument("--no-download-sd-model", action='store_true', help="Disable download of default model even if no model is found", default=False)
 parser.add_argument("--profile", action='store_true', help="Run profiler, default: %(default)s")
 parser.add_argument("--disable-queue", action='store_true', help="Disable Gradio queues and force use of HTTP instead of WebSockets, default: %(default)s")
+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)
 parser.add_argument("--token-merging", action='store_true', help="Provides speed and memory improvements by merging redundant tokens. This has a more pronounced effect on higher resolutions.", default=False)
 
 

modules/processing.py

@@ -689,8 +689,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 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)
+                    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)

modules/sd_vae.py

@@ -197,6 +197,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)
 

webui.py

@@ -64,7 +64,19 @@ else:
     server_name = "0.0.0.0" if cmd_opts.listen else None
 
 
+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_rollback_vae()
+
     extensions.list_extensions()
     startup_timer.record("extensions")