diff --git a/html/locale_en.json b/html/locale_en.json
index 790d09adc..6a13e9cd4 100644
--- a/html/locale_en.json
+++ b/html/locale_en.json
@@ -565,6 +565,7 @@
{"id":"","label":"Select diffuser pipeline when loading from safetensors","localized":"","hint":""},
{"id":"","label":"Move base model to CPU when using refiner","localized":"","hint":""},
{"id":"","label":"Move refiner model to CPU when not in use","localized":"","hint":""},
+ {"id":"","label":"Move UNet to CPU while VAE decoding","localized":"","hint":""},
{"id":"","label":"Use model EMA weights when possible","localized":"","hint":""},
{"id":"","label":"Generator device","localized":"","hint":""},
{"id":"","label":"Enable sequential CPU offload","localized":"","hint":"Reduces GPU memory usage by transferring weights to the CPU. Increases inference time approximately 10%. Use with Enable Attention slicing for minimal memory consumption"},
diff --git a/html/locale_ko.json b/html/locale_ko.json
index 1b0a05f8e..b2a2807d0 100644
--- a/html/locale_ko.json
+++ b/html/locale_ko.json
@@ -565,6 +565,7 @@
{"id":"","label":"Select diffuser pipeline when loading from safetensors","localized":"safetensors 파일에서 로드할 때 사용할 파이프라인 선택","hint":""},
{"id":"","label":"Move base model to CPU when using refiner","localized":"리파이너를 사용 중일 때 base 모델을 CPU로 이동","hint":""},
{"id":"","label":"Move refiner model to CPU when not in use","localized":"사용 중이지 않을 때 리파이너 모델을 CPU로 이동","hint":""},
+ {"id":"","label":"Move UNet to CPU while VAE decoding","localized":"","hint":""},
{"id":"","label":"Use model EMA weights when possible","localized":"가능하다면 모델의 EMA 가중치 사용","hint":""},
{"id":"","label":"Generator device","localized":"Generator 디바이스","hint":""},
{"id":"","label":"Enable sequential CPU offload","localized":"순차 CPU 오프로드 활성화","hint":"가중치를 CPU로 옮겨 GPU 메모리 사용률을 낮춘다. 속도가 약 10% 느려진다. Use with 어텐션 슬라이싱 활성화 for minimal memory consumption"},
diff --git a/modules/processing_diffusers.py b/modules/processing_diffusers.py
index 52175889c..ce1457967 100644
--- a/modules/processing_diffusers.py
+++ b/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
diff --git a/modules/shared.py b/modules/shared.py
index df7a23771..711a77802 100644
--- a/modules/shared.py
+++ b/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"),