From cc89ed84f9578e5e81d765c4324cf8b1884a4b03 Mon Sep 17 00:00:00 2001
From: Disty0 <disty@disty.xyz>
Date: Mon, 19 Aug 2024 18:48:09 +0300
Subject: [PATCH] Rename cpu offload to model offload

---
 modules/model_pixart.py | 2 +-
 modules/model_t5.py     | 2 +-
 modules/sd_models.py    | 2 +-
 modules/shared.py       | 6 +++---
 4 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/modules/model_pixart.py b/modules/model_pixart.py
index f7918c54b..560fcccd9 100644
--- a/modules/model_pixart.py
+++ b/modules/model_pixart.py
@@ -5,7 +5,7 @@ def load_pixart(checkpoint_info, diffusers_load_config={}):
     from modules import shared, devices, modelloader, model_t5
     modelloader.hf_login()
     # shared.opts.data['cuda_dtype'] = 'FP32' # override
-    # shared.opts.data['diffusers_offload_mode}'] = "cpu" # override
+    # shared.opts.data['diffusers_offload_mode}'] = "model" # override
     # devices.set_cuda_params()
     fn = checkpoint_info.path.replace('huggingface/', '')
     t5 = model_t5.load_t5(shared.opts.sd_text_encoder, cache_dir=shared.opts.diffusers_dir)
diff --git a/modules/model_t5.py b/modules/model_t5.py
index e5201e675..d79a27c34 100644
--- a/modules/model_t5.py
+++ b/modules/model_t5.py
@@ -83,7 +83,7 @@ def set_t5(pipe, module, t5=None, cache_dir=None):
         from accelerate import cpu_offload
         getattr(pipe, module).to("cpu")
         cpu_offload(getattr(pipe, module), devices.device, offload_buffers=len(getattr(pipe, module)._parameters) > 0) # pylint: disable=protected-access
-    elif shared.opts.diffusers_offload_mode == "cpu":
+    elif shared.opts.diffusers_offload_mode == "model":
         if not hasattr(pipe, "_all_hooks") or len(pipe._all_hooks) == 0: # pylint: disable=protected-access
             pipe.enable_model_cpu_offload(device=devices.device)
     if hasattr(pipe, "maybe_free_model_hooks"):
diff --git a/modules/sd_models.py b/modules/sd_models.py
index 12533bdbe..be98c2f1f 100644
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -732,7 +732,7 @@ def set_diffuser_offload(sd_model, op: str = 'model'):
     if not (hasattr(sd_model, "has_accelerate") and sd_model.has_accelerate):
         sd_model.has_accelerate = False
     if hasattr(sd_model, "enable_model_cpu_offload"):
-        if shared.opts.diffusers_offload_mode == "cpu":
+        if shared.opts.diffusers_offload_mode == "model":
             shared.log.debug(f'Setting {op}: enable model CPU offload')
             if shared.opts.diffusers_move_base or shared.opts.diffusers_move_unet or shared.opts.diffusers_move_refiner:
                 shared.opts.diffusers_move_base = False
diff --git a/modules/shared.py b/modules/shared.py
index 852db5a2d..a839adcc8 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -369,13 +369,13 @@ if not (cmd_opts.lowvram or cmd_opts.medvram):
             log.info(f"VRAM: Detected={gpu_memory} GB Optimization=lowvram")
         elif gpu_memory <= 8:
             cmd_opts.medvram = True
-            offload_mode_default = "cpu"
+            offload_mode_default = "model"
             log.info(f"VRAM: Detected={gpu_memory} GB Optimization=medvram")
         else:
             offload_mode_default = "none"
             log.info(f"VRAM: Detected={gpu_memory} GB Optimization=none")
 elif cmd_opts.medvram:
-    offload_mode_default = "cpu"
+    offload_mode_default = "model"
 elif cmd_opts.lowvram:
     offload_mode_default = "sequential"
 
@@ -542,7 +542,7 @@ options_templates.update(options_section(('diffusers', "Diffusers Settings"), {
     "diffusers_move_refiner": OptionInfo(False, "Move refiner model to CPU when not in use"),
     "diffusers_extract_ema": OptionInfo(False, "Use model EMA weights when possible"),
     "diffusers_generator_device": OptionInfo("GPU", "Generator device", gr.Radio, {"choices": ["GPU", "CPU", "Unset"]}),
-    "diffusers_offload_mode": OptionInfo(offload_mode_default, "Model offload mode", gr.Radio, {"choices": ['none', 'balanced', 'cpu', 'sequential']}),
+    "diffusers_offload_mode": OptionInfo(offload_mode_default, "Model offload mode", gr.Radio, {"choices": ['none', 'balanced', 'model', 'sequential']}),
     "diffusers_offload_max_gpu_memory": OptionInfo(gpu_memory * 0.75, "Max GPU memory for balanced offload mode in GB", gr.Slider, {"minimum": 0, "maximum": gpu_memory, "step": 0.1,}),
     "diffusers_offload_max_cpu_memory": OptionInfo(cpu_memory * 0.75, "Max CPU memory for balanced offload mode in GB", gr.Slider, {"minimum": 0, "maximum": cpu_memory, "step": 0.1,}),
     "diffusers_vae_upcast": OptionInfo("default", "VAE upcasting", gr.Radio, {"choices": ['default', 'true', 'false']}),