More offload checks

Author: Disty0

modules/model_stablecascade.py

@@ -2,7 +2,7 @@
 import copy
 import torch
 import diffusers
-from modules import shared, devices
+from modules import shared, devices, sd_models
 
 def get_timestep_ratio_conditioning(t, alphas_cumprod):
     s = torch.tensor([0.008]) # diffusers uses 0.003 while the original is 0.008
@@ -191,7 +191,8 @@ def __call__(
         callback_on_step_end=None,
         callback_on_step_end_tensor_inputs=["latents"],
     ):
-
+        if shared.opts.diffusers_offload_mode == "balanced":
+            shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model)
         # 0. Define commonly used variables
         self.guidance_scale = guidance_scale
         self.do_classifier_free_guidance = self.guidance_scale > 1
@@ -334,6 +335,8 @@ def __call__(
 
         # Offload all models
         self.maybe_free_model_hooks()
+        if shared.opts.diffusers_offload_mode == "balanced":
+            shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model)
 
         if not return_dict:
             return images

modules/processing_vae.py

@@ -69,7 +69,9 @@ def full_vae_decode(latents, model):
             model.vae.apply(sd_models.convert_to_faketensors)
             devices.torch_gc(force=True)
 
-    if shared.opts.diffusers_move_unet and not getattr(model, 'has_accelerate', False) and base_device is not None:
+    if shared.opts.diffusers_offload_mode == "balanced":
+        shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model)
+    elif shared.opts.diffusers_move_unet and not getattr(model, 'has_accelerate', False) and base_device is not None:
         sd_models.move_base(model, base_device)
     t1 = time.time()
     debug(f'VAE decode: name={sd_vae.loaded_vae_file if sd_vae.loaded_vae_file is not None else "baked"} dtype={model.vae.dtype} upcast={upcast} images={latents.shape[0]} latents={latents.shape} time={round(t1-t0, 3)}')