allow lora specifier to select model component. see changelog/wiki for docs on how to.

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

CHANGELOG.md

@@ -130,6 +130,13 @@ And check out new **history** tab in the right panel, it now shows visualization
     *experimental*: requires new pydantic package which *may* break other things, to enable start sdnext with `--new` flag  
     *note*: this is model quantization only, no support for tensorRT inference yet  
 - **Other**
+  - **LoRA** allow specifying module to apply lora on  
+    *example*: `<lora:mylora:1.0:module=unet>` would apply lora *only* on unet regardless of lora content  
+    this is particularly useful when you have multiple loras and you want to apply them on different parts of the model  
+    *example*: `<lora:firstlora:1.0:high>` and `<lora:secondlora:1.0:low>`  
+    *note*: `low` is shorthand for `module=transformer_2` and `high` is shortcut for `module=transformer`  
+  - **Detailer** allow manually setting processing resolution  
+    *note*: this does not impact the actual image resolution, only the resolution at which detailer internally operates  
   - refactor reuse-seed and add functionality to all tabs  
   - refactor modernui js codebase  
   - move zluda flash attenion to *Triton Flash attention* option  
@@ -144,8 +151,6 @@ And check out new **history** tab in the right panel, it now shows visualization
   - add deprecation warning for `python==3.9`  
   - allow setting denoise strength to 0 in control/img2img  
     this allows to run workflows which only refine or detail existing image without changing it   
-  - **Detailer** allow manually setting processing resolution  
-    *note*: this does not impact the actual image resolution, only the resolution at which detailer internally operates  
 - **Fixes**
   - normalize path hanlding when deleting images  
   - unified compile upscalers  

modules/cachedit.py

@@ -39,12 +39,16 @@ def apply_cache_dit(pipe):
     else:
         calibrator_config = None
     shared.log.info(f'Apply Cache-DiT: config="{cache_config.strify()}" calibrator="{calibrator_config.strify() if calibrator_config else "None"}"')
-    cache_dit.enable_cache(
-        pipe,
-        cache_config=cache_config,
-        calibrator_config=calibrator_config,
-    )
-    shared.sd_model.has_cache_dit = True
+    try:
+        cache_dit.enable_cache(
+            pipe,
+            cache_config=cache_config,
+            calibrator_config=calibrator_config,
+        )
+        shared.sd_model.has_cache_dit = True
+    except Exception as e:
+        shared.log.error(f'Cache-DiT: {e}')
+        return
 
 
 def unapply_cache_dir(pipe):

modules/lora/extra_networks_lora.py

@@ -2,7 +2,7 @@ from typing import List
 import os
 import re
 import numpy as np
-from modules.lora import networks, lora_overrides, lora_load
+from modules.lora import networks, lora_overrides, lora_load, lora_diffusers
 from modules.lora import lora_common as l
 from modules import extra_networks, shared, sd_models
 
@@ -88,16 +88,21 @@ def parse(p, params_list, step=0):
     te_multipliers = []
     unet_multipliers = []
     dyn_dims = []
+    lora_modules = []
     for params in params_list:
-        assert params.items
         names.append(params.positional[0])
-        te_multiplier = params.named.get("te", params.positional[1] if len(params.positional) > 1 else shared.opts.extra_networks_default_multiplier)
+
+        default_multiplier = params.positional[1] if len(params.positional) > 1 else shared.opts.extra_networks_default_multiplier
+        if isinstance(default_multiplier, str) and "@" not in default_multiplier:
+            default_multiplier = shared.opts.extra_networks_default_multiplier
+
+        te_multiplier = params.named.get("te", default_multiplier)
         if isinstance(te_multiplier, str) and "@" in te_multiplier:
             te_multiplier = get_stepwise(te_multiplier, step, p.steps)
         else:
             te_multiplier = float(te_multiplier)
-        unet_multiplier = [params.positional[2] if len(params.positional) > 2 else te_multiplier] * 3
-        unet_multiplier = [params.named.get("unet", unet_multiplier[0])] * 3
+
+        unet_multiplier = 3 * [params.named.get("unet", te_multiplier)] # fill all 3 with same value
         unet_multiplier[0] = params.named.get("in", unet_multiplier[0])
         unet_multiplier[1] = params.named.get("mid", unet_multiplier[1])
         unet_multiplier[2] = params.named.get("out", unet_multiplier[2])
@@ -106,12 +111,22 @@ def parse(p, params_list, step=0):
                 unet_multiplier[i] = get_stepwise(unet_multiplier[i], step, p.steps)
             else:
                 unet_multiplier[i] = float(unet_multiplier[i])
-        dyn_dim = int(params.positional[3]) if len(params.positional) > 3 else None
-        dyn_dim = int(params.named["dyn"]) if "dyn" in params.named else dyn_dim
+
+        dyn_dim = int(params.named["dyn"]) if "dyn" in params.named else None
         te_multipliers.append(te_multiplier)
         unet_multipliers.append(unet_multiplier)
         dyn_dims.append(dyn_dim)
-    return names, te_multipliers, unet_multipliers, dyn_dims
+
+        lora_module = []
+        if 'high' in params.positional or 'HIGH 14B' in params.positional[0]:
+            lora_module.append('transformer')
+        if 'low' in params.positional or 'LOW 14B' in params.positional[0]:
+            lora_module.append('transformer_2')
+        if params.named.get('module', None) is not None:
+            lora_module.append(params.named['module'].lower())
+        lora_modules.append(lora_module)
+
+    return names, te_multipliers, unet_multipliers, dyn_dims, lora_modules
 
 
 def unload_diffusers():
@@ -168,7 +183,7 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork):
         if len(params_list) > 0 and not self.active: # activate patches once
             self.active = True
             self.model = shared.opts.sd_model_checkpoint
-        names, te_multipliers, unet_multipliers, dyn_dims = parse(p, params_list, step)
+        names, te_multipliers, unet_multipliers, dyn_dims, lora_modules = parse(p, params_list, step)
         requested = self.signature(names, te_multipliers, unet_multipliers)
 
         load_method = lora_overrides.get_method()
@@ -181,7 +196,7 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork):
             has_changed = False # diffusers handles its own loading
             if len(exclude) == 0:
                 jobid = shared.state.begin('LoRA')
-                lora_load.network_load(names, te_multipliers, unet_multipliers, dyn_dims) # load only on first call
+                lora_load.network_load(names, te_multipliers, unet_multipliers, dyn_dims, lora_modules) # load only on first call
                 sd_models.set_diffuser_offload(shared.sd_model, op="model")
                 shared.state.end(jobid)
         elif load_method == 'nunchaku':
@@ -208,7 +223,7 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork):
                 shared.log.info(f'Network load: type=LoRA apply={[n.name for n in l.loaded_networks]} method={load_method} 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 len(lora_load.diffuser_loaded) > 0:
+        if len(lora_diffusers.diffuser_loaded) > 0:
             if not (shared.compiled_model_state is not None and shared.compiled_model_state.is_compiled is True):
                 unload_diffusers()
         if self.active and l.debug:

modules/lora/lora_diffusers.py

@@ -0,0 +1,92 @@
+from typing import Union
+import os
+import time
+import diffusers
+from modules import shared, errors
+from modules.lora import network
+from modules.lora import lora_common as l
+
+
+diffuser_loaded = []
+diffuser_scales = []
+
+
+def load_per_module(sd_model: diffusers.DiffusionPipeline, filename: str, adapter_name: str, lora_modules: list[str]):
+    shared.log.debug(f'LoRA load: modules={lora_modules}')
+    try:
+        state_dict = sd_model.lora_state_dict(filename)
+        if isinstance(state_dict, tuple) and len(state_dict) == 2:
+            state_dict, network_alphas = state_dict
+        else:
+            network_alphas = {}
+    except Exception as e:
+        shared.log.error(f'LoRA load: {e}')
+        if l.debug:
+            errors.display(e, "LoRA")
+        return None
+    for lora_module in lora_modules:
+        if lora_module == 'transformer':
+            if hasattr(sd_model, 'transformer') and sd_model.transformer is not None:
+                sd_model.load_lora_into_transformer(state_dict, transformer=sd_model.transformer, adapter_name=adapter_name)
+            else:
+                shared.log.warning(f'LoRA load: requested={lora_module} missing')
+        elif lora_module == 'transformer_2':
+            if hasattr(sd_model, 'transformer_2') and sd_model.transformer_2 is not None:
+                sd_model.load_lora_into_transformer(state_dict, transformer=sd_model.transformer_2, adapter_name=adapter_name)
+            else:
+                shared.log.warning(f'LoRA load: requested={lora_module} missing')
+        elif lora_module == 'unet':
+            if hasattr(sd_model, 'unet') and sd_model.unet is not None:
+                sd_model.load_lora_into_unet(state_dict, network_alphas, unet=sd_model.unet, adapter_name=adapter_name)
+            else:
+                shared.log.warning(f'LoRA load: requested={lora_module} missing')
+        elif lora_module == 'text_encoder' or lora_module == 'te':
+            if hasattr(sd_model, 'text_encoder') and sd_model.text_encoder is not None:
+                sd_model.load_lora_into_text_encoder(state_dict, network_alphas, text_encoder=sd_model.text_encoder, adapter_name=adapter_name)
+            else:
+                shared.log.warning(f'LoRA load: requested={lora_module} missing')
+        else:
+            shared.log.warning(f'LoRA load: requested={lora_module} unknown')
+    return adapter_name
+
+
+def load_diffusers(name: str, network_on_disk: network.NetworkOnDisk, lora_scale:float=shared.opts.extra_networks_default_multiplier, lora_module=None) -> Union[network.Network, None]:
+    t0 = time.time()
+    name = name.replace(".", "_")
+    sd_model: diffusers.DiffusionPipeline = getattr(shared.sd_model, "pipe", shared.sd_model)
+    shared.log.debug(f'Network load: type=LoRA name="{name}" file="{network_on_disk.filename}" detected={network_on_disk.sd_version} method=diffusers scale={lora_scale} fuse={shared.opts.lora_fuse_diffusers}')
+    if not hasattr(sd_model, 'load_lora_weights'):
+        shared.log.error(f'Network load: type=LoRA class={sd_model.__class__} does not implement load lora')
+        return None
+    try:
+        if lora_module is not None and isinstance(lora_module, list) and len(lora_module) > 0:
+            name = load_per_module(sd_model, network_on_disk.filename, adapter_name=name, lora_modules=lora_module)
+        else:
+            sd_model.load_lora_weights(network_on_disk.filename, adapter_name=name)
+    except Exception as e:
+        if 'already in use' in str(e):
+            pass
+        else:
+            if 'following keys have not been correctly renamed' in str(e):
+                shared.log.error(f'Network load: type=LoRA name="{name}" diffusers unsupported format')
+            elif 'object has no attribute' in str(e):
+                shared.log.error(f'Network load: type=LoRA name="{name}" diffusers empty module')
+            else:
+                shared.log.error(f'Network load: type=LoRA name="{name}" {e}')
+            if l.debug:
+                errors.display(e, "LoRA")
+            return None
+    if name is None:
+        return None
+    if name not in diffuser_loaded:
+        list_adapters = sd_model.get_list_adapters()
+        list_adapters = [adapter for adapters in list_adapters.values() for adapter in adapters]
+        if name not in list_adapters:
+            shared.log.error(f'Network load: type=LoRA name="{name}" adapters={list_adapters} not loaded')
+        else:
+            diffuser_loaded.append(name)
+            diffuser_scales.append(lora_scale)
+    net = network.Network(name, network_on_disk)
+    net.mtime = os.path.getmtime(network_on_disk.filename)
+    l.timer.activate += time.time() - t0
+    return net

modules/lora/lora_load.py

@@ -3,12 +3,10 @@ import os
 import time
 import concurrent
 from modules import shared, errors, sd_models, sd_models_compile, files_cache
-from modules.lora import network, lora_overrides, lora_convert
+from modules.lora import network, lora_overrides, lora_convert, lora_diffusers
 from modules.lora import lora_common as l
 
 
-diffuser_loaded = []
-diffuser_scales = []
 lora_cache = {}
 available_networks = {}
 available_network_aliases = {}
@@ -17,41 +15,6 @@ available_network_hash_lookup = {}
 dump_lora_keys = os.environ.get('SD_LORA_DUMP', None) is not None
 
 
-def load_diffusers(name, network_on_disk, lora_scale=shared.opts.extra_networks_default_multiplier) -> Union[network.Network, None]:
-    t0 = time.time()
-    name = name.replace(".", "_")
-    sd_model = getattr(shared.sd_model, "pipe", shared.sd_model)
-    shared.log.debug(f'Network load: type=LoRA name="{name}" file="{network_on_disk.filename}" detected={network_on_disk.sd_version} method=diffusers scale={lora_scale} fuse={shared.opts.lora_fuse_diffusers}')
-    if not hasattr(sd_model, 'load_lora_weights'):
-        shared.log.error(f'Network load: type=LoRA class={sd_model.__class__} does not implement load lora')
-        return None
-    try:
-        sd_model.load_lora_weights(network_on_disk.filename, adapter_name=name)
-    except Exception as e:
-        if 'already in use' in str(e):
-            pass
-        else:
-            if 'The following keys have not been correctly renamed' in str(e):
-                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 l.debug:
-                errors.display(e, "LoRA")
-            return None
-    if name not in diffuser_loaded:
-        list_adapters = sd_model.get_list_adapters()
-        list_adapters = [adapter for adapters in list_adapters.values() for adapter in adapters]
-        if name not in list_adapters:
-            shared.log.error(f'Network load: type=LoRA name="{name}" adapters={list_adapters} not loaded')
-        else:
-            diffuser_loaded.append(name)
-            diffuser_scales.append(lora_scale)
-    net = network.Network(name, network_on_disk)
-    net.mtime = os.path.getmtime(network_on_disk.filename)
-    l.timer.activate += time.time() - t0
-    return net
-
-
 def lora_dump(lora, dct):
     import tempfile
     sd_model = getattr(shared.sd_model, "pipe", shared.sd_model)
@@ -73,7 +36,7 @@ def lora_dump(lora, dct):
             f.write(line + "\n")
 
 
-def load_safetensors(name, network_on_disk) -> Union[network.Network, None]:
+def load_safetensors(name, network_on_disk: network.NetworkOnDisk) -> Union[network.Network, None]:
     if not shared.sd_loaded:
         return None
 
@@ -261,15 +224,15 @@ def gather_networks(names):
     return networks_on_disk
 
 
-def network_load(names, te_multipliers=None, unet_multipliers=None, dyn_dims=None):
+def network_load(names, te_multipliers=None, unet_multipliers=None, dyn_dims=None, lora_modules=None):
     networks_on_disk = gather_networks(names)
     failed_to_load_networks = []
     recompile_model, skip_lora_load = maybe_recompile_model(names, te_multipliers)
     sd_model = getattr(shared.sd_model, "pipe", shared.sd_model)
 
     l.loaded_networks.clear()
-    diffuser_loaded.clear()
-    diffuser_scales.clear()
+    lora_diffusers.diffuser_loaded.clear()
+    lora_diffusers.diffuser_scales.clear()
     t0 = time.time()
 
     for i, (network_on_disk, name) in enumerate(zip(networks_on_disk, names)):
@@ -279,11 +242,13 @@ def network_load(names, te_multipliers=None, unet_multipliers=None, dyn_dims=Non
             if l.debug:
                 shared.log.debug(f'Network load: type=LoRA name="{name}" file="{network_on_disk.filename}" hash="{shorthash}"')
             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
                 if recompile_model:
-                    shared.compiled_model_state.lora_model.append(f"{name}:{te_multipliers[i] if te_multipliers else shared.opts.extra_networks_default_multiplier}")
+                    shared.compiled_model_state.lora_model.append(f"{name}:{lora_scale}")
                 lora_method = lora_overrides.get_method(shorthash)
                 if lora_method == 'diffusers':
-                    net = load_diffusers(name, network_on_disk, lora_scale=te_multipliers[i] if te_multipliers else shared.opts.extra_networks_default_multiplier)
+                    net = lora_diffusers.load_diffusers(name, network_on_disk, lora_scale, lora_module)
                 elif lora_method == 'nunchaku':
                     pass # handled directly from extra_networks_lora.load_nunchaku
                 else:
@@ -311,21 +276,21 @@ def network_load(names, te_multipliers=None, unet_multipliers=None, dyn_dims=Non
         name = next(iter(lora_cache))
         lora_cache.pop(name, None)
 
-    if not skip_lora_load and len(diffuser_loaded) > 0:
-        shared.log.debug(f'Network load: type=LoRA loaded={diffuser_loaded} available={sd_model.get_list_adapters()} active={sd_model.get_active_adapters()} scales={diffuser_scales}')
+    if not skip_lora_load and len(lora_diffusers.diffuser_loaded) > 0:
+        shared.log.debug(f'Network load: type=LoRA loaded={lora_diffusers.diffuser_loaded} available={sd_model.get_list_adapters()} active={sd_model.get_active_adapters()} scales={lora_diffusers.diffuser_scales}')
         try:
             t1 = time.time()
             if l.debug:
                 shared.log.trace(f'Network load: type=LoRA list={sd_model.get_list_adapters()}')
                 shared.log.trace(f'Network load: type=LoRA active={sd_model.get_active_adapters()}')
-            sd_model.set_adapters(adapter_names=diffuser_loaded, adapter_weights=diffuser_scales)
+            sd_model.set_adapters(adapter_names=lora_diffusers.diffuser_loaded, adapter_weights=lora_diffusers.diffuser_scales)
         except Exception as e:
             shared.log.error(f'Network load: type=LoRA action=set {e}')
             if l.debug:
                 errors.display(e, 'LoRA')
         try:
             if shared.opts.lora_fuse_diffusers and not lora_overrides.disable_fuse():
-                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
+                sd_model.fuse_lora(adapter_names=lora_diffusers.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
                 sd_model.unload_lora_weights()
             l.timer.activate += time.time() - t1
         except Exception as e:

pipelines/model_wanai.py

@@ -65,16 +65,20 @@ def load_wan(checkpoint_info, diffusers_load_config={}):
     repo_id = sd_models.path_to_repo(checkpoint_info)
     sd_models.hf_auth_check(checkpoint_info)
 
+    boundary_ratio = None
     if 'a14b' in repo_id.lower() or 'fun-14b' in repo_id.lower():
         if shared.opts.model_wan_stage == 'high noise' or shared.opts.model_wan_stage == 'first':
             transformer = load_transformer(repo_id, diffusers_load_config, 'transformer')
             transformer_2 = None
+            boundary_ratio = 0.0
         elif shared.opts.model_wan_stage == 'low noise' or shared.opts.model_wan_stage == 'second':
-            transformer = load_transformer(repo_id, diffusers_load_config, 'transformer_2')
-            transformer_2 = None
+            transformer = None
+            transformer_2 = load_transformer(repo_id, diffusers_load_config, 'transformer_2')
+            boundary_ratio = 1.0
         elif shared.opts.model_wan_stage == 'combined' or shared.opts.model_wan_stage == 'both':
             transformer = load_transformer(repo_id, diffusers_load_config, 'transformer')
             transformer_2 = load_transformer(repo_id, diffusers_load_config, 'transformer_2')
+            boundary_ratio = shared.opts.model_wan_boundary
         else:
             shared.log.error(f'Load model: type=WanAI stage="{shared.opts.model_wan_stage}" unsupported')
             return None
@@ -85,7 +89,6 @@ def load_wan(checkpoint_info, diffusers_load_config={}):
     text_encoder = load_text_encoder(repo_id, diffusers_load_config)
 
     load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, module='Model')
-    boundary_ratio = shared.opts.model_wan_boundary if transformer_2 is not None else None
 
     if 'Wan2.2-I2V' in repo_id:
         pipe_cls = diffusers.WanImageToVideoPipeline