Merge pull request #4987 from vladmandic/fix/lora-te-fixes

Fix/lora te fixes

modules/detailer/detailer.py

@@ -247,7 +247,7 @@ class Detailer():
                 pc.disable_extra_networks = True # disable processing_diffusers from handling network activation since its handled here
                 network_same = len(p.network_data.values()) == len(pc.network_data.values()) and all(x == y for x, y in zip(p.network_data.values(), pc.network_data.values()))
                 if not network_same:
-                    extra_networks.activate(pc, pc.network_data)
+                    extra_networks.activate_filtered(pc, pc.network_data)
                 log.debug(f'Detail: model="{i+1}:{name}" item={j+1}/{len(items)} box={item.box} label="{item.label}" score={item.score:.2f} seg={detailer_opt(p, "detailer_segmentation")} network={network_same} prompt="{pc.prompt}"')
                 pc.init_images = [image]
                 pc.image_mask = [item.mask]

modules/extra_networks.py

@@ -121,6 +121,15 @@ def activate(p: StableDiffusionProcessing, extra_network_data: defaultdict[str,
     p.network_data = extra_network_data
 
 
+def activate_filtered(p: StableDiffusionProcessing, extra_network_data: defaultdict[str, list[ExtraNetworkParams]] | None = None, step=0):
+    """activate with text encoder components gated on lora_apply_te; must run before prompt encode so te networks affect embeds"""
+    apply_te = getattr(p, 'lora_apply_te', None)
+    if apply_te is None:
+        apply_te = shared.opts.lora_apply_te
+    exclude = [] if apply_te else ['text_encoder', 'text_encoder_2', 'text_encoder_3']
+    activate(p, extra_network_data, step=step, exclude=exclude)
+
+
 def deactivate(p: StableDiffusionProcessing, extra_network_data: defaultdict[str, list[ExtraNetworkParams]] | None = None, force: bool | None = None):
     """call deactivate for extra networks in extra_network_data in specified order, then call deactivate for all remaining registered networks"""
     if p.disable_extra_networks:

modules/face/faceid.py

@@ -216,7 +216,7 @@ def face_id(
                 p.subseeds = p.all_subseeds[n * p.batch_size:(n+1) * p.batch_size]
                 p.prompts, p.network_data = extra_networks.parse_prompts(p.prompts, p.network_data)
 
-                extra_networks.activate(p, p.network_data)
+                extra_networks.activate_filtered(p, p.network_data)
                 ip_model_dict.update({
                         "prompt": p.prompts[0],
                         "negative_prompt": p.negative_prompts[0],

modules/lora/extra_networks_lora.py

@@ -191,11 +191,13 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork):
         key = f'include={",".join(include)}:exclude={",".join(exclude)}'
         loaded = sd_model.loaded_loras.get(key, [])
         if len(requested) != len(loaded):
+            sd_model.loaded_loras.clear() # single-entry cache: any activation invalidates state recorded under other filter keys
             sd_model.loaded_loras[key] = requested
             debug_log(f'Network check: type=LoRA key="{key}" requested={requested} loaded={loaded} status="num changed"')
             return True, "num changed"
         for req, load in zip(requested, loaded, strict=False):
             if req != load:
+                sd_model.loaded_loras.clear()
                 sd_model.loaded_loras[key] = requested
                 debug_log(f'Network check: type=LoRA key="{key}" requested={requested} loaded={loaded} status="content changed"')
                 return True, "content changed"
@@ -229,6 +231,14 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork):
             if has_changed:
                 jobid = shared.state.begin('LoRA')
                 lora_load.network_load(names, te_multipliers, unet_multipliers, dyn_dims, lora_modules) # load only on first call
+                if len(names) == 0: # removal disables adapters in place, unload_lora_weights would unwrap modules and detach offload hooks
+                    sd_model = getattr(shared.sd_model, "pipe", shared.sd_model)
+                    if hasattr(sd_model, 'disable_lora'):
+                        try:
+                            sd_model.disable_lora()
+                            log.info('Network unload: type=LoRA mode=diffusers')
+                        except Exception as e:
+                            log.error(f'Network unload: type=LoRA {e}')
                 sd_models.set_diffuser_offload(shared.sd_model, op="model")
                 shared.state.end(jobid)
 
@@ -237,7 +247,7 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork):
             has_changed = lora_nunchaku.load_nunchaku(names, unet_multipliers)
 
         else: # native
-            lora_load.network_load(names, te_multipliers, unet_multipliers, dyn_dims) # load
+            lora_load.network_load(names, te_multipliers, unet_multipliers, dyn_dims, activate=False) # load only, activation below honors include/exclude
             has_changed, reason = self.changed(requested, include, exclude)
             if has_changed:
                 jobid = shared.state.begin('LoRA')

modules/lora/lora_convert.py

@@ -130,6 +130,12 @@ class KeyConvert:
         sd_module = shared.sd_model.network_layer_mapping.get(key, None)
         if sd_module is None:
             sd_module = shared.sd_model.network_layer_mapping.get(key.replace("guidance", "timestep"), None)  # FLUX1 fix
+        if sd_module is None and key.startswith("lora_te"):
+            # transformers >=5.6 flattened CLIPTextModel; kohya te keys still carry the text_model wrapper
+            flat_key = key.replace("_text_model_", "_", 1)
+            sd_module = shared.sd_model.network_layer_mapping.get(flat_key, None)
+            if sd_module is not None:
+                key = flat_key
         if debug and sd_module is None:
             raise RuntimeError(f"LoRA key not found in network_layer_mapping: key={key} mapping={shared.sd_model.network_layer_mapping.keys()}")
         return key, sd_module

modules/lora/lora_extract.py

@@ -142,14 +142,17 @@ def make_lora(fn, maxrank, auto_rank, rank_ratio, modules, overwrite):
 
         if 'te' in modules and getattr(shared.sd_model, 'text_encoder', None) is not None:
             task = progress.add_task(description="te1 decompose", total=len(list(shared.sd_model.text_encoder.named_modules())))
+            # transformers >=5.6 flattened CLIPTextModel; kohya naming keeps the text_model wrapper
+            flattened_clip = 'CLIPTextModel' in shared.sd_model.text_encoder.__class__.__name__ and not hasattr(shared.sd_model.text_encoder, 'text_model')
             for name, module in shared.sd_model.text_encoder.named_modules():
                 progress.update(task, advance=1)
                 weights_backup = getattr(module, "network_weights_backup", None)
                 if weights_backup is None or getattr(module, "network_current_names", None) is None:
                     continue
                 prefix = "lora_te1_" if hasattr(shared.sd_model, 'text_encoder_2') else "lora_te_"
+                key_name = f'text_model.{name}' if flattened_clip else name
                 module.svdhandler = SVDHandler(maxrank, rank_ratio)
-                module.svdhandler.network_name = prefix + name.replace(".", "_")
+                module.svdhandler.network_name = prefix + key_name.replace(".", "_")
                 with devices.inference_context():
                     module.svdhandler.decompose(module.weight, weights_backup)
             progress.remove_task(task)

modules/lora/lora_load.py

@@ -263,7 +263,7 @@ def gather_networks(names):
     return networks_on_disk
 
 
-def network_load(names, te_multipliers=None, unet_multipliers=None, dyn_dims=None, lora_modules=None):
+def network_load(names, te_multipliers=None, unet_multipliers=None, dyn_dims=None, lora_modules=None, activate=True):
     networks_on_disk = gather_networks(names)
     failed_to_load_networks = []
     recompile_model, skip_lora_load = maybe_recompile_model(names, te_multipliers)
@@ -324,6 +324,7 @@ def network_load(names, te_multipliers=None, unet_multipliers=None, dyn_dims=Non
                 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()}')
             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:
             if str(e) not in exclude_errors:
                 log.error(f'Network load: type=LoRA action=strength {str(e)}')
@@ -342,9 +343,10 @@ def network_load(names, te_multipliers=None, unet_multipliers=None, dyn_dims=Non
 
     # Activate native modules loaded via diffusers path (e.g., LoKR on Flux2)
     # Also restore backed-up weights when previously active native modules are removed
+    # Callers that run their own deactivate/activate sequence pass activate=False
     from modules.lora import networks
     native_nets = [net for net in l.loaded_networks if len(net.modules) > 0]
-    if native_nets or networks.native_active:
+    if activate and (native_nets or networks.native_active):
         networks.network_activate()
 
     if len(l.loaded_networks) > 0 and l.debug:

modules/lora/network.py

@@ -181,7 +181,7 @@ class NetworkModule:
 
     def multiplier(self):
         unet_multiplier = 3 * [self.network.unet_multiplier] if not isinstance(self.network.unet_multiplier, list) else self.network.unet_multiplier
-        if 'transformer' in self.sd_key[:20]:
+        if self.sd_key.startswith('lora_te') or 'transformer' in self.sd_key[:20]:
             return self.network.te_multiplier
         if "down_blocks" in self.sd_key:
             return unet_multiplier[0]

modules/lora/networks.py

@@ -28,11 +28,13 @@ def network_activate(include=None, exclude=None):
         modules = {}
         components = include if len(include) > 0 else default_components
         components = [x for x in components if x not in exclude]
+        filtered_components = [x for x in default_components if x not in components] # filtered components restore to backup so a filter means detached, not frozen with stale weights
         active_components = []
-        for name in components:
+        for name in components + filtered_components:
             component = getattr(sd_model, name, None)
             if component is not None and hasattr(component, 'named_modules'):
-                active_components.append(name)
+                if name in components:
+                    active_components.append(name)
                 modules[name] = list(component.named_modules())
         total = sum(len(x) for x in modules.values())
         if len(l.loaded_networks) > 0:
@@ -48,16 +50,25 @@ def network_activate(include=None, exclude=None):
             applied_layers.clear()
             backup_size = 0
             for component in modules.keys():
+                component_wanted = wanted_names if component in components else ()
                 device = getattr(sd_model, component, None).device
                 for _, module in modules[component]:
                     network_layer_name = getattr(module, 'network_layer_name', None)
                     current_names = getattr(module, "network_current_names", ())
-                    if getattr(module, 'weight', None) is None or shared.state.interrupted or (network_layer_name is None) or (current_names == wanted_names):
+                    if getattr(module, 'weight', None) is None or shared.state.interrupted or (network_layer_name is None) or (current_names == component_wanted):
                         if task is not None:
                             pbar.update(task, advance=1)
                         continue
-                    backup_size += network_backup_weights(module, network_layer_name, wanted_names)
-                    batch_updown, batch_ex_bias = network_calc_weights(module, network_layer_name, elimit=elimit)
+                    backup_size += network_backup_weights(module, network_layer_name, component_wanted)
+                    if component_wanted == ():
+                        weights_backup = getattr(module, "network_weights_backup", None)
+                        if weights_backup is None or isinstance(weights_backup, bool): # fuse mode has no tensor backup, restore stays with network_deactivate
+                            if task is not None:
+                                pbar.update(task, advance=1)
+                            continue
+                        batch_updown, batch_ex_bias = None, None # restore-only pass, apply with no weights reverts to backup
+                    else:
+                        batch_updown, batch_ex_bias = network_calc_weights(module, network_layer_name, elimit=elimit)
                     if shared.opts.lora_fuse_native:
                         network_apply_direct(module, batch_updown, batch_ex_bias, device=device)
                     else:
@@ -68,7 +79,7 @@ def network_activate(include=None, exclude=None):
                         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
+                    module.network_current_names = component_wanted
                     if task is not None:
                         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}')

modules/processing_args.py

@@ -6,7 +6,7 @@ import inspect
 import torch
 import numpy as np
 from PIL import Image
-from modules import shared, sd_models, processing, processing_vae, processing_helpers, sd_hijack_hypertile, extra_networks, sd_vae
+from modules import shared, sd_models, processing, processing_vae, processing_helpers, sd_hijack_hypertile, sd_vae
 from modules.logger import log
 from modules.processing_callbacks import diffusers_callback_legacy, diffusers_callback, set_callbacks_p
 from modules.processing_helpers import get_generator, apply_circular # pylint: disable=unused-import
@@ -241,9 +241,6 @@ def set_pipeline_args(p, model, prompts:list, negative_prompts:list, prompts_2:l
         else:
             args['clip_skip'] = clip_skip - 1
 
-    if shared.opts.lora_apply_te:
-        extra_networks.activate(p, include=['text_encoder', 'text_encoder_2', 'text_encoder_3'])
-
     if 'complex_human_instruction' in possible:
         chi = shared.opts.te_complex_human_instruction
         p.extra_generation_params["CHI"] = chi

modules/processing_diffusers.py

@@ -147,6 +147,8 @@ def process_base(p: processing.StableDiffusionProcessing):
     desc = 'Base'
     if 'detailer' in p.ops:
         desc = 'Detail'
+    p.prompts, p.network_data = extra_networks.parse_prompts(p.prompts, p.network_data)
+    extra_networks.activate_filtered(p) # networks must patch weights before prompt encode so te loras affect embeds
     base_args = set_pipeline_args(
         p=p,
         model=shared.sd_model,
@@ -176,9 +178,6 @@ def process_base(p: processing.StableDiffusionProcessing):
         modelstats.analyze()
     try:
         t0 = time.time()
-        p.prompts, p.network_data = extra_networks.parse_prompts(p.prompts, p.network_data)
-        extra_networks.activate(p, exclude=['text_encoder', 'text_encoder_2', 'text_encoder_3'])
-
         if hasattr(shared.sd_model, 'tgate') and getattr(p, 'gate_step', -1) > 0:
             base_args['gate_step'] = p.gate_step
             output = shared.sd_model.tgate(**base_args) # pylint: disable=not-callable
@@ -311,7 +310,7 @@ def process_hires(p: processing.StableDiffusionProcessing, output):
                 prompts, p.network_data = extra_networks.parse_prompts(prompts)
                 reset_prompts = True
             if reset_prompts or ('base' in p.skip):
-                extra_networks.activate(p)
+                extra_networks.activate_filtered(p)
 
             hires_args = set_pipeline_args(
                 p=p,

modules/prompt_parser_diffusers.py

@@ -130,8 +130,11 @@ class PromptEmbedder:
         # unpack EN data in case of TE LoRA
         en_data = p.network_data
         en_data = [idx.items for item in en_data.values() for idx in item]
+        apply_te = getattr(p, 'lora_apply_te', None)
+        if apply_te is None:
+            apply_te = shared.opts.lora_apply_te
         effective_batch = 1 if self.allsame else self.batchsize
-        key = str([self.prompts, self.negative_prompts, effective_batch, self.clip_skip, self.steps, en_data])
+        key = str([self.prompts, self.negative_prompts, effective_batch, self.clip_skip, self.steps, en_data, apply_te])
         item = cache.get(key)
         if not item:
             if not any(flatten(emb) for emb in [self.prompt_embeds,