Add bf16 support.

This commit is contained in:
Sakura-Luna
2023-04-02 17:28:44 +08:00
committed by GitHub
parent 22bcc7be42
commit 4aa115d5c1
+14 -2
View File
@@ -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)