lora debug logging and support minimax-ref2va

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

modules/lora/lora_convert.py

@@ -487,28 +487,35 @@ def assign_network_names_to_compvis_modules(sd_model):
     sd_model = getattr(shared.sd_model, "pipe", shared.sd_model)  # wrapped model compatibility
     network_layer_mapping = {}
     if hasattr(sd_model, 'text_encoder') and sd_model.text_encoder is not None:
-        for name, module in sd_model.text_encoder.named_modules():
+        for name, module in sd_model.text_encoder.named_modules() :
             prefix = "lora_te1_" if hasattr(sd_model, 'text_encoder_2') else "lora_te_"
             network_name = prefix + name.replace(".", "_")
             network_layer_mapping[network_name] = module
             module.network_layer_name = network_name
-    if hasattr(sd_model, 'text_encoder_2'):
+    if hasattr(sd_model, 'text_encoder_2') and sd_model.text_encoder_2 is not None:
         for name, module in sd_model.text_encoder_2.named_modules():
             network_name = "lora_te2_" + name.replace(".", "_")
             network_layer_mapping[network_name] = module
             module.network_layer_name = network_name
-    if hasattr(sd_model, 'unet'):
+    if hasattr(sd_model, 'unet') and sd_model.unet is not None:
         for name, module in sd_model.unet.named_modules():
             network_name = "lora_unet_" + name.replace(".", "_")
             network_layer_mapping[network_name] = module
             module.network_layer_name = network_name
-    if hasattr(sd_model, 'transformer'):
+    if hasattr(sd_model, 'transformer') and sd_model.transformer is not None:
         for name, module in sd_model.transformer.named_modules():
             network_name = "lora_transformer_" + name.replace(".", "_")
             network_layer_mapping[network_name] = module
             if "norm" in network_name and "linear" not in network_name and shared.sd_model_type != "sd3":
                 continue
             module.network_layer_name = network_name
+    if hasattr(sd_model, 'transformer_ref') and sd_model.transformer_ref is not None:
+        for name, module in sd_model.transformer_ref.named_modules():
+            network_name = "lora_transformer_" + name.replace(".", "_")
+            network_layer_mapping[network_name] = module
+            if "norm" in network_name and "linear" not in network_name and shared.sd_model_type != "sd3":
+                continue
+            module.network_layer_name = network_name
     if hasattr(sd_model, 'llm_adapter') and sd_model.llm_adapter is not None:
         for name, module in sd_model.llm_adapter.named_modules():
             network_name = "lora_llm_adapter_" + name.replace(".", "_")

modules/lora/lora_load.py

@@ -56,12 +56,19 @@ def lora_dump(lora, dct):
 def load_safetensors(name, network_on_disk: network.NetworkOnDisk) -> network.Network | None:
     if not shared.sd_loaded:
         return None
-
     sd_model = getattr(shared.sd_model, "pipe", shared.sd_model)
+
+    # cached
     cached = lora_cache.get(name, None)
     if cached is not None:
+        if l.debug:
+            log.trace(f'LoRA load: name="{name}" fn="{network_on_disk.filename}" cache=True')
         return cached
+
+    # native dispatch
     native_module = NATIVE_DISPATCH.get(shared.sd_model_type)
+    if l.debug:
+        log.trace(f'LoRA load: name="{name}" fn="{network_on_disk.filename}" native={native_module}')
     if native_module is not None:
         import importlib
         mod = importlib.import_module(native_module)
@@ -69,6 +76,10 @@ def load_safetensors(name, network_on_disk: network.NetworkOnDisk) -> network.Ne
         if net is not None:
             lora_cache[name] = net
         return net
+
+    # fallback to standard network loading
+    if l.debug:
+        log.trace(f'LoRA load: name="{name}" network_on_disk="{network_on_disk.filename}" safetensors')
     net = network.Network(name, network_on_disk)
     net.mtime = os.path.getmtime(network_on_disk.filename)
     state_dict = sd_models.read_state_dict(network_on_disk.filename, what='network')
@@ -96,7 +107,7 @@ def load_safetensors(name, network_on_disk: network.NetworkOnDisk) -> network.Ne
             emb_dict[vec_name] = weight
             bundle_embeddings[emb_name] = emb_dict
             continue
-        if parts[0] in ["clip_l","clip_g","t5","unet","transformer"]:
+        if parts[0] in ["clip_l", "clip_g", "t5", "unet", "transformer", "transformer_2"]:
             network_part = []
             while parts and parts[-1] in ["alpha","weight","lora_up","lora_down"]:
                 network_part.insert(0,parts[-1])
@@ -147,7 +158,7 @@ def load_safetensors(name, network_on_disk: network.NetworkOnDisk) -> network.Ne
     if len(keys_failed_to_match) > 0:
         log.warning(f'Network load: type=LoRA name="{name}" type={set(network_types)} unmatched={len(keys_failed_to_match)} matched={len(matched_networks)}')
         if l.debug:
-            log.debug(f'Network load: type=LoRA name="{name}" unmatched={keys_failed_to_match}')
+            log.trace(f'Network load: type=LoRA name="{name}" unmatched={keys_failed_to_match}')
     else:
         log.debug(f'Network load: type=LoRA name="{name}" type={set(network_types)} keys={len(matched_networks)} dtypes={dtypes} fuse={lora_overrides.fuse_native()}:{shared.opts.lora_fuse_diffusers}')
     if len(matched_networks) == 0:
@@ -172,13 +183,13 @@ def maybe_recompile_model(names, te_multipliers):
             if not recompile_model:
                 skip_lora_load = True
                 if len(l.loaded_networks) > 0 and l.debug:
-                    log.debug('Model Compile: Skipping LoRa loading')
+                    log.trace('LoRA: recompile required, skip loading')
                 return recompile_model, skip_lora_load
         else:
             recompile_model = True
             shared.compiled_model_state.lora_model = []
     if l.debug:
-        log.debug(f'Model recompile check: task={sd_models.get_diffusers_task(shared.sd_model)} recompile={recompile_model} load={skip_lora_load}')
+        log.trace(f'LoRA recompile check: task={sd_models.get_diffusers_task(shared.sd_model)} recompile={recompile_model} load={skip_lora_load}')
     if recompile_model:
         current_task = sd_models.get_diffusers_task(shared.sd_model)
         log.debug(f'Compile: task={current_task} force model reload')
@@ -285,7 +296,7 @@ def network_load(names, te_multipliers=None, unet_multipliers=None, dyn_dims=Non
         if network_on_disk is not None:
             shorthash = getattr(network_on_disk, 'shorthash', '').lower()
             if l.debug:
-                log.debug(f'Network load: type=LoRA name="{name}" file="{network_on_disk.filename}" hash="{shorthash}" cached={name in lora_cache}')
+                log.trace(f'LoRA: name="{name}" fn="{network_on_disk.filename}" hash="{shorthash}" cached={name in lora_cache}')
             try:
                 lora_scale = te_multipliers[i] if te_multipliers else shared.opts.extra_networks_default_multiplier
                 lora_module = lora_modules[i] if lora_modules and len(lora_modules) > i else None
@@ -330,8 +341,8 @@ def network_load(names, te_multipliers=None, unet_multipliers=None, dyn_dims=Non
         try:
             t1 = time.time()
             if l.debug:
-                log.trace(f'Network load: type=LoRA list={sd_model.get_list_adapters()}')
-                log.trace(f'Network load: type=LoRA active={sd_model.get_active_adapters()}')
+                log.trace(f'LoRA: list={sd_model.get_list_adapters()}')
+                log.trace(f'LoRA: active={sd_model.get_active_adapters()}')
             sd_model.set_adapters(adapter_names=lora_diffusers.diffuser_loaded, adapter_weights=lora_diffusers.diffuser_scales)
             sd_model.enable_lora() # set_adapters does not clear the disabled flag left by a prior removal
         except Exception as e:
@@ -359,7 +370,7 @@ def network_load(names, te_multipliers=None, unet_multipliers=None, dyn_dims=Non
         networks.network_activate()
 
     if len(l.loaded_networks) > 0 and l.debug:
-        log.debug(f'Network load: type=LoRA loaded={[n.name for n in l.loaded_networks]} cache={list(lora_cache)} fuse={lora_overrides.fuse_native()}:{shared.opts.lora_fuse_diffusers}')
+        log.trace(f'LoRA: loaded={[n.name for n in l.loaded_networks]} cache={list(lora_cache)} fuse={lora_overrides.fuse_native()}:{shared.opts.lora_fuse_diffusers}')
 
     if recompile_model:
         log.info("Network load: type=LoRA model recompile required")

modules/lora/lora_overrides.py

@@ -1,4 +1,9 @@
+import os
 from modules import shared
+from modules.logger import log
+
+
+debug_log = log.debug if os.environ.get('SD_LORA_DEBUG', None) is not None else lambda *args, **kwargs: None
 
 
 force_hashes_diffusers = [ # forced always
@@ -116,15 +121,22 @@ def disable_fuse():
     from modules.lora import lora_common as l
     from modules.lora import lora_stack
     if lora_stack.select_possible(len(l.loaded_networks)) or lora_stack.select_engaged():
+        debug_log('LoRA: fuse=False reason="active select mode"')
         return True # select flips per-layer winners against the pristine backup; a dormant select mode leaves fuse alone
     sd_model = getattr(shared.sd_model, 'pipe', shared.sd_model)
     if is_quantized(sd_model):
+        debug_log('LoRA: fuse=False reason="model is quantized"')
         return True
     if any(is_quantized(getattr(sd_model, name, None)) for name in fuse_components(sd_model)):
+        debug_log('LoRA: fuse=False reason="component is quantized"')
         return True
     if hasattr(sd_model, '_lora_partial'):
+        debug_log('LoRA: fuse=False reason="partial lora applied"')
         return True
-    return shared.sd_model_type in fuse_ignore
+    if shared.sd_model_type in fuse_ignore:
+        debug_log(f'LoRA: fuse=False reason="model type {shared.sd_model_type} in fuse_ignore"')
+        return True
+    return False
 
 
 def fuse_native():
@@ -134,4 +146,7 @@ def fuse_native():
     the backup, activate and deactivate passes, since backup mode restores from a
     stored tensor while fuse mode restores by subtracting the delta.
     """
-    return shared.opts.lora_fuse_native and not disable_fuse()
+    result = shared.opts.lora_fuse_native and not disable_fuse()
+    force = os.environ.get('SD_LORA_FUSE', None) is not None
+    debug_log(f'LoRA: native fuse={result} force={force}')
+    return (result or force)

modules/lora/networks.py

@@ -56,7 +56,7 @@ from modules.logger import log, console
 applied_layers: list[str] = []
 refused_writes: int = 0 # deltas the modules would not take on the last activate pass; infotext reports the network as partial
 native_active: bool = False
-default_components = ['text_encoder', 'text_encoder_2', 'text_encoder_3', 'text_encoder_4', 'unet', 'transformer', 'transformer_2', 'llm_adapter']
+default_components = ['text_encoder', 'text_encoder_2', 'text_encoder_3', 'text_encoder_4', 'unet', 'transformer', 'transformer_2', 'transformer_ref', 'llm_adapter']
 
 
 class ActivationPass: