Add Move UNet to CPU option

modules/processing_diffusers.py

@@ -36,9 +36,16 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro
     def vae_decode(latents, model, output_type='np'):
         if hasattr(model, 'vae') and torch.is_tensor(latents):
             shared.log.debug(f'Diffusers VAE decode: name={model.vae.config.get("_name_or_path", "default")} dtype={model.vae.dtype} upcast={model.vae.config.get("force_upcast", None)}')
+            if shared.opts.diffusers_move_unet:
+                shared.log.debug('Diffusers: Moving UNet to CPU')
+                unet_device = model.unet.device
+                model.unet.to(devices.cpu)
+                devices.torch_gc()
             latents.to(model.vae.device)
             decoded = model.vae.decode(latents / model.vae.config.scaling_factor, return_dict=False)[0]
             imgs = model.image_processor.postprocess(decoded, output_type=output_type)
+            if shared.opts.diffusers_move_unet:
+                model.unet.to(unet_device)
             return imgs
         else:
             return latents

modules/shared.py

@@ -398,6 +398,7 @@ options_templates.update(options_section(('diffusers', "Diffusers Settings"), {
     "diffusers_refiner_latents": OptionInfo(True, "Use latents when using refiner"),
     "diffusers_move_base": OptionInfo(False, "Move base model to CPU when using refiner"),
     "diffusers_move_refiner": OptionInfo(True, "Move refiner model to CPU when not in use"),
+    "diffusers_move_unet": OptionInfo(False, "Move UNet to CPU while VAE decoding"),
     "diffusers_extract_ema": OptionInfo(True, "Use model EMA weights when possible"),
     "diffusers_generator_device": OptionInfo("default", "Generator device", gr.Radio, lambda: {"choices": ["default", "cpu"]}),
     "diffusers_seq_cpu_offload": OptionInfo(False, "Enable sequential CPU offload"),