offloading improvements

Signed-off-by: Vladimir Mandic <mandic00@live.com>

modules/sd_hijack_te.py

@@ -11,6 +11,9 @@ def hijack_encode_prompt(*args, **kwargs):
     # if hasattr(shared.sd_model, 'text_encoder') and shared.sd_model.text_encoder is not None:
     #     sd_models.move_model(shared.sd_model.text_encoder, devices.device)
     try:
+        prompt = kwargs.get('prompt', None) or (args[0] if len(args) > 0 else None)
+        if prompt is not None:
+            shared.log.debug(f'Encode: prompt="{prompt}" hijack=True')
         res = shared.sd_model.orig_encode_prompt(*args, **kwargs)
     except Exception as e:
         shared.log.error(f'Encode prompt: {e}')

modules/sd_offload.py

@@ -211,7 +211,7 @@ class OffloadHook(accelerate.hooks.ModelHook):
                     for module_name in get_module_names(pipe):
                         module_instance = getattr(pipe, module_name, None)
                         module_cls = module_instance.__class__.__name__
-                        if (module_cls != module.__class__.__name__) and (module_cls not in self.offload_never) and (not devices.same_device(module_instance.device, devices.cpu)):
+                        if (id(module) != id(module_instance)) and (module_cls not in self.offload_never) and (not devices.same_device(module_instance.device, devices.cpu)):
                             apply_balanced_offload_to_module(module_instance, op='pre')
 
         if not devices.same_device(module.device, devices.device):

modules/shared_defaults.py

@@ -20,10 +20,12 @@ def get_default_modes(cmd_opts, mem_stat):
                 cmd_opts.medvram = True # VAE Tiling and other stuff
                 default_offload_mode = "balanced"
                 default_diffusers_offload_min_gpu_memory = 0
+                default_diffusers_offload_always = ', '.join(['T5EncoderModel', 'UMT5EncoderModel'])
                 log.info(f"Device detect: memory={gpu_memory:.1f} default=balanced optimization=medvram")
             elif gpu_memory >= 24:
                 default_offload_mode = "balanced"
                 default_diffusers_offload_max_gpu_memory = 0.8
+                default_diffusers_offload_always = ', '.join(['T5EncoderModel', 'UMT5EncoderModel'])
                 default_diffusers_offload_never = ', '.join(['CLIPTextModel', 'CLIPTextModelWithProjection', 'AutoencoderKL'])
                 log.info(f"Device detect: memory={gpu_memory:.1f} default=balanced optimization=highvram")
             else: