add monitor cli option and finish lora refactor

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

CHANGELOG.md

@@ -122,6 +122,7 @@ Plus...
   - add Flash Attention 2 support under [triton for ZLUDA v3.9.2](https://github.com/vladmandic/sdnext/wiki/ZLUDA#how-to-enable-triton)  
   - add Sage Attention support  
 - **Other**  
+  - new command line option `--monitor PERIOD` to monitor CPU and GPU memory ever n seconds  
   - **upscale**: new [asymmetric vae v2](https://huggingface.co/Heasterian/AsymmetricAutoencoderKLUpscaler_v2) upscaling method  
   - **upscale**: new experimental support for `libvips` upscaling  
   - **quantization**: add support for `optimum-quanto` on-the-fly quantization during load for all models  

installer.py

@@ -517,7 +517,7 @@ def check_python(supported_minors=[9, 10, 11, 12], reason=None):
         log.error(f"Python version incompatible: {sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro} required 3.{supported_minors}")
         if reason is not None:
             log.error(reason)
-        if not args.ignore:
+        if not args.ignore and not args.experimental:
             sys.exit(1)
     if int(sys.version_info.minor) == 12:
         os.environ.setdefault('SETUPTOOLS_USE_DISTUTILS', 'local') # hack for python 3.11 setuptools
@@ -1492,6 +1492,7 @@ def add_args(parser):
     group_log.add_argument("--log", type=str, default=os.environ.get("SD_LOG", None), help="Set log file, default: %(default)s")
     group_log.add_argument('--debug', default=os.environ.get("SD_DEBUG",False), action='store_true', help="Run installer with debug logging, default: %(default)s")
     group_log.add_argument("--profile", default=os.environ.get("SD_PROFILE", False), action='store_true', help="Run profiler, default: %(default)s")
+    group_log.add_argument("--monitor", default=os.environ.get("SD_PROFILE", 0), help="Run memory monitor, default: %(default)s")
     group_log.add_argument('--docs', default=os.environ.get("SD_DOCS", False), action='store_true', help="Mount API docs, default: %(default)s")
     group_log.add_argument("--api-log", default=os.environ.get("SD_APILOG", True), action='store_true', help="Log all API requests")
 

launch.py

@@ -150,10 +150,14 @@ def run_extension_installer(ext_dir): # compatbility function
     installer.run_extension_installer(ext_dir)
 
 
-def get_memory_stats():
-    from modules.memstats import ram_stats
-    res = ram_stats()
-    return f'{res["used"]}/{res["total"]}'
+def get_memory_stats(detailed:bool=False):
+    from modules.memstats import ram_stats, memory_stats
+    if not detailed:
+        res = ram_stats()
+        return f'{res["used"]}/{res["total"]}'
+    else:
+        res = memory_stats()
+        return res
 
 
 def start_server(immediate=True, server=None):
@@ -260,6 +264,8 @@ def main():
     get_custom_args()
 
     uv, instance = start_server(immediate=True, server=None)
+    t_server = time.time()
+    t_monitor = time.time()
     while True:
         try:
             alive = uv.thread.is_alive()
@@ -267,8 +273,13 @@ def main():
         except Exception:
             alive = False
             requests = 0
-        if round(time.time()) % 120 == 0:
-            installer.log.debug(f'Server: alive={alive} requests={requests} memory={get_memory_stats()} {instance.state.status()}')
+        t_current = time.time()
+        if t_current - t_server > 120:
+            installer.log.trace(f'Server: alive={alive} requests={requests} memory={get_memory_stats()} {instance.state.status()}')
+            t_server = t_current
+        if float(args.monitor) > 0 and t_current - t_monitor > float(args.monitor):
+            installer.log.trace(f'Monitor: {get_memory_stats(detailed=True)}')
+            t_monitor = t_current
         if not alive:
             if uv is not None and uv.wants_restart:
                 installer.log.info('Server restarting...')

modules/cmd_args.py

@@ -37,6 +37,7 @@ def main_args():
     group_diag.add_argument("--no-hashing", default=os.environ.get("SD_NOHASHING", False), action='store_true', help="Disable hashing of checkpoints, default: %(default)s")
     group_diag.add_argument("--no-metadata", default=os.environ.get("SD_NOMETADATA", False), action='store_true', help="Disable reading of metadata from models, default: %(default)s")
     group_diag.add_argument("--profile", default=os.environ.get("SD_PROFILE", False), action='store_true', help="Run profiler, default: %(default)s")
+    group_diag.add_argument("--monitor", default=os.environ.get("SD_PROFILE", 0), help="Run memory monitor, default: %(default)s")
 
     group_http = parser.add_argument_group('HTTP')
     group_http.add_argument('--theme', type=str, default=os.environ.get("SD_THEME", None), help='Override UI theme')

modules/lora/extra_networks_lora.py

@@ -3,7 +3,8 @@ import os
 import re
 import numpy as np
 from modules.lora import networks, lora_overrides, lora_load
-from modules import extra_networks, shared, sd_models
+from modules.lora import lora_common as l
+from modules import extra_networks, shared
 
 
 debug = os.environ.get('SD_LORA_DEBUG', None) is not None
@@ -39,7 +40,7 @@ def prompt(p):
     if shared.opts.lora_apply_tags == 0:
         return
     all_tags = []
-    for loaded in networks.loaded_networks:
+    for loaded in l.loaded_networks:
         page = [en for en in shared.extra_networks if en.name == 'lora'][0]
         item = page.create_item(loaded.name)
         tags = (item or {}).get("tags", {})
@@ -69,12 +70,12 @@ def prompt(p):
 
 
 def infotext(p):
-    names = [i.name for i in networks.loaded_networks]
+    names = [i.name for i in l.loaded_networks]
     if len(names) > 0:
         p.extra_generation_params["LoRA networks"] = ", ".join(names)
     if shared.opts.lora_add_hashes_to_infotext:
         network_hashes = []
-        for item in networks.loaded_networks:
+        for item in l.loaded_networks:
             if not item.network_on_disk.shorthash:
                 continue
             network_hashes.append(item.network_on_disk.shorthash)
@@ -113,6 +114,19 @@ def parse(p, params_list, step=0):
     return names, te_multipliers, unet_multipliers, dyn_dims
 
 
+def unload_diffusers():
+    if hasattr(shared.sd_model, "unfuse_lora"):
+        try:
+            shared.sd_model.unfuse_lora()
+        except Exception:
+            pass
+    if hasattr(shared.sd_model, "unload_lora_weights"):
+        try:
+            shared.sd_model.unload_lora_weights() # fails for non-CLIP models
+        except Exception:
+            pass
+
+
 class ExtraNetworkLora(extra_networks.ExtraNetwork):
 
     def __init__(self):
@@ -131,11 +145,11 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork):
         key = f'{",".join(include)}:{",".join(exclude)}'
         loaded = sd_model.loaded_loras.get(key, [])
         # shared.log.trace(f'Network load: type=LoRA key="{key}" requested={requested} loaded={loaded}')
-        if (len(requested) == 0) or (len(requested) != len(loaded)):
+        if len(requested) != len(loaded):
             sd_model.loaded_loras[key] = requested
             return True
-        for r, l in zip(requested, loaded):
-            if r != l:
+        for req, load in zip(requested, loaded):
+            if req != load:
                 sd_model.loaded_loras[key] = requested
                 return True
         return False
@@ -160,40 +174,35 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork):
         if force_diffusers:
             has_changed = False # diffusers handle their own loading
             if len(exclude) == 0:
+                shared.state.begin('LoRA')
                 lora_load.network_load(names, te_multipliers, unet_multipliers, dyn_dims) # load only on first call
+                shared.state.end()
         else:
             lora_load.network_load(names, te_multipliers, unet_multipliers, dyn_dims) # load
             has_changed = self.changed(requested, include, exclude)
             if has_changed:
-                networks.network_deactivate(include, exclude)
+                shared.state.begin('LoRA')
+                if len(l.previously_loaded_networks) > 0:
+                    shared.log.info(f'Network unload: type=LoRA apply={[n.name for n in l.previously_loaded_networks]} mode={"fuse" if shared.opts.lora_fuse_diffusers else "backup"}')
+                    networks.network_deactivate(include, exclude)
                 networks.network_activate(include, exclude)
-                shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model) # TODO lora: required for flux to reapply offload after lora has been applied, but fails with oom
-                debug_log(f'Network load: type=LoRA previous={[n.name for n in networks.previously_loaded_networks]} current={[n.name for n in networks.loaded_networks]} changed')
+                if len(exclude) > 0: # only update on last activation
+                    l.previously_loaded_networks = l.loaded_networks.copy()
+                shared.state.end()
+                debug_log(f'Network load: type=LoRA previous={[n.name for n in l.previously_loaded_networks]} current={[n.name for n in l.loaded_networks]} changed')
 
-        if len(networks.loaded_networks) > 0 and (len(networks.applied_layers) > 0 or force_diffusers) and step == 0:
+        if len(l.loaded_networks) > 0 and (len(networks.applied_layers) > 0 or force_diffusers) and step == 0:
             infotext(p)
             prompt(p)
             if (has_changed or force_diffusers) and len(include) == 0: # print only once
-                shared.log.info(f'Network load: type=LoRA apply={[n.name for n in networks.loaded_networks]} mode={"fuse" if shared.opts.lora_fuse_diffusers else "backup"} te={te_multipliers} unet={unet_multipliers} time={networks.timer.summary}')
+                shared.log.info(f'Network load: type=LoRA apply={[n.name for n in l.loaded_networks]} mode={"fuse" if shared.opts.lora_fuse_diffusers else "backup"} te={te_multipliers} unet={unet_multipliers} time={l.timer.summary}')
 
     def deactivate(self, p):
-        if shared.native:
-            networks.previously_loaded_networks = networks.loaded_networks.copy()
-            debug_log(f'Network load: type=LoRA active={[n.name for n in networks.previously_loaded_networks]} deactivate')
         if shared.native and len(lora_load.diffuser_loaded) > 0:
             if not (shared.compiled_model_state is not None and shared.compiled_model_state.is_compiled is True):
-                if hasattr(shared.sd_model, "unfuse_lora"):
-                    try:
-                        shared.sd_model.unfuse_lora()
-                    except Exception:
-                        pass
-                if hasattr(shared.sd_model, "unload_lora_weights"):
-                    try:
-                        shared.sd_model.unload_lora_weights() # fails for non-CLIP models
-                    except Exception:
-                        pass
-        if self.active and networks.debug:
-            shared.log.debug(f"Network end: type=LoRA time={networks.timer.summary}")
+                unload_diffusers()
+        if self.active and l.debug:
+            shared.log.debug(f"Network end: type=LoRA time={l.timer.summary}")
         if self.errors:
             for k, v in self.errors.items():
                 shared.log.error(f'LoRA: name="{k}" errors={v}')

modules/lora/lora_apply.py

@@ -3,7 +3,7 @@ import re
 import time
 import torch
 import diffusers.models.lora
-from modules.lora.lora_common import timer, debug, loaded_networks, previously_loaded_networks, extra_network_lora
+from modules.lora import lora_common as l
 from modules import shared, devices, errors, model_quant
 
 
@@ -14,7 +14,7 @@ re_network_name = re.compile(r"(.*)\s*\([0-9a-fA-F]+\)")
 def network_backup_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.GroupNorm, torch.nn.LayerNorm, diffusers.models.lora.LoRACompatibleLinear, diffusers.models.lora.LoRACompatibleConv], network_layer_name: str, wanted_names: tuple):
     global bnb # pylint: disable=W0603
     backup_size = 0
-    if len(loaded_networks) > 0 and network_layer_name is not None and any([net.modules.get(network_layer_name, None) for net in loaded_networks]): # noqa: C419 # pylint: disable=R1729
+    if len(l.loaded_networks) > 0 and network_layer_name is not None and any([net.modules.get(network_layer_name, None) for net in l.loaded_networks]): # noqa: C419 # pylint: disable=R1729
         t0 = time.time()
 
         weights_backup = getattr(self, "network_weights_backup", None)
@@ -33,25 +33,15 @@ def network_backup_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.n
                 if bnb is None:
                     bnb = model_quant.load_bnb('Network load: type=LoRA', silent=True)
                 if bnb is not None:
-                    with devices.inference_context():
-                        if shared.opts.lora_fuse_diffusers:
-                            self.network_weights_backup = True
-                        else:
-                            self.network_weights_backup = bnb.functional.dequantize_4bit(weight, quant_state=weight.quant_state, quant_type=weight.quant_type, blocksize=weight.blocksize,)
-                        self.quant_state = weight.quant_state
-                        self.quant_type = weight.quant_type
-                        self.blocksize = weight.blocksize
-                else:
                     if shared.opts.lora_fuse_diffusers:
                         self.network_weights_backup = True
                     else:
-                        weights_backup = weight.clone()
-                        self.network_weights_backup = weights_backup.to(devices.cpu)
-            else:
-                if shared.opts.lora_fuse_diffusers:
-                    self.network_weights_backup = True
+                        self.network_weights_backup = bnb.functional.dequantize_4bit(weight, quant_state=weight.quant_state, quant_type=weight.quant_type, blocksize=weight.blocksize,)
+                    self.quant_state, self.quant_type, self.blocksize = weight.quant_state, weight.quant_type, weight.blocksize
                 else:
-                    self.network_weights_backup = weight.clone().to(devices.cpu)
+                    self.network_weights_backup = weight.clone().to(devices.cpu) if not shared.opts.lora_fuse_diffusers else True
+            else:
+                self.network_weights_backup = weight.clone().to(devices.cpu) if not shared.opts.lora_fuse_diffusers else True
 
         if bias_backup is None:
             if getattr(self, 'bias', None) is not None:
@@ -65,7 +55,7 @@ def network_backup_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.n
             backup_size += self.network_weights_backup.numel() * self.network_weights_backup.element_size() if isinstance(self.network_weights_backup, torch.Tensor) else 0
         if getattr(self, 'network_bias_backup', None) is not None:
             backup_size += self.network_bias_backup.numel() * self.network_bias_backup.element_size() if isinstance(self.network_bias_backup, torch.Tensor) else 0
-        timer.backup += time.time() - t0
+        l.timer.backup += time.time() - t0
     return backup_size
 
 
@@ -77,7 +67,7 @@ def network_calc_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.
             pass
     batch_updown = None
     batch_ex_bias = None
-    loaded = loaded_networks if not use_previous else previously_loaded_networks
+    loaded = l.loaded_networks if not use_previous else l.previously_loaded_networks
     for net in loaded:
         module = net.modules.get(network_layer_name, None)
         if module is None:
@@ -88,8 +78,8 @@ def network_calc_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.
                 weight = self.weight.to(devices.device)
             except Exception:
                 weight = self.weight
-
             updown, ex_bias = module.calc_updown(weight)
+            del module
             if updown is not None:
                 if batch_updown is not None:
                     batch_updown += updown.to(batch_updown.device)
@@ -100,8 +90,7 @@ def network_calc_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.
                     batch_ex_bias += ex_bias.to(batch_ex_bias.device)
                 else:
                     batch_ex_bias = ex_bias.to(devices.device)
-            timer.calc += time.time() - t0
-
+            l.timer.calc += time.time() - t0
             if shared.opts.diffusers_offload_mode == "sequential":
                 t0 = time.time()
                 if batch_updown is not None:
@@ -109,10 +98,10 @@ def network_calc_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.
                 if batch_ex_bias is not None:
                     batch_ex_bias = batch_ex_bias.to(devices.cpu)
                 t1 = time.time()
-                timer.move += t1 - t0
+                l.timer.move += t1 - t0
         except RuntimeError as e:
-            extra_network_lora.errors[net.name] = extra_network_lora.errors.get(net.name, 0) + 1
-            if debug:
+            l.extra_network_lora.errors[net.name] = l.extra_network_lora.errors.get(net.name, 0) + 1
+            if l.debug:
                 module_name = net.modules.get(network_layer_name, None)
                 shared.log.error(f'LoRA apply weight name="{net.name}" module="{module_name}" layer="{network_layer_name}" {e}')
                 errors.display(e, 'LoRA')
@@ -121,7 +110,7 @@ def network_calc_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.
     return batch_updown, batch_ex_bias
 
 
-def network_add_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.GroupNorm, torch.nn.LayerNorm, diffusers.models.lora.LoRACompatibleLinear, diffusers.models.lora.LoRACompatibleConv], model_weights: Union[None, torch.Tensor] = None, lora_weights: torch.Tensor = None, deactivate: bool = False, device: torch.device = devices.device):
+def network_add_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.GroupNorm, torch.nn.LayerNorm, diffusers.models.lora.LoRACompatibleLinear, diffusers.models.lora.LoRACompatibleConv], model_weights: Union[None, torch.Tensor] = None, lora_weights: torch.Tensor = None, deactivate: bool = False, bias: bool = False):
     if lora_weights is None:
         return None
     if deactivate:
@@ -135,20 +124,25 @@ def network_add_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.G
             dequant_weight = bnb.functional.dequantize_4bit(model_weights.to(devices.device), quant_state=self.quant_state, quant_type=self.quant_type, blocksize=self.blocksize)
             new_weight = dequant_weight.to(devices.device) + lora_weights.to(devices.device)
             weight = bnb.nn.Params4bit(new_weight, quant_state=self.quant_state, quant_type=self.quant_type, blocksize=self.blocksize, requires_grad=False)
-            # weight._quantize(devices.device) # TODO force imediate quantization
+            # TODO lora: maybe force imediate quantization
+            # weight._quantize(devices.device) / weight.to(device=device)
         except Exception as e:
             shared.log.error(f'Network load: type=LoRA quant=bnb cls={self.__class__.__name__} type={self.quant_type} blocksize={self.blocksize} state={vars(self.quant_state)} weight={self.weight} bias={lora_weights} {e}')
     else:
         try:
             new_weight = model_weights.to(devices.device) + lora_weights.to(devices.device)
-        except Exception:
+        except Exception as e:
+            shared.log.warning(f'Network load: {e}')
             new_weight = model_weights + lora_weights # try without device cast
+        del model_weights
+        del lora_weights
         weight = torch.nn.Parameter(new_weight, requires_grad=False)
-    try:
-        # weight.to(device=device) # TODO required since quantization happens only during .to call, not during params creation
-        pass
-    except Exception:
-        pass # may fail if weights is meta tensor
+        del new_weight # without this its a massive memory leak
+    if weight is not None:
+        if not bias:
+            self.weight = weight
+        else:
+            self.bias = weight
     return weight
 
 
@@ -166,22 +160,18 @@ def network_apply_direct(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.
 
     if weights_backup:
         if updown is not None and len(self.weight.shape) == 4 and self.weight.shape[1] == 9: # inpainting model so zero pad updown to make channel 4 to 9
-            updown = torch.nn.functional.pad(updown, (0, 0, 0, 0, 0, 5))  # pylint: disable=not-callable
+            updown = torch.nn.functional.pad(updown, (0, 0, 0, 0, 0, 5)) # pylint: disable=not-callable
         if updown is not None:
-            weight = network_add_weights(self, lora_weights=updown, deactivate=deactivate, device=device)
-            if weight is not None:
-                self.weight = weight
+            network_add_weights(self, lora_weights=updown, deactivate=deactivate, bias=False)
 
     if bias_backup:
         if ex_bias is not None:
-            bias = network_add_weights(self, lora_weights=ex_bias, deactivate=deactivate, device=device)
-            if bias is not None:
-                self.bias = bias
+            network_add_weights(self, lora_weights=ex_bias, deactivate=deactivate, bias=True)
 
     if hasattr(self, "qweight") and hasattr(self, "freeze"):
         self.freeze()
 
-    timer.apply += time.time() - t0
+    l.timer.apply += time.time() - t0
 
 
 def network_apply_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.GroupNorm, torch.nn.LayerNorm, diffusers.models.lora.LoRACompatibleLinear, diffusers.models.lora.LoRACompatibleConv], updown: torch.Tensor, ex_bias: torch.Tensor, device: torch.device, deactivate: bool = False):
@@ -194,24 +184,20 @@ def network_apply_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn
     if weights_backup is not None:
         self.weight = None
         if updown is not None and len(weights_backup.shape) == 4 and weights_backup.shape[1] == 9: # inpainting model. zero pad updown to make channel[1]  4 to 9
-            updown = torch.nn.functional.pad(updown, (0, 0, 0, 0, 0, 5))  # pylint: disable=not-callable
+            updown = torch.nn.functional.pad(updown, (0, 0, 0, 0, 0, 5)) # pylint: disable=not-callable
         if updown is not None:
-            weight = network_add_weights(self, model_weights=weights_backup, lora_weights=updown, deactivate=deactivate, device=device)
-            if weight is not None:
-                self.weight = weight
+            network_add_weights(self, model_weights=weights_backup, lora_weights=updown, deactivate=deactivate, bias=False)
         else:
             self.weight = torch.nn.Parameter(weights_backup.to(device), requires_grad=False)
 
     if bias_backup is not None:
         self.bias = None
         if ex_bias is not None:
-            bias = network_add_weights(self, model_weights=weights_backup, lora_weights=ex_bias, deactivate=deactivate, device=device)
-            if bias:
-                self.weight = bias
+            network_add_weights(self, model_weights=bias_backup, lora_weights=ex_bias, deactivate=deactivate, bias=True)
         else:
             self.bias = torch.nn.Parameter(bias_backup.to(device), requires_grad=False)
 
     if hasattr(self, "qweight") and hasattr(self, "freeze"):
         self.freeze()
 
-    timer.apply += time.time() - t0
+    l.timer.apply += time.time() - t0

modules/lora/lora_load.py

@@ -4,7 +4,7 @@ import time
 import concurrent
 from modules import shared, errors, devices, sd_models, sd_models_compile, files_cache
 from modules.lora import network, lora_overrides, lora_convert
-from modules.lora.lora_common import timer, debug, module_types, loaded_networks
+from modules.lora import lora_common as l
 
 
 diffuser_loaded = []
@@ -35,7 +35,7 @@ def load_diffusers(name, network_on_disk, lora_scale=shared.opts.extra_networks_
                 shared.log.error(f'Network load: type=LoRA name="{name}" diffusers unsupported format')
             else:
                 shared.log.error(f'Network load: type=LoRA name="{name}" {e}')
-            if debug:
+            if l.debug:
                 errors.display(e, "LoRA")
             return None
     if name not in diffuser_loaded:
@@ -43,7 +43,7 @@ def load_diffusers(name, network_on_disk, lora_scale=shared.opts.extra_networks_
         diffuser_scales.append(lora_scale)
     net = network.Network(name, network_on_disk)
     net.mtime = os.path.getmtime(network_on_disk.filename)
-    timer.activate += time.time() - t0
+    l.timer.activate += time.time() - t0
     return net
 
 
@@ -52,19 +52,19 @@ def load_safetensors(name, network_on_disk) -> Union[network.Network, None]:
         return None
 
     cached = lora_cache.get(name, None)
-    if debug:
+    if l.debug:
         shared.log.debug(f'Network load: type=LoRA name="{name}" file="{network_on_disk.filename}" type=lora {"cached" if cached else ""}')
     if cached is not None:
         return cached
     net = network.Network(name, network_on_disk)
     net.mtime = os.path.getmtime(network_on_disk.filename)
     sd = sd_models.read_state_dict(network_on_disk.filename, what='network')
-    if shared.sd_model_type == 'f1':  # if kohya flux lora, convert state_dict
-        sd = lora_convert._convert_kohya_flux_lora_to_diffusers(sd) or sd  # pylint: disable=protected-access
-    if shared.sd_model_type == 'sd3':  # if kohya flux lora, convert state_dict
+    if shared.sd_model_type == 'f1': # if kohya flux lora, convert state_dict
+        sd = lora_convert._convert_kohya_flux_lora_to_diffusers(sd) or sd # pylint: disable=protected-access
+    if shared.sd_model_type == 'sd3': # if kohya flux lora, convert state_dict
         try:
-            sd = lora_convert._convert_kohya_sd3_lora_to_diffusers(sd) or sd  # pylint: disable=protected-access
-        except ValueError:  # EAFP for diffusers PEFT keys
+            sd = lora_convert._convert_kohya_sd3_lora_to_diffusers(sd) or sd # pylint: disable=protected-access
+        except ValueError: # EAFP for diffusers PEFT keys
             pass
     lora_convert.assign_network_names_to_compvis_modules(shared.sd_model)
     keys_failed_to_match = {}
@@ -72,6 +72,7 @@ def load_safetensors(name, network_on_disk) -> Union[network.Network, None]:
     bundle_embeddings = {}
     dtypes = []
     convert = lora_convert.KeyConvert()
+    device = devices.device if shared.opts.lora_apply_gpu else devices.cpu
     for key_network, weight in sd.items():
         parts = key_network.split('.')
         if parts[0] == "bundle_emb":
@@ -99,7 +100,7 @@ def load_safetensors(name, network_on_disk) -> Union[network.Network, None]:
     network_types = []
     for key, weights in matched_networks.items():
         net_module = None
-        for nettype in module_types:
+        for nettype in l.module_types:
             net_module = nettype.create_module(net, weights)
             if net_module is not None:
                 network_types.append(nettype.__class__.__name__)
@@ -110,10 +111,10 @@ def load_safetensors(name, network_on_disk) -> Union[network.Network, None]:
             net.modules[key] = net_module
     if len(keys_failed_to_match) > 0:
         shared.log.warning(f'Network load: type=LoRA name="{name}" type={set(network_types)} unmatched={len(keys_failed_to_match)} matched={len(matched_networks)}')
-        if debug:
+        if l.debug:
             shared.log.debug(f'Network load: type=LoRA name="{name}" unmatched={keys_failed_to_match}')
     else:
-        shared.log.debug(f'Network load: type=LoRA name="{name}" type={set(network_types)} keys={len(matched_networks)} dtypes={dtypes} direct={shared.opts.lora_fuse_diffusers}')
+        shared.log.debug(f'Network load: type=LoRA name="{name}" type={set(network_types)} keys={len(matched_networks)} device={device} dtypes={dtypes} direct={shared.opts.lora_fuse_diffusers}')
     if len(matched_networks) == 0:
         return None
     lora_cache[name] = net
@@ -134,7 +135,7 @@ def maybe_recompile_model(names, te_multipliers):
                     break
             if not recompile_model:
                 skip_lora_load = True
-                if len(loaded_networks) > 0 and debug:
+                if len(l.loaded_networks) > 0 and l.debug:
                     shared.log.debug('Model Compile: Skipping LoRa loading')
                 return recompile_model, skip_lora_load
         else:
@@ -178,7 +179,7 @@ def list_available_networks():
                 available_network_aliases[entry.alias] = entry
             if entry.shorthash:
                 available_network_hash_lookup[entry.shorthash] = entry
-        except OSError as e:  # should catch FileNotFoundError and PermissionError etc.
+        except OSError as e: # should catch FileNotFoundError and PermissionError etc.
             shared.log.error(f'LoRA: filename="{filename}" {e}')
 
     candidates = sorted(files_cache.list_files(shared.cmd_opts.lora_dir, ext_filter=[".pt", ".ckpt", ".safetensors"]))
@@ -186,7 +187,7 @@ def list_available_networks():
         for fn in candidates:
             executor.submit(add_network, fn)
     t1 = time.time()
-    timer.list = t1 - t0
+    l.timer.list = t1 - t0
     shared.log.info(f'Available LoRAs: path="{shared.cmd_opts.lora_dir}" items={len(available_networks)} folders={len(forbidden_network_aliases)} time={t1 - t0:.2f}')
 
 
@@ -214,7 +215,7 @@ def network_load(names, te_multipliers=None, unet_multipliers=None, dyn_dims=Non
     failed_to_load_networks = []
     recompile_model, skip_lora_load = maybe_recompile_model(names, te_multipliers)
 
-    loaded_networks.clear()
+    l.loaded_networks.clear()
     diffuser_loaded.clear()
     diffuser_scales.clear()
     t0 = time.time()
@@ -223,7 +224,7 @@ def network_load(names, te_multipliers=None, unet_multipliers=None, dyn_dims=Non
         net = None
         if network_on_disk is not None:
             shorthash = getattr(network_on_disk, 'shorthash', '').lower()
-            if debug:
+            if l.debug:
                 shared.log.debug(f'Network load: type=LoRA name="{name}" file="{network_on_disk.filename}" hash="{shorthash}"')
             try:
                 if recompile_model:
@@ -237,7 +238,7 @@ def network_load(names, te_multipliers=None, unet_multipliers=None, dyn_dims=Non
                     network_on_disk.read_hash()
             except Exception as e:
                 shared.log.error(f'Network load: type=LoRA file="{network_on_disk.filename}" {e}')
-                if debug:
+                if l.debug:
                     errors.display(e, 'LoRA')
                 continue
         if net is None:
@@ -249,7 +250,7 @@ def network_load(names, te_multipliers=None, unet_multipliers=None, dyn_dims=Non
         net.te_multiplier = te_multipliers[i] if te_multipliers else shared.opts.extra_networks_default_multiplier
         net.unet_multiplier = unet_multipliers[i] if unet_multipliers else shared.opts.extra_networks_default_multiplier
         net.dyn_dim = dyn_dims[i] if dyn_dims else shared.opts.extra_networks_default_multiplier
-        loaded_networks.append(net)
+        l.loaded_networks.append(net)
 
     while len(lora_cache) > shared.opts.lora_in_memory_limit:
         name = next(iter(lora_cache))
@@ -261,16 +262,16 @@ def network_load(names, te_multipliers=None, unet_multipliers=None, dyn_dims=Non
             t0 = time.time()
             shared.sd_model.set_adapters(adapter_names=diffuser_loaded, adapter_weights=diffuser_scales)
             if shared.opts.lora_fuse_diffusers and not lora_overrides.check_fuse():
-                shared.sd_model.fuse_lora(adapter_names=diffuser_loaded, lora_scale=1.0, fuse_unet=True, fuse_text_encoder=True) # fuse uses fixed scale since later apply does the scaling
+                shared.sd_model.fuse_lora(adapter_names=diffuser_loaded, lora_scale=1.0, fuse_unet=True, fuse_text_encoder=True) # diffusers with fuse uses fixed scale since later apply does the scaling
                 shared.sd_model.unload_lora_weights()
-            timer.activate += time.time() - t0
+            l.timer.activate += time.time() - t0
         except Exception as e:
             shared.log.error(f'Network load: type=LoRA {e}')
-            if debug:
+            if l.debug:
                 errors.display(e, 'LoRA')
 
-    if len(loaded_networks) > 0 and debug:
-        shared.log.debug(f'Network load: type=LoRA loaded={[n.name for n in loaded_networks]} cache={list(lora_cache)}')
+    if len(l.loaded_networks) > 0 and l.debug:
+        shared.log.debug(f'Network load: type=LoRA loaded={[n.name for n in l.loaded_networks]} cache={list(lora_cache)}')
 
     if recompile_model:
         shared.log.info("Network load: type=LoRA recompiling model")
@@ -279,7 +280,4 @@ def network_load(names, te_multipliers=None, unet_multipliers=None, dyn_dims=Non
             shared.sd_model = sd_models_compile.compile_diffusers(shared.sd_model)
         shared.compiled_model_state.lora_model = backup_lora_model
 
-    if len(loaded_networks) > 0:
-        devices.torch_gc()
-
-    timer.load = time.time() - t0
+    l.timer.load = time.time() - t0

modules/lora/networks.py

@@ -1,7 +1,7 @@
 from contextlib import nullcontext
 import time
 import rich.progress as rp
-from modules.lora.lora_common import timer, debug, loaded_networks, previously_loaded_networks
+from modules.lora import lora_common as l
 from modules.lora.lora_apply import network_apply_weights, network_apply_direct, network_backup_weights, network_calc_weights
 from modules import shared, devices, sd_models
 
@@ -11,7 +11,7 @@ applied_layers: list[str] = []
 
 def network_activate(include=[], exclude=[]):
     t0 = time.time()
-    sd_model = getattr(shared.sd_model, "pipe", shared.sd_model)  # wrapped model compatiblility
+    sd_model = getattr(shared.sd_model, "pipe", shared.sd_model)
     if shared.opts.diffusers_offload_mode == "sequential":
         sd_models.disable_offload(sd_model)
         sd_models.move_model(sd_model, device=devices.cpu)
@@ -25,7 +25,7 @@ def network_activate(include=[], exclude=[]):
             active_components.append(name)
             modules[name] = list(component.named_modules())
     total = sum(len(x) for x in modules.values())
-    if len(loaded_networks) > 0:
+    if len(l.loaded_networks) > 0:
         pbar = rp.Progress(rp.TextColumn('[cyan]Network: type=LoRA action=activate'), rp.BarColumn(), rp.TaskProgressColumn(), rp.TimeRemainingColumn(), rp.TimeElapsedColumn(), rp.TextColumn('[cyan]{task.description}'), console=shared.console)
         task = pbar.add_task(description='' , total=total)
     else:
@@ -33,9 +33,9 @@ def network_activate(include=[], exclude=[]):
         pbar = nullcontext()
     applied_weight = 0
     applied_bias = 0
-    device = devices.device if shared.opts.lora_apply_gpu else devices.cpu
+    device = devices.device if shared.opts.lora_apply_gpu or shared.opts.diffusers_offload_mode == 'none' else devices.cpu
     with devices.inference_context(), pbar:
-        wanted_names = tuple((x.name, x.te_multiplier, x.unet_multiplier, x.dyn_dim) for x in loaded_networks) if len(loaded_networks) > 0 else ()
+        wanted_names = tuple((x.name, x.te_multiplier, x.unet_multiplier, x.dyn_dim) for x in l.loaded_networks) if len(l.loaded_networks) > 0 else ()
         applied_layers.clear()
         backup_size = 0
         for component in modules.keys():
@@ -55,32 +55,32 @@ def network_activate(include=[], exclude=[]):
                     network_apply_weights(module, batch_updown, batch_ex_bias, device=orig_device)
                 if batch_updown is not None or batch_ex_bias is not None:
                     applied_layers.append(network_layer_name)
-                    # module.to(device) # TODO maybe
-                    if batch_updown is not None:
-                        applied_weight += 1
-                    if batch_ex_bias is not None:
-                        applied_bias += 1
+                    applied_weight += 1 if batch_updown is not None else 0
+                    applied_bias += 1 if batch_ex_bias is not None else 0
+                batch_updown, batch_ex_bias = None, None
                 del batch_updown, batch_ex_bias
                 module.network_current_names = wanted_names
                 if task is not None:
-                    pbar.update(task, advance=1, description=f'networks={len(loaded_networks)} modules={active_components} layers={total} weights={applied_weight} bias={applied_bias} backup={backup_size}')
+                    bs = round(backup_size/1024/1024/1024, 2) if backup_size > 0 else None
+                    pbar.update(task, advance=1, description=f'networks={len(l.loaded_networks)} modules={active_components} layers={total} weights={applied_weight} bias={applied_bias} backup={bs} device={device}')
 
         if task is not None and len(applied_layers) == 0:
             pbar.remove_task(task) # hide progress bar for no action
-    timer.activate += time.time() - t0
-    if debug and len(loaded_networks) > 0:
-        shared.log.debug(f'Network load: type=LoRA networks={[n.name for n in loaded_networks]} modules={active_components} layers={total} weights={applied_weight} bias={applied_bias} backup={backup_size} fuse={shared.opts.lora_fuse_diffusers} device={device} time={timer.summary}')
+    l.timer.activate += time.time() - t0
+    if l.debug and len(l.loaded_networks) > 0:
+        shared.log.debug(f'Network load: type=LoRA networks={[n.name for n in l.loaded_networks]} modules={active_components} layers={total} weights={applied_weight} bias={applied_bias} backup={round(backup_size/1024/1024/1024, 2)} fuse={shared.opts.lora_fuse_diffusers} device={device} time={l.timer.summary}')
     modules.clear()
-    if len(loaded_networks) > 0 and (applied_weight > 0 or applied_bias > 0):
-        if shared.opts.diffusers_offload_mode == "sequential":
-            sd_models.set_diffuser_offload(sd_model, op="model")
+    if len(applied_layers) > 0 or shared.opts.diffusers_offload_mode == "sequential":
+        sd_models.set_diffuser_offload(sd_model, op="model")
 
 
 def network_deactivate(include=[], exclude=[]):
     if not shared.opts.lora_fuse_diffusers or shared.opts.lora_force_diffusers:
         return
+    if len(l.previously_loaded_networks) == 0:
+        return
     t0 = time.time()
-    sd_model = getattr(shared.sd_model, "pipe", shared.sd_model)  # wrapped model compatiblility
+    sd_model = getattr(shared.sd_model, "pipe", shared.sd_model)
     if shared.opts.diffusers_offload_mode == "sequential":
         sd_models.disable_offload(sd_model)
         sd_models.move_model(sd_model, device=devices.cpu)
@@ -96,7 +96,7 @@ def network_deactivate(include=[], exclude=[]):
             active_components.append(name)
     total = sum(len(x) for x in modules.values())
     device = devices.device if shared.opts.lora_apply_gpu else devices.cpu
-    if len(previously_loaded_networks) > 0 and debug:
+    if len(l.previously_loaded_networks) > 0 and l.debug:
         pbar = rp.Progress(rp.TextColumn('[cyan]Network: type=LoRA action=deactivate'), rp.BarColumn(), rp.TaskProgressColumn(), rp.TimeRemainingColumn(), rp.TimeElapsedColumn(), rp.TextColumn('[cyan]{task.description}'), console=shared.console)
         task = pbar.add_task(description='', total=total)
     else:
@@ -118,16 +118,15 @@ def network_deactivate(include=[], exclude=[]):
                 else:
                     network_apply_weights(module, batch_updown, batch_ex_bias, device=orig_device, deactivate=True)
                 if batch_updown is not None or batch_ex_bias is not None:
-                    # module.to(device) # TODO maybe
                     applied_layers.append(network_layer_name)
                 del batch_updown, batch_ex_bias
                 module.network_current_names = ()
                 if task is not None:
-                    pbar.update(task, advance=1, description=f'networks={len(previously_loaded_networks)} modules={active_components} layers={total} unapply={len(applied_layers)}')
+                    pbar.update(task, advance=1, description=f'networks={len(l.previously_loaded_networks)} modules={active_components} layers={total} unapply={len(applied_layers)}')
 
-    timer.deactivate = time.time() - t0
-    if debug and len(previously_loaded_networks) > 0:
-        shared.log.debug(f'Network deactivate: type=LoRA networks={[n.name for n in previously_loaded_networks]} modules={active_components} layers={total} apply={len(applied_layers)} fuse={shared.opts.lora_fuse_diffusers} time={timer.summary}')
+    l.timer.deactivate = time.time() - t0
+    if l.debug and len(l.previously_loaded_networks) > 0:
+        shared.log.debug(f'Network deactivate: type=LoRA networks={[n.name for n in l.previously_loaded_networks]} modules={active_components} layers={total} apply={len(applied_layers)} fuse={shared.opts.lora_fuse_diffusers} time={l.timer.summary}')
     modules.clear()
-    if shared.opts.diffusers_offload_mode == "sequential":
+    if len(applied_layers) > 0 or shared.opts.diffusers_offload_mode == "sequential":
         sd_models.set_diffuser_offload(sd_model, op="model")

modules/memstats.py

@@ -56,15 +56,17 @@ def memory_stats():
             fail_once = True
         mem.update({ 'ram': { 'error': str(e) } })
     try:
-        s = torch.cuda.mem_get_info()
-        gpu = { 'used': gb(s[1] - s[0]), 'total': gb(s[1]) }
-        s = dict(torch.cuda.memory_stats())
-        if s.get('num_ooms', 0) > 0:
+        free, total = torch.cuda.mem_get_info()
+        gpu = { 'used': gb(total - free), 'total': gb(total) }
+        stats = dict(torch.cuda.memory_stats())
+        if stats.get('num_ooms', 0) > 0:
             shared.state.oom = True
         mem.update({
             'gpu': gpu,
-            'retries': s.get('num_alloc_retries', 0),
-            'oom': s.get('num_ooms', 0)
+            'active': gb(stats.get('active_bytes.all.current', 0)),
+            'peak': gb(stats.get('active_bytes.all.peak', 0)),
+            'retries': stats.get('num_alloc_retries', 0),
+            'oom': stats.get('num_ooms', 0),
         })
         return mem
     except Exception:

modules/model_quant.py

@@ -3,7 +3,7 @@ import sys
 import copy
 import time
 import diffusers
-from installer import install, log, setup_logging
+from installer import installed, install, log, setup_logging
 
 
 ao = None
@@ -116,7 +116,9 @@ def load_torchao(msg='', silent=False):
     global ao # pylint: disable=global-statement
     if ao is not None:
         return ao
-    install('torchao==0.8.0', quiet=True)
+    if not installed('torchao'):
+        install('torchao==0.8.0', quiet=True)
+        log.warning('Quantization: torchao installed please restart')
     try:
         import torchao
         ao = torchao
@@ -140,9 +142,11 @@ def load_bnb(msg='', silent=False):
     global bnb # pylint: disable=global-statement
     if bnb is not None:
         return bnb
-    if devices.backend == 'cuda':
-        # forcing a version will uninstall the multi-backend-refactor branch of bnb
-        install('bitsandbytes==0.45.1', quiet=True)
+    if not installed('bitsandbytes'):
+        if devices.backend == 'cuda':
+            # forcing a version will uninstall the multi-backend-refactor branch of bnb
+            install('bitsandbytes==0.45.1', quiet=True)
+            log.warning('Quantization: bitsandbytes installed please restart')
     try:
         import bitsandbytes
         bnb = bitsandbytes
@@ -165,7 +169,9 @@ def load_quanto(msg='', silent=False):
     global optimum_quanto # pylint: disable=global-statement
     if optimum_quanto is not None:
         return optimum_quanto
-    install('optimum-quanto==0.2.7', quiet=True)
+    if not installed('optimum-quanto'):
+        install('optimum-quanto==0.2.7', quiet=True)
+        log.warning('Quantization: optimum-quanto installed please restart')
     try:
         from optimum import quanto # pylint: disable=no-name-in-module
         optimum_quanto = quanto
@@ -190,7 +196,9 @@ def load_nncf(msg='', silent=False):
     global intel_nncf # pylint: disable=global-statement
     if intel_nncf is not None:
         return intel_nncf
-    install('nncf==2.7.0', quiet=True)
+    if not installed('nncf'):
+        install('nncf==2.7.0', quiet=True)
+        log.warning('Quantization: nncf installed please restart')
     try:
         import nncf
         intel_nncf = nncf

modules/processing_diffusers.py

@@ -9,7 +9,7 @@ from modules import shared, devices, processing, sd_models, errors, sd_hijack_hy
 from modules.processing_helpers import resize_hires, calculate_base_steps, calculate_hires_steps, calculate_refiner_steps, save_intermediate, update_sampler, is_txt2img, is_refiner_enabled, get_job_name
 from modules.processing_args import set_pipeline_args
 from modules.onnx_impl import preprocess_pipeline as preprocess_onnx_pipeline, check_parameters_changed as olive_check_parameters_changed
-from modules.lora import networks
+from modules.lora import lora_common
 
 
 debug = shared.log.trace if os.environ.get('SD_DIFFUSERS_DEBUG', None) is not None else lambda *args, **kwargs: None
@@ -478,8 +478,8 @@ def process_diffusers(p: processing.StableDiffusionProcessing):
             return results
 
     extra_networks.deactivate(p)
-    timer.process.add('lora', networks.timer.total)
-    networks.timer.clear(complete=True)
+    timer.process.add('lora', lora_common.timer.total)
+    lora_common.timer.clear(complete=True)
 
     results = process_decode(p, output)
     timer.process.record('decode')

modules/sd_models.py

@@ -1051,7 +1051,6 @@ def reload_model_weights(sd_model=None, info=None, reuse_dict=False, op='model',
 
 
 def clear_caches():
-    # shared.log.debug('Cache clear')
     if not shared.opts.lora_legacy:
         from modules.lora import lora_common, lora_load
         lora_common.loaded_networks.clear()

modules/sd_offload.py

@@ -166,9 +166,7 @@ class OffloadHook(accelerate.hooks.ModelHook):
                 keys = device_map.keys()
                 for v in keys:
                     if isinstance(device_map[v], int):
-                        # int implies CUDA or XPU device, but it will break DirectML backend.
-                        # Therefore, the type of device should be added.
-                        device_map[v] = f"{devices.device.type}:{device_map[v]}"
+                        device_map[v] = f"{devices.device.type}:{device_map[v]}" # int implies CUDA or XPU device, but it will break DirectML backend so we add type
             module = accelerate.dispatch_model(module, device_map=device_map, offload_dir=offload_dir)
             module._hf_hook.execution_device = torch.device(devices.device) # pylint: disable=protected-access
             module.balanced_offload_device_map = device_map