Composable LoRA

extensions-builtin/Lora/extra_networks_lora.py

@@ -1,8 +1,43 @@
 import time
+import numpy as np
+import re
 import networks
 import lora_patches
 from modules import extra_networks, shared
 
+# from https://github.com/cheald/sd-webui-loractl/blob/master/loractl/lib/utils.py
+def get_stepwise(param, step, steps):
+    def sorted_positions(raw_steps):
+        steps = [[float(s.strip()) for s in re.split("[@~]", x)]
+                 for x in re.split("[,;]", str(raw_steps))]
+        # If we just got a single number, just return it
+        if len(steps[0]) == 1:
+            return steps[0][0]
+
+        # Add implicit 1s to any steps which don't have a weight
+        steps = [[s[0], s[1] if len(s) == 2 else 1] for s in steps]
+
+        # Sort by index
+        steps.sort(key=lambda k: k[1])
+
+        steps = [list(v) for v in zip(*steps)]
+        return steps
+
+    def calculate_weight(m, step, max_steps, step_offset=2):
+        if isinstance(m, list):
+            if m[1][-1] <= 1.0:
+                if max_steps > 0:
+                    step = (step) / (max_steps - step_offset)
+                else:
+                    step = 1.0
+            else:
+                step = step
+            v = np.interp(step, m[1], m[0])
+            return v
+        else:
+            return m
+    return calculate_weight(sorted_positions(param), step, steps)
+
 
 class ExtraNetworkLora(extra_networks.ExtraNetwork):
 
@@ -14,7 +49,7 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork):
 
         """mapping of network names to the number of errors the network had during operation"""
 
-    def activate(self, p, params_list):
+    def activate(self, p, params_list, step=0):
         t0 = time.time()
         self.errors.clear()
         if len(params_list) > 0:
@@ -29,13 +64,21 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork):
         for params in params_list:
             assert params.items
             names.append(params.positional[0])
-            te_multiplier = float(params.positional[1]) if len(params.positional) > 1 else 1.0
-            te_multiplier = float(params.named.get("te", te_multiplier))
-            unet_multiplier = [float(params.positional[2]) if len(params.positional) > 2 else te_multiplier] * 3
-            unet_multiplier = [float(params.named.get("unet", unet_multiplier[0]))] * 3
-            unet_multiplier[0] = float(params.named.get("in", unet_multiplier[0]))
-            unet_multiplier[1] = float(params.named.get("mid", unet_multiplier[1]))
-            unet_multiplier[2] = float(params.named.get("out", unet_multiplier[2]))
+            te_multiplier = params.named.get("te", params.positional[1] if len(params.positional) > 1 else 1.0)
+            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[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])
+            for i in range(len(unet_multiplier)):
+                if isinstance(unet_multiplier[i], str) and "@" in unet_multiplier[i]:
+                    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
             te_multipliers.append(te_multiplier)

modules/extra_networks.py

@@ -1,6 +1,6 @@
 import re
 from collections import defaultdict
-from modules import errors
+from modules import errors, shared
 
 
 extra_network_registry = {}
@@ -62,17 +62,29 @@ class ExtraNetwork:
         raise NotImplementedError
 
 
-def activate(p, extra_network_data):
+def is_stepwise(en_obj):
+    all_args = []
+    for en in en_obj:
+        all_args.extend(en.positional[1:])
+        all_args.extend(en.named.values())
+    return any([len(str(x).split("@")) > 1 for x in all_args])
+
+def activate(p, extra_network_data, step=0):
     """call activate for extra networks in extra_network_data in specified order, then call activate for all remaining registered networks with an empty argument list"""
     if extra_network_data is None:
         return
+    stepwise = False
+    for extra_network_args in extra_network_data.values():
+        stepwise = stepwise or is_stepwise(extra_network_args)
+    functional = shared.opts.lora_functional
+    shared.opts.lora_functional = stepwise or functional
     for extra_network_name, extra_network_args in extra_network_data.items():
         extra_network = extra_network_registry.get(extra_network_name, None)
         if extra_network is None:
             errors.log.warning(f"Skipping unknown extra network: {extra_network_name}")
             continue
         try:
-            extra_network.activate(p, extra_network_args)
+            extra_network.activate(p, extra_network_args, step=step)
         except Exception as e:
             errors.display(e, f"activating extra network: name={extra_network_name} args:{extra_network_args}")
 
@@ -84,6 +96,9 @@ def activate(p, extra_network_data):
             extra_network.activate(p, [])
         except Exception as e:
             errors.display(e, f"activating extra network: name={extra_network_name}")
+    if stepwise:
+        p.extra_network_data = extra_network_data
+        shared.opts.lora_functional = functional
 
 
 def deactivate(p, extra_network_data):

modules/processing_diffusers.py

@@ -8,7 +8,7 @@ import numpy as np
 import torch
 import torchvision.transforms.functional as TF
 import diffusers
-from modules import shared, devices, processing, sd_samplers, sd_models, images, errors, prompt_parser_diffusers, sd_hijack_hypertile, processing_correction, processing_vae, sd_models_compile
+from modules import shared, devices, processing, sd_samplers, sd_models, images, errors, prompt_parser_diffusers, sd_hijack_hypertile, processing_correction, processing_vae, sd_models_compile, extra_networks
 from modules.processing_helpers import resize_init_images, resize_hires, fix_prompts, calculate_base_steps, calculate_hires_steps, calculate_refiner_steps
 from modules.onnx_impl import preprocess_pipeline as preprocess_onnx_pipeline, check_parameters_changed as olive_check_parameters_changed
 
@@ -73,6 +73,14 @@ def process_diffusers(p: processing.StableDiffusionProcessing):
                 if shared.state.interrupted or shared.state.skipped:
                     raise AssertionError('Interrupted...')
                 time.sleep(0.1)
+        if hasattr(p, "extra_network_data"):
+            if shared.opts.lora_force_diffusers:
+                shared.log.warning("Composable LoRA not compatible with 'lora_force_diffusers'")
+            else:
+                functional = shared.opts.lora_functional
+                shared.opts.lora_functional = True
+                extra_networks.activate(p, p.extra_network_data, step=step)
+                shared.opts.lora_functional = functional
         if latents is None:
             return kwargs
         elif shared.opts.nan_skip: