lora load cache state_dict

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

modules/lora/native_adapter.py

@@ -33,7 +33,7 @@ from dataclasses import dataclass
 
 import torch
 
-from modules import shared, sd_models
+from modules import shared, sd_models, sd_models_utils
 from modules.logger import log
 from modules.lora import (
     lora_convert, network, network_boft, network_full, network_glora,
@@ -483,7 +483,7 @@ def try_load_lokr(name, network_on_disk, lora_scale, *,
                   resolve_targets, prefixes=KNOWN_PREFIXES_DEFAULT,
                   bare_prefixes=(), bare_diffusers_prefixes=(),
                   network_prefix=NETWORK_PREFIX_DEFAULT,
-                  arch_name="generic"):
+                  arch_name="generic"): # pylint: disable=unused-argument
     """Generic LoKR loader.
 
     Stores only the compact LoKR factors and dispatches to
@@ -759,7 +759,7 @@ def try_load_norm(name, network_on_disk, lora_scale, *,
                   resolve_targets, prefixes=KNOWN_PREFIXES_DEFAULT,
                   bare_prefixes=(), bare_diffusers_prefixes=(),
                   network_prefix=NETWORK_PREFIX_DEFAULT,
-                  arch_name="generic"):
+                  arch_name="generic"): # pylint: disable=unused-argument
     """Generic Norm (LayerNorm / RMSNorm weight + bias delta) loader.
 
     Norm targets are never fused, so the chunk dispatch is dropped.
@@ -873,6 +873,7 @@ def try_load_chain(name, network_on_disk, lora_scale, family_loaders):
     tuple of partial-applied generic loaders, each already bound to the arch's
     ``resolve_targets`` and prefix tuples.
     """
+    sd_models_utils.state_dict_cache.enable()
     net = None
     for try_fn in family_loaders:
         sub = try_fn(name, network_on_disk, lora_scale)
@@ -882,4 +883,5 @@ def try_load_chain(name, network_on_disk, lora_scale, family_loaders):
             net = sub
         else:
             net.modules.update(sub.modules)
+    sd_models_utils.state_dict_cache.disable()
     return net

modules/sd_detect.py

@@ -92,7 +92,7 @@ def guess_by_name(fn, current_guess):
         new_guess = 'ZetaChroma'
     elif 'chroma' in fn.lower() and 'xl' not in fn.lower():
         new_guess = 'Chroma'
-    elif 'flux.2' in fn.lower() and 'klein' in fn.lower():
+    elif ('flux.2' in fn.lower() or 'f2' in fn.lower()) and 'klein' in fn.lower():
         new_guess = 'FLUX2Klein'
     elif 'flux.2' in fn.lower():
         new_guess = 'FLUX2'

modules/sd_models_utils.py

@@ -15,6 +15,30 @@ from modules.sd_checkpoint import CheckpointInfo # pylint: disable=unused-import
 debug = log.trace if os.environ.get('SD_LOAD_DEBUG', None) is not None else lambda *args, **kwargs: None
 
 
+class StateDictCache:
+    _enabled: bool = True
+    _cache: dict[str, dict] = {}
+
+    def get(self, key: str):
+        if not self._enabled:
+            return None
+        return self._cache.get(key, None)
+
+    def set(self, key: str, value: dict):
+        if not self._enabled:
+            return
+        self._cache[key] = value
+
+    def enable(self):
+        self._enabled = True
+
+    def disable(self):
+        self._enabled = False
+        self._cache.clear()
+
+state_dict_cache = StateDictCache()
+
+
 class NoWatermark:
     def apply_watermark(self, img):
         return img
@@ -97,12 +121,18 @@ def convert_to_faketensors(tensor):
 
 
 def read_state_dict(checkpoint_file, map_location=None, what:str='model'): # pylint: disable=unused-argument
+    cached = state_dict_cache.get(checkpoint_file)
+    if cached is not None:
+        return cached
     if not os.path.isfile(checkpoint_file):
-        log.error(f'Load dict: path="{checkpoint_file}" not a file')
+        log.error(f'Load dict: file="{checkpoint_file}" not a file')
         return None
     _, extension = os.path.splitext(checkpoint_file)
     if extension.lower() == ".ckpt" and shared.opts.sd_disable_ckpt:
-        log.warning(f"Checkpoint loading disabled: {checkpoint_file}")
+        log.warning(f'Load dict: file="{checkpoint_file}" checkpoint loading disabled')
+        return None
+    if shared.state.interrupted:
+        log.warning(f'Load dict: file="{checkpoint_file}" interrupted before read')
         return None
     try:
         pl_sd = None
@@ -123,6 +153,7 @@ def read_state_dict(checkpoint_file, map_location=None, what:str='model'): # pyl
                 else:
                     pl_sd = torch.load(f, map_location='cpu')
         sd = get_state_dict_from_checkpoint(pl_sd)
+        state_dict_cache.set(checkpoint_file, sd)
         del pl_sd
     except Exception as e:
         errors.display(e, f'Load model: {checkpoint_file}')