Merge pull request #3593 from vladmandic/lora-refactor

Major lora refactor: SD3.5 Large included

CHANGELOG.md

@@ -1,6 +1,6 @@
 # Change Log for SD.Next
 
-## Update for 2024-11-28
+## Update for 2024-12-02
 
 ### New models and integrations
 
@@ -33,6 +33,12 @@
 
 ### UI and workflow improvements
 
+- **LoRA** handler rewrite:  
+  - LoRA weights are no longer calculated on-the-fly during model execution, but are pre-calculated at the start  
+    this results in perceived overhead on generate startup, but results in overall faster execution as LoRA does not need to be processed on each step  
+  - *note*: LoRA weights backups are required so LoRA can be unapplied, but can take quite a lot of system memory  
+    if you know you will not need to unapply LoRA, you can disable backups in *settings -> networks -> lora fuse*  
+    in which case, you need to reload model to unapply LoRA  
 - **Model loader** improvements:  
   - detect model components on model load fail  
   - allow passing absolute path to model loader  
@@ -41,6 +47,10 @@
   - Flux: all-in-one safetensors  
     example: <https://civitai.com/models/646328?modelVersionId=1040235>  
   - Flux: do not recast quants  
+- **Offload** improvements:  
+  - faster and more compatible *balanced* mode  
+  - balanced offload: units are now in percentage instead of bytes  
+  - balanced offload: add both high and low watermark  
 - **UI**:  
   - improved stats on generate completion  
   - improved live preview display and performance  
@@ -54,6 +64,7 @@
 - **Sampler** improvements  
   - Euler FlowMatch: add sigma methods (*karras/exponential/betas*)  
   - DPM FlowMatch: update all and add sigma methods  
+  - BDIA-DDIM: *experimental*  
 
 ### Fixes  
 
@@ -68,7 +79,9 @@
 - fix xyz-grid with lora  
 - fix api script callbacks  
 - fix gpu memory monitoring  
-- simplify img2img/inpaint/sketch canvas handling
+- simplify img2img/inpaint/sketch canvas handling  
+- fix prompt caching  
+- fix xyz grid skip final pass  
 
 ## Update for 2024-11-21
 

extensions-builtin/Lora/network_overrides.py

@@ -26,7 +26,6 @@ force_diffusers = [ # forced always
 
 force_models = [ # forced always
     'sc',
-    # 'sd3',
     'kandinsky',
     'hunyuandit',
     'auraflow',

extensions-builtin/Lora/scripts/lora_script.py

@@ -5,7 +5,7 @@ from lora_extract import create_ui
 from network import NetworkOnDisk
 from ui_extra_networks_lora import ExtraNetworksPageLora
 from extra_networks_lora import ExtraNetworkLora
-from modules import script_callbacks, extra_networks, ui_extra_networks, ui_models # pylint: disable=unused-import
+from modules import script_callbacks, extra_networks, ui_extra_networks, ui_models, shared # pylint: disable=unused-import
 
 
 re_lora = re.compile("<lora:([^:]+):")
@@ -57,8 +57,9 @@ def infotext_pasted(infotext, d): # pylint: disable=unused-argument
     d["Prompt"] = re.sub(re_lora, network_replacement, d["Prompt"])
 
 
-script_callbacks.on_app_started(api_networks)
-script_callbacks.on_before_ui(before_ui)
-script_callbacks.on_model_loaded(networks.assign_network_names_to_compvis_modules)
-script_callbacks.on_infotext_pasted(networks.infotext_pasted)
-script_callbacks.on_infotext_pasted(infotext_pasted)
+if not shared.native:
+    script_callbacks.on_app_started(api_networks)
+    script_callbacks.on_before_ui(before_ui)
+    script_callbacks.on_model_loaded(networks.assign_network_names_to_compvis_modules)
+    script_callbacks.on_infotext_pasted(networks.infotext_pasted)
+    script_callbacks.on_infotext_pasted(infotext_pasted)

installer.py

@@ -459,7 +459,7 @@ def check_python(supported_minors=[9, 10, 11, 12], reason=None):
 def check_diffusers():
     if args.skip_all or args.skip_requirements:
         return
-    sha = '069186fac510d6f6f88a5e435523b235c823a8a0'
+    sha = 'c96bfa5c80eca798d555a79a491043c311d0f608'
     pkg = pkg_resources.working_set.by_key.get('diffusers', None)
     minor = int(pkg.version.split('.')[1] if pkg is not None else 0)
     cur = opts.get('diffusers_version', '') if minor > 0 else ''

launch.py

@@ -192,6 +192,9 @@ def main():
     global args # pylint: disable=global-statement
     installer.ensure_base_requirements()
     init_args() # setup argparser and default folders
+    if args.malloc:
+        import tracemalloc
+        tracemalloc.start()
     installer.args = args
     installer.setup_logging()
     installer.log.info('Starting SD.Next')

modules/api/api.py

@@ -91,6 +91,11 @@ class Api:
         self.add_api_route("/sdapi/v1/history", endpoints.get_history, methods=["GET"], response_model=List[str])
         self.add_api_route("/sdapi/v1/history", endpoints.post_history, methods=["POST"], response_model=int)
 
+        # lora api
+        if shared.native:
+            self.add_api_route("/sdapi/v1/loras", endpoints.get_loras, methods=["GET"], response_model=List[dict])
+            self.add_api_route("/sdapi/v1/refresh-loras", endpoints.post_refresh_loras, methods=["POST"])
+
         # gallery api
         gallery.register_api(app)
 

modules/api/endpoints.py

@@ -40,6 +40,12 @@ def get_embeddings():
 
     return {"loaded": convert_embeddings(db.word_embeddings), "skipped": convert_embeddings(db.skipped_embeddings)}
 
+def get_loras():
+    from modules.lora import network, networks
+    def create_lora_json(obj: network.NetworkOnDisk):
+        return { "name": obj.name, "alias": obj.alias, "path": obj.filename, "metadata": obj.metadata }
+    return [create_lora_json(obj) for obj in networks.available_networks.values()]
+
 def get_extra_networks(page: Optional[str] = None, name: Optional[str] = None, filename: Optional[str] = None, title: Optional[str] = None, fullname: Optional[str] = None, hash: Optional[str] = None): # pylint: disable=redefined-builtin
     res = []
     for pg in shared.extra_networks:
@@ -126,6 +132,10 @@ def post_refresh_checkpoints():
 def post_refresh_vae():
     return shared.refresh_vaes()
 
+def post_refresh_loras():
+    from modules.lora import networks
+    return networks.list_available_networks()
+
 def get_extensions_list():
     from modules import extensions
     extensions.list_extensions()

modules/call_queue.py

@@ -73,16 +73,20 @@ def wrap_gradio_call(func, extra_outputs=None, add_stats=False, name=None):
         elapsed_m = int(elapsed // 60)
         elapsed_s = elapsed % 60
         elapsed_text = f"{elapsed_m}m {elapsed_s:.2f}s" if elapsed_m > 0 else f"{elapsed_s:.2f}s"
-        summary = timer.process.summary(min_time=0.1, total=False).replace('=', ' ')
-        vram_html = ''
+        summary = timer.process.summary(min_time=0.25, total=False).replace('=', ' ')
+        gpu = ''
+        cpu = ''
         if not shared.mem_mon.disabled:
             vram = {k: -(v//-(1024*1024)) for k, v in shared.mem_mon.read().items()}
-            used = round(100 * vram['used'] / (vram['total'] + 0.001))
-            if vram.get('active_peak', 0) > 0:
-                vram_html = " | "
-                vram_html += f"GPU {max(vram['active_peak'], vram['reserved_peak'])} MB {used}%"
-                vram_html += f" | retries {vram['retries']} oom {vram['oom']}" if vram.get('retries', 0) > 0 or vram.get('oom', 0) > 0 else ''
+            peak = max(vram['active_peak'], vram['reserved_peak'], vram['used'])
+            used = round(100.0 * peak / vram['total']) if vram['total'] > 0 else 0
+            if used > 0:
+                gpu += f"| GPU {peak} MB {used}%"
+                gpu += f" | retries {vram['retries']} oom {vram['oom']}" if vram.get('retries', 0) > 0 or vram.get('oom', 0) > 0 else ''
+            ram = shared.ram_stats()
+            if ram['used'] > 0:
+                cpu += f"| RAM {ram['used']} GB {round(100.0 * ram['used'] / ram['total'])}%"
         if isinstance(res, list):
-            res[-1] += f"<div class='performance'><p>Time: {elapsed_text} | {summary}{vram_html}</p></div>"
+            res[-1] += f"<div class='performance'><p>Time: {elapsed_text} | {summary} {gpu} {cpu}</p></div>"
         return tuple(res)
     return f

modules/cmd_args.py

@@ -26,6 +26,7 @@ def main_args():
     group_diag.add_argument("--no-hashing", default=os.environ.get("SD_NOHASHING", False), action='store_true', help="Disable hashing of checkpoints, default: %(default)s")
     group_diag.add_argument("--no-metadata", default=os.environ.get("SD_NOMETADATA", False), action='store_true', help="Disable reading of metadata from models, default: %(default)s")
     group_diag.add_argument("--profile", default=os.environ.get("SD_PROFILE", False), action='store_true', help="Run profiler, default: %(default)s")
+    group_diag.add_argument("--malloc", default=os.environ.get("SD_PROFILE", False), action='store_true', help="Trace memory ops, default: %(default)s")
     group_diag.add_argument("--disable-queue", default=os.environ.get("SD_DISABLEQUEUE", False), action='store_true', help="Disable queues, default: %(default)s")
     group_diag.add_argument('--debug', default=os.environ.get("SD_DEBUG", False), action='store_true', help = "Run installer with debug logging, default: %(default)s")
 

modules/devices.py

@@ -224,7 +224,7 @@ def torch_gc(force=False, fast=False):
         timer.process.records['gc'] = 0
     timer.process.records['gc'] += t1 - t0
     if not force or collected == 0:
-        return
+        return used_gpu, used_ram
     mem = memstats.memory_stats()
     saved = round(gpu.get('used', 0) - mem.get('gpu', {}).get('used', 0), 2)
     before = { 'gpu': gpu.get('used', 0), 'ram': ram.get('used', 0) }
@@ -233,6 +233,7 @@ def torch_gc(force=False, fast=False):
     results = { 'collected': collected, 'saved': saved }
     fn = f'{sys._getframe(2).f_code.co_name}:{sys._getframe(1).f_code.co_name}' # pylint: disable=protected-access
     log.debug(f'GC: utilization={utilization} gc={results} before={before} after={after} device={torch.device(get_optimal_device_name())} fn={fn} time={round(t1 - t0, 2)}') # pylint: disable=protected-access
+    return used_gpu, used_ram
 
 
 def set_cuda_sync_mode(mode):

modules/extensions.py

@@ -154,4 +154,4 @@ def list_extensions():
     for dirname, path, is_builtin in extension_paths:
         extension = Extension(name=dirname, path=path, enabled=dirname not in disabled_extensions, is_builtin=is_builtin)
         extensions.append(extension)
-    shared.log.info(f'Disabled extensions: {[e.name for e in extensions if not e.enabled]}')
+    shared.log.debug(f'Disabled extensions: {[e.name for e in extensions if not e.enabled]}')

modules/extra_networks.py

@@ -15,10 +15,14 @@ def register_extra_network(extra_network):
 
 
 def register_default_extra_networks():
-    from modules.ui_extra_networks_hypernet import ExtraNetworkHypernet
-    register_extra_network(ExtraNetworkHypernet())
     from modules.ui_extra_networks_styles import ExtraNetworkStyles
     register_extra_network(ExtraNetworkStyles())
+    if shared.native:
+        from modules.lora.networks import extra_network_lora
+        register_extra_network(extra_network_lora)
+    if shared.opts.hypernetwork_enabled:
+        from modules.ui_extra_networks_hypernet import ExtraNetworkHypernet
+        register_extra_network(ExtraNetworkHypernet())
 
 
 class ExtraNetworkParams:

modules/hashes.py

@@ -9,6 +9,13 @@ cache_filename = os.path.join(data_path, "cache.json")
 cache_data = None
 progress_ok = True
 
+
+def init_cache():
+    global cache_data # pylint: disable=global-statement
+    if cache_data is None:
+        cache_data = {} if not os.path.isfile(cache_filename) else shared.readfile(cache_filename, lock=True)
+
+
 def dump_cache():
     shared.writefile(cache_data, cache_filename)
 

modules/infotext.py

@@ -10,6 +10,7 @@ else:
     debug = lambda *args, **kwargs: None # pylint: disable=unnecessary-lambda-assignment
 re_size = re.compile(r"^(\d+)x(\d+)$") # int x int
 re_param = re.compile(r'\s*([\w ]+):\s*("(?:\\"[^,]|\\"|\\|[^\"])+"|[^,]*)(?:,|$)') # multi-word: value
+re_lora = re.compile("<lora:([^:]+):")
 
 
 def quote(text):
@@ -27,6 +28,40 @@ def unquote(text):
         return text
 
 
+# disabled by default can be enabled if needed
+def check_lora(params):
+    try:
+        import modules.lora.networks as networks
+        from modules.errors import log # pylint: disable=redefined-outer-name
+    except Exception:
+        return
+    loras = [s.strip() for s in params.get('LoRA hashes', '').split(',')]
+    found = []
+    missing = []
+    for l in loras:
+        lora = networks.available_network_hash_lookup.get(l, None)
+        if lora is not None:
+            found.append(lora.name)
+        else:
+            missing.append(l)
+    loras = [s.strip() for s in params.get('LoRA networks', '').split(',')]
+    for l in loras:
+        lora = networks.available_network_aliases.get(l, None)
+        if lora is not None:
+            found.append(lora.name)
+        else:
+            missing.append(l)
+    # networks.available_network_aliases.get(name, None)
+    loras = re_lora.findall(params.get('Prompt', ''))
+    for l in loras:
+        lora = networks.available_network_aliases.get(l, None)
+        if lora is not None:
+            found.append(lora.name)
+        else:
+            missing.append(l)
+    log.debug(f'LoRA: found={list(set(found))} missing={list(set(missing))}')
+
+
 def parse(infotext):
     if not isinstance(infotext, str):
         return {}
@@ -75,6 +110,7 @@ def parse(infotext):
             params[key] = val
         debug(f'Param parsed: type={type(params[key])} {key}={params[key]} raw="{val}"')
 
+    # check_lora(params)
     return params
 
 

modules/loader.py

@@ -14,7 +14,7 @@ errors.install()
 logging.getLogger("DeepSpeed").disabled = True
 
 
-os.environ.setdefault('TORCH_LOGS', '-all')
+# os.environ.setdefault('TORCH_LOGS', '-all')
 import torch # pylint: disable=C0411
 if torch.__version__.startswith('2.5.0'):
     errors.log.warning(f'Disabling cuDNN for SDP on torch={torch.__version__}')

modules/lora/extra_networks_lora.py

@@ -0,0 +1,152 @@
+import re
+import time
+import numpy as np
+import modules.lora.networks as networks
+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 len(steps[0]) == 1: # If we just got a single number, just return it
+            return steps[0][0]
+        steps = [[s[0], s[1] if len(s) == 2 else 1] for s in steps] # Add implicit 1s to any steps which don't have a weight
+        steps.sort(key=lambda k: k[1]) # Sort by index
+        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:
+                step = step / (max_steps - step_offset) if max_steps > 0 else 1.0
+            v = np.interp(step, m[1], m[0])
+            return v
+        else:
+            return m
+
+    stepwise = calculate_weight(sorted_positions(param), step, steps)
+    return stepwise
+
+
+def prompt(p):
+    if shared.opts.lora_apply_tags == 0:
+        return
+    all_tags = []
+    for loaded in networks.loaded_networks:
+        page = [en for en in shared.extra_networks if en.name == 'lora'][0]
+        item = page.create_item(loaded.name)
+        tags = (item or {}).get("tags", {})
+        loaded.tags = list(tags)
+        if len(loaded.tags) == 0:
+            loaded.tags.append(loaded.name)
+        if shared.opts.lora_apply_tags > 0:
+            loaded.tags = loaded.tags[:shared.opts.lora_apply_tags]
+        all_tags.extend(loaded.tags)
+    if len(all_tags) > 0:
+        shared.log.debug(f"Load network: type=LoRA tags={all_tags} max={shared.opts.lora_apply_tags} apply")
+        all_tags = ', '.join(all_tags)
+        p.extra_generation_params["LoRA tags"] = all_tags
+        if '_tags_' in p.prompt:
+            p.prompt = p.prompt.replace('_tags_', all_tags)
+        else:
+            p.prompt = f"{p.prompt}, {all_tags}"
+        if p.all_prompts is not None:
+            for i in range(len(p.all_prompts)):
+                if '_tags_' in p.all_prompts[i]:
+                    p.all_prompts[i] = p.all_prompts[i].replace('_tags_', all_tags)
+                else:
+                    p.all_prompts[i] = f"{p.all_prompts[i]}, {all_tags}"
+
+
+def infotext(p):
+    names = [i.name for i in networks.loaded_networks]
+    if len(names) > 0:
+        p.extra_generation_params["LoRA networks"] = ", ".join(names)
+    if shared.opts.lora_add_hashes_to_infotext:
+        network_hashes = []
+        for item in networks.loaded_networks:
+            if not item.network_on_disk.shorthash:
+                continue
+            network_hashes.append(item.network_on_disk.shorthash)
+        if len(network_hashes) > 0:
+            p.extra_generation_params["LoRA hashes"] = ", ".join(network_hashes)
+
+
+def parse(p, params_list, step=0):
+    names = []
+    te_multipliers = []
+    unet_multipliers = []
+    dyn_dims = []
+    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)
+        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)
+        unet_multipliers.append(unet_multiplier)
+        dyn_dims.append(dyn_dim)
+    return names, te_multipliers, unet_multipliers, dyn_dims
+
+
+class ExtraNetworkLora(extra_networks.ExtraNetwork):
+
+    def __init__(self):
+        super().__init__('lora')
+        self.active = False
+        self.model = None
+        self.errors = {}
+
+    def activate(self, p, params_list, step=0):
+        self.errors.clear()
+        if self.active:
+            if self.model != shared.opts.sd_model_checkpoint: # reset if model changed
+                self.active = False
+        if len(params_list) > 0 and not self.active: # activate patches once
+            # shared.log.debug(f'Activate network: type=LoRA model="{shared.opts.sd_model_checkpoint}"')
+            self.active = True
+            self.model = shared.opts.sd_model_checkpoint
+        names, te_multipliers, unet_multipliers, dyn_dims = parse(p, params_list, step)
+        networks.network_load(names, te_multipliers, unet_multipliers, dyn_dims) # load
+        networks.network_activate()
+        if len(networks.loaded_networks) > 0 and step == 0:
+            infotext(p)
+            prompt(p)
+            shared.log.info(f'Load network: type=LoRA apply={[n.name for n in networks.loaded_networks]} te={te_multipliers} unet={unet_multipliers} time={networks.get_timers()}')
+
+    def deactivate(self, p):
+        t0 = time.time()
+        if shared.native and len(networks.diffuser_loaded) > 0:
+            if hasattr(shared.sd_model, "unload_lora_weights") and hasattr(shared.sd_model, "text_encoder"):
+                if not (shared.compiled_model_state is not None and shared.compiled_model_state.is_compiled is True):
+                    try:
+                        if shared.opts.lora_fuse_diffusers:
+                            shared.sd_model.unfuse_lora()
+                        shared.sd_model.unload_lora_weights() # fails for non-CLIP models
+                    except Exception:
+                        pass
+        networks.network_deactivate()
+        t1 = time.time()
+        networks.timer['restore'] += t1 - t0
+        if self.active and networks.debug:
+            shared.log.debug(f"Network end: type=LoRA load={networks.timer['load']:.2f} apply={networks.timer['apply']:.2f} restore={networks.timer['restore']:.2f}")
+        if self.errors:
+            for k, v in self.errors.items():
+                shared.log.error(f'LoRA: name="{k}" errors={v}')
+            self.errors.clear()

modules/lora/lora_convert.py

@@ -0,0 +1,509 @@
+import os
+import re
+import bisect
+from typing import Dict
+import torch
+from modules import shared
+
+
+debug = os.environ.get('SD_LORA_DEBUG', None) is not None
+suffix_conversion = {
+    "attentions": {},
+    "resnets": {
+        "conv1": "in_layers_2",
+        "conv2": "out_layers_3",
+        "norm1": "in_layers_0",
+        "norm2": "out_layers_0",
+        "time_emb_proj": "emb_layers_1",
+        "conv_shortcut": "skip_connection",
+    }
+}
+re_digits = re.compile(r"\d+")
+re_x_proj = re.compile(r"(.*)_([qkv]_proj)$")
+re_compiled = {}
+
+
+def make_unet_conversion_map() -> Dict[str, str]:
+    unet_conversion_map_layer = []
+
+    for i in range(3):  # num_blocks is 3 in sdxl
+        # loop over downblocks/upblocks
+        for j in range(2):
+            # loop over resnets/attentions for downblocks
+            hf_down_res_prefix = f"down_blocks.{i}.resnets.{j}."
+            sd_down_res_prefix = f"input_blocks.{3 * i + j + 1}.0."
+            unet_conversion_map_layer.append((sd_down_res_prefix, hf_down_res_prefix))
+            if i < 3:
+                # no attention layers in down_blocks.3
+                hf_down_atn_prefix = f"down_blocks.{i}.attentions.{j}."
+                sd_down_atn_prefix = f"input_blocks.{3 * i + j + 1}.1."
+                unet_conversion_map_layer.append((sd_down_atn_prefix, hf_down_atn_prefix))
+
+        for j in range(3):
+            # loop over resnets/attentions for upblocks
+            hf_up_res_prefix = f"up_blocks.{i}.resnets.{j}."
+            sd_up_res_prefix = f"output_blocks.{3 * i + j}.0."
+            unet_conversion_map_layer.append((sd_up_res_prefix, hf_up_res_prefix))
+            # if i > 0: commentout for sdxl
+            # no attention layers in up_blocks.0
+            hf_up_atn_prefix = f"up_blocks.{i}.attentions.{j}."
+            sd_up_atn_prefix = f"output_blocks.{3 * i + j}.1."
+            unet_conversion_map_layer.append((sd_up_atn_prefix, hf_up_atn_prefix))
+
+        if i < 3:
+            # no downsample in down_blocks.3
+            hf_downsample_prefix = f"down_blocks.{i}.downsamplers.0.conv."
+            sd_downsample_prefix = f"input_blocks.{3 * (i + 1)}.0.op."
+            unet_conversion_map_layer.append((sd_downsample_prefix, hf_downsample_prefix))
+            # no upsample in up_blocks.3
+            hf_upsample_prefix = f"up_blocks.{i}.upsamplers.0."
+            sd_upsample_prefix = f"output_blocks.{3 * i + 2}.{2}."  # change for sdxl
+            unet_conversion_map_layer.append((sd_upsample_prefix, hf_upsample_prefix))
+
+    hf_mid_atn_prefix = "mid_block.attentions.0."
+    sd_mid_atn_prefix = "middle_block.1."
+    unet_conversion_map_layer.append((sd_mid_atn_prefix, hf_mid_atn_prefix))
+
+    for j in range(2):
+        hf_mid_res_prefix = f"mid_block.resnets.{j}."
+        sd_mid_res_prefix = f"middle_block.{2 * j}."
+        unet_conversion_map_layer.append((sd_mid_res_prefix, hf_mid_res_prefix))
+
+    unet_conversion_map_resnet = [
+        # (stable-diffusion, HF Diffusers)
+        ("in_layers.0.", "norm1."),
+        ("in_layers.2.", "conv1."),
+        ("out_layers.0.", "norm2."),
+        ("out_layers.3.", "conv2."),
+        ("emb_layers.1.", "time_emb_proj."),
+        ("skip_connection.", "conv_shortcut."),
+    ]
+
+    unet_conversion_map = []
+    for sd, hf in unet_conversion_map_layer:
+        if "resnets" in hf:
+            for sd_res, hf_res in unet_conversion_map_resnet:
+                unet_conversion_map.append((sd + sd_res, hf + hf_res))
+        else:
+            unet_conversion_map.append((sd, hf))
+
+    for j in range(2):
+        hf_time_embed_prefix = f"time_embedding.linear_{j + 1}."
+        sd_time_embed_prefix = f"time_embed.{j * 2}."
+        unet_conversion_map.append((sd_time_embed_prefix, hf_time_embed_prefix))
+
+    for j in range(2):
+        hf_label_embed_prefix = f"add_embedding.linear_{j + 1}."
+        sd_label_embed_prefix = f"label_emb.0.{j * 2}."
+        unet_conversion_map.append((sd_label_embed_prefix, hf_label_embed_prefix))
+
+    unet_conversion_map.append(("input_blocks.0.0.", "conv_in."))
+    unet_conversion_map.append(("out.0.", "conv_norm_out."))
+    unet_conversion_map.append(("out.2.", "conv_out."))
+
+    sd_hf_conversion_map = {sd.replace(".", "_")[:-1]: hf.replace(".", "_")[:-1] for sd, hf in unet_conversion_map}
+    return sd_hf_conversion_map
+
+
+class KeyConvert:
+    def __init__(self):
+        self.is_sdxl = True if shared.sd_model_type == "sdxl" else False
+        self.UNET_CONVERSION_MAP = make_unet_conversion_map() if self.is_sdxl else None
+        self.LORA_PREFIX_UNET = "lora_unet_"
+        self.LORA_PREFIX_TEXT_ENCODER = "lora_te_"
+        self.OFT_PREFIX_UNET = "oft_unet_"
+        # SDXL: must starts with LORA_PREFIX_TEXT_ENCODER
+        self.LORA_PREFIX_TEXT_ENCODER1 = "lora_te1_"
+        self.LORA_PREFIX_TEXT_ENCODER2 = "lora_te2_"
+
+    def __call__(self, key):
+        if self.is_sdxl:
+            if "diffusion_model" in key:  # Fix NTC Slider naming error
+                key = key.replace("diffusion_model", "lora_unet")
+            map_keys = list(self.UNET_CONVERSION_MAP.keys())  # prefix of U-Net modules
+            map_keys.sort()
+            search_key = key.replace(self.LORA_PREFIX_UNET, "").replace(self.OFT_PREFIX_UNET, "").replace(self.LORA_PREFIX_TEXT_ENCODER1, "").replace(self.LORA_PREFIX_TEXT_ENCODER2, "")
+            position = bisect.bisect_right(map_keys, search_key)
+            map_key = map_keys[position - 1]
+            if search_key.startswith(map_key):
+                key = key.replace(map_key, self.UNET_CONVERSION_MAP[map_key]).replace("oft", "lora") # pylint: disable=unsubscriptable-object
+        if "lycoris" in key and "transformer" in key:
+            key = key.replace("lycoris", "lora_transformer")
+        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 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
+
+
+# Taken from https://github.com/huggingface/diffusers/blob/main/src/diffusers/loaders/lora_conversion_utils.py
+# Modified from 'lora_A' and 'lora_B' to 'lora_down' and 'lora_up'
+# Added early exit
+# The utilities under `_convert_kohya_flux_lora_to_diffusers()`
+# are taken from https://github.com/kohya-ss/sd-scripts/blob/a61cf73a5cb5209c3f4d1a3688dd276a4dfd1ecb/networks/convert_flux_lora.py
+# All credits go to `kohya-ss`.
+def _convert_to_ai_toolkit(sds_sd, ait_sd, sds_key, ait_key):
+    if sds_key + ".lora_down.weight" not in sds_sd:
+        return
+    down_weight = sds_sd.pop(sds_key + ".lora_down.weight")
+
+    # scale weight by alpha and dim
+    rank = down_weight.shape[0]
+    alpha = sds_sd.pop(sds_key + ".alpha").item()  # alpha is scalar
+    scale = alpha / rank  # LoRA is scaled by 'alpha / rank' in forward pass, so we need to scale it back here
+
+    # calculate scale_down and scale_up to keep the same value. if scale is 4, scale_down is 2 and scale_up is 2
+    scale_down = scale
+    scale_up = 1.0
+    while scale_down * 2 < scale_up:
+        scale_down *= 2
+        scale_up /= 2
+
+    ait_sd[ait_key + ".lora_down.weight"] = down_weight * scale_down
+    ait_sd[ait_key + ".lora_up.weight"] = sds_sd.pop(sds_key + ".lora_up.weight") * scale_up
+
+def _convert_to_ai_toolkit_cat(sds_sd, ait_sd, sds_key, ait_keys, dims=None):
+    if sds_key + ".lora_down.weight" not in sds_sd:
+        return
+    down_weight = sds_sd.pop(sds_key + ".lora_down.weight")
+    up_weight = sds_sd.pop(sds_key + ".lora_up.weight")
+    sd_lora_rank = down_weight.shape[0]
+
+    # scale weight by alpha and dim
+    alpha = sds_sd.pop(sds_key + ".alpha")
+    scale = alpha / sd_lora_rank
+
+    # calculate scale_down and scale_up
+    scale_down = scale
+    scale_up = 1.0
+    while scale_down * 2 < scale_up:
+        scale_down *= 2
+        scale_up /= 2
+
+    down_weight = down_weight * scale_down
+    up_weight = up_weight * scale_up
+
+    # calculate dims if not provided
+    num_splits = len(ait_keys)
+    if dims is None:
+        dims = [up_weight.shape[0] // num_splits] * num_splits
+    else:
+        assert sum(dims) == up_weight.shape[0]
+
+    # check upweight is sparse or not
+    is_sparse = False
+    if sd_lora_rank % num_splits == 0:
+        ait_rank = sd_lora_rank // num_splits
+        is_sparse = True
+        i = 0
+        for j in range(len(dims)):
+            for k in range(len(dims)):
+                if j == k:
+                    continue
+                is_sparse = is_sparse and torch.all(
+                    up_weight[i : i + dims[j], k * ait_rank : (k + 1) * ait_rank] == 0
+                )
+            i += dims[j]
+        # if is_sparse:
+        #     print(f"weight is sparse: {sds_key}")
+
+    # make ai-toolkit weight
+    ait_down_keys = [k + ".lora_down.weight" for k in ait_keys]
+    ait_up_keys = [k + ".lora_up.weight" for k in ait_keys]
+    if not is_sparse:
+        # down_weight is copied to each split
+        ait_sd.update({k: down_weight for k in ait_down_keys})
+
+        # up_weight is split to each split
+        ait_sd.update({k: v for k, v in zip(ait_up_keys, torch.split(up_weight, dims, dim=0))})  # noqa: C416 # pylint: disable=unnecessary-comprehension
+    else:
+        # down_weight is chunked to each split
+        ait_sd.update({k: v for k, v in zip(ait_down_keys, torch.chunk(down_weight, num_splits, dim=0))})  # noqa: C416 # pylint: disable=unnecessary-comprehension
+
+        # up_weight is sparse: only non-zero values are copied to each split
+        i = 0
+        for j in range(len(dims)):
+            ait_sd[ait_up_keys[j]] = up_weight[i : i + dims[j], j * ait_rank : (j + 1) * ait_rank].contiguous()
+            i += dims[j]
+
+def _convert_text_encoder_lora_key(key, lora_name):
+    """
+    Converts a text encoder LoRA key to a Diffusers compatible key.
+    """
+    if lora_name.startswith(("lora_te_", "lora_te1_")):
+        key_to_replace = "lora_te_" if lora_name.startswith("lora_te_") else "lora_te1_"
+    else:
+        key_to_replace = "lora_te2_"
+
+    diffusers_name = key.replace(key_to_replace, "").replace("_", ".")
+    diffusers_name = diffusers_name.replace("text.model", "text_model")
+    diffusers_name = diffusers_name.replace("self.attn", "self_attn")
+    diffusers_name = diffusers_name.replace("q.proj.lora", "to_q_lora")
+    diffusers_name = diffusers_name.replace("k.proj.lora", "to_k_lora")
+    diffusers_name = diffusers_name.replace("v.proj.lora", "to_v_lora")
+    diffusers_name = diffusers_name.replace("out.proj.lora", "to_out_lora")
+    diffusers_name = diffusers_name.replace("text.projection", "text_projection")
+
+    if "self_attn" in diffusers_name or "text_projection" in diffusers_name:
+        pass
+    elif "mlp" in diffusers_name:
+        # Be aware that this is the new diffusers convention and the rest of the code might
+        # not utilize it yet.
+        diffusers_name = diffusers_name.replace(".lora.", ".lora_linear_layer.")
+    return diffusers_name
+
+def _convert_kohya_flux_lora_to_diffusers(state_dict):
+    def _convert_sd_scripts_to_ai_toolkit(sds_sd):
+        ait_sd = {}
+        for i in range(19):
+            _convert_to_ai_toolkit(
+                sds_sd,
+                ait_sd,
+                f"lora_unet_double_blocks_{i}_img_attn_proj",
+                f"transformer.transformer_blocks.{i}.attn.to_out.0",
+            )
+            _convert_to_ai_toolkit_cat(
+                sds_sd,
+                ait_sd,
+                f"lora_unet_double_blocks_{i}_img_attn_qkv",
+                [
+                    f"transformer.transformer_blocks.{i}.attn.to_q",
+                    f"transformer.transformer_blocks.{i}.attn.to_k",
+                    f"transformer.transformer_blocks.{i}.attn.to_v",
+                ],
+            )
+            _convert_to_ai_toolkit(
+                sds_sd,
+                ait_sd,
+                f"lora_unet_double_blocks_{i}_img_mlp_0",
+                f"transformer.transformer_blocks.{i}.ff.net.0.proj",
+            )
+            _convert_to_ai_toolkit(
+                sds_sd,
+                ait_sd,
+                f"lora_unet_double_blocks_{i}_img_mlp_2",
+                f"transformer.transformer_blocks.{i}.ff.net.2",
+            )
+            _convert_to_ai_toolkit(
+                sds_sd,
+                ait_sd,
+                f"lora_unet_double_blocks_{i}_img_mod_lin",
+                f"transformer.transformer_blocks.{i}.norm1.linear",
+            )
+            _convert_to_ai_toolkit(
+                sds_sd,
+                ait_sd,
+                f"lora_unet_double_blocks_{i}_txt_attn_proj",
+                f"transformer.transformer_blocks.{i}.attn.to_add_out",
+            )
+            _convert_to_ai_toolkit_cat(
+                sds_sd,
+                ait_sd,
+                f"lora_unet_double_blocks_{i}_txt_attn_qkv",
+                [
+                    f"transformer.transformer_blocks.{i}.attn.add_q_proj",
+                    f"transformer.transformer_blocks.{i}.attn.add_k_proj",
+                    f"transformer.transformer_blocks.{i}.attn.add_v_proj",
+                ],
+            )
+            _convert_to_ai_toolkit(
+                sds_sd,
+                ait_sd,
+                f"lora_unet_double_blocks_{i}_txt_mlp_0",
+                f"transformer.transformer_blocks.{i}.ff_context.net.0.proj",
+            )
+            _convert_to_ai_toolkit(
+                sds_sd,
+                ait_sd,
+                f"lora_unet_double_blocks_{i}_txt_mlp_2",
+                f"transformer.transformer_blocks.{i}.ff_context.net.2",
+            )
+            _convert_to_ai_toolkit(
+                sds_sd,
+                ait_sd,
+                f"lora_unet_double_blocks_{i}_txt_mod_lin",
+                f"transformer.transformer_blocks.{i}.norm1_context.linear",
+            )
+
+        for i in range(38):
+            _convert_to_ai_toolkit_cat(
+                sds_sd,
+                ait_sd,
+                f"lora_unet_single_blocks_{i}_linear1",
+                [
+                    f"transformer.single_transformer_blocks.{i}.attn.to_q",
+                    f"transformer.single_transformer_blocks.{i}.attn.to_k",
+                    f"transformer.single_transformer_blocks.{i}.attn.to_v",
+                    f"transformer.single_transformer_blocks.{i}.proj_mlp",
+                ],
+                dims=[3072, 3072, 3072, 12288],
+            )
+            _convert_to_ai_toolkit(
+                sds_sd,
+                ait_sd,
+                f"lora_unet_single_blocks_{i}_linear2",
+                f"transformer.single_transformer_blocks.{i}.proj_out",
+            )
+            _convert_to_ai_toolkit(
+                sds_sd,
+                ait_sd,
+                f"lora_unet_single_blocks_{i}_modulation_lin",
+                f"transformer.single_transformer_blocks.{i}.norm.linear",
+            )
+
+        if len(sds_sd) > 0:
+            return None
+
+        return ait_sd
+
+    return _convert_sd_scripts_to_ai_toolkit(state_dict)
+
+def _convert_kohya_sd3_lora_to_diffusers(state_dict):
+    def _convert_sd_scripts_to_ai_toolkit(sds_sd):
+        ait_sd = {}
+        for i in range(38):
+            _convert_to_ai_toolkit_cat(
+                sds_sd,
+                ait_sd,
+                f"lora_unet_joint_blocks_{i}_context_block_attn_qkv",
+                [
+                    f"transformer.transformer_blocks.{i}.attn.to_q",
+                    f"transformer.transformer_blocks.{i}.attn.to_k",
+                    f"transformer.transformer_blocks.{i}.attn.to_v",
+                ],
+            )
+            _convert_to_ai_toolkit(
+                sds_sd,
+                ait_sd,
+                f"lora_unet_joint_blocks_{i}_context_block_mlp_fc1",
+                f"transformer.transformer_blocks.{i}.ff_context.net.0.proj",
+            )
+            _convert_to_ai_toolkit(
+                sds_sd,
+                ait_sd,
+                f"lora_unet_joint_blocks_{i}_context_block_mlp_fc2",
+                f"transformer.transformer_blocks.{i}.ff_context.net.2",
+            )
+            _convert_to_ai_toolkit(
+                sds_sd,
+                ait_sd,
+                f"lora_unet_joint_blocks_{i}_x_block_mlp_fc1",
+                f"transformer.transformer_blocks.{i}.ff.net.0.proj",
+            )
+            _convert_to_ai_toolkit(
+                sds_sd,
+                ait_sd,
+                f"lora_unet_joint_blocks_{i}_x_block_mlp_fc2",
+                f"transformer.transformer_blocks.{i}.ff.net.2",
+            )
+            _convert_to_ai_toolkit(
+                sds_sd,
+                ait_sd,
+                f"lora_unet_joint_blocks_{i}_context_block_adaLN_modulation_1",
+                f"transformer.transformer_blocks.{i}.norm1_context.linear",
+            )
+            _convert_to_ai_toolkit(
+                sds_sd,
+                ait_sd,
+                f"lora_unet_joint_blocks_{i}_x_block_adaLN_modulation_1",
+                f"transformer.transformer_blocks.{i}.norm1.linear",
+            )
+            _convert_to_ai_toolkit(
+                sds_sd,
+                ait_sd,
+                f"lora_unet_joint_blocks_{i}_context_block_attn_proj",
+                f"transformer.transformer_blocks.{i}.attn.to_add_out",
+            )
+            _convert_to_ai_toolkit(
+                sds_sd,
+                ait_sd,
+                f"lora_unet_joint_blocks_{i}_x_block_attn_proj",
+                f"transformer.transformer_blocks.{i}.attn.to_out_0",
+            )
+
+            _convert_to_ai_toolkit_cat(
+                sds_sd,
+                ait_sd,
+                f"lora_unet_joint_blocks_{i}_x_block_attn_qkv",
+                [
+                    f"transformer.transformer_blocks.{i}.attn.add_q_proj",
+                    f"transformer.transformer_blocks.{i}.attn.add_k_proj",
+                    f"transformer.transformer_blocks.{i}.attn.add_v_proj",
+                ],
+            )
+        remaining_keys = list(sds_sd.keys())
+        te_state_dict = {}
+        if remaining_keys:
+            if not all(k.startswith("lora_te1") for k in remaining_keys):
+                raise ValueError(f"Incompatible keys detected: \n\n {', '.join(remaining_keys)}")
+            for key in remaining_keys:
+                if not key.endswith("lora_down.weight"):
+                    continue
+
+                lora_name = key.split(".")[0]
+                lora_name_up = f"{lora_name}.lora_up.weight"
+                lora_name_alpha = f"{lora_name}.alpha"
+                diffusers_name = _convert_text_encoder_lora_key(key, lora_name)
+
+                sd_lora_rank = 1
+                if lora_name.startswith(("lora_te_", "lora_te1_")):
+                    down_weight = sds_sd.pop(key)
+                    sd_lora_rank = down_weight.shape[0]
+                    te_state_dict[diffusers_name] = down_weight
+                    te_state_dict[diffusers_name.replace(".down.", ".up.")] = sds_sd.pop(lora_name_up)
+
+                if lora_name_alpha in sds_sd:
+                    alpha = sds_sd.pop(lora_name_alpha).item()
+                    scale = alpha / sd_lora_rank
+
+                    scale_down = scale
+                    scale_up = 1.0
+                    while scale_down * 2 < scale_up:
+                        scale_down *= 2
+                        scale_up /= 2
+
+                    te_state_dict[diffusers_name] *= scale_down
+                    te_state_dict[diffusers_name.replace(".down.", ".up.")] *= scale_up
+
+        if len(sds_sd) > 0:
+            print(f"Unsupported keys for ai-toolkit: {sds_sd.keys()}")
+
+        if te_state_dict:
+            te_state_dict = {f"text_encoder.{module_name}": params for module_name, params in te_state_dict.items()}
+
+        new_state_dict = {**ait_sd, **te_state_dict}
+        return new_state_dict
+
+    return _convert_sd_scripts_to_ai_toolkit(state_dict)
+
+
+def assign_network_names_to_compvis_modules(sd_model):
+    if sd_model is None:
+        return
+    sd_model = getattr(shared.sd_model, "pipe", shared.sd_model)  # wrapped model compatiblility
+    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():
+            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'):
+        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'):
+        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'):
+        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
+    shared.sd_model.network_layer_mapping = network_layer_mapping

modules/lora/lora_extract.py

@@ -0,0 +1,271 @@
+import os
+import time
+import json
+import datetime
+import torch
+from safetensors.torch import save_file
+import gradio as gr
+from rich import progress as p
+from modules import shared, devices
+from modules.ui_common import create_refresh_button
+from modules.call_queue import wrap_gradio_gpu_call
+
+
+class SVDHandler:
+    def __init__(self, maxrank=0, rank_ratio=1):
+        self.network_name: str = None
+        self.U: torch.Tensor = None
+        self.S: torch.Tensor = None
+        self.Vh: torch.Tensor = None
+        self.maxrank: int = maxrank
+        self.rank_ratio: float = rank_ratio
+        self.rank: int = 0
+        self.out_size: int = None
+        self.in_size: int = None
+        self.kernel_size: tuple[int, int] = None
+        self.conv2d: bool = False
+
+    def decompose(self, weight, backupweight):
+        self.conv2d = len(weight.size()) == 4
+        self.kernel_size = None if not self.conv2d else weight.size()[2:4]
+        self.out_size, self.in_size = weight.size()[0:2]
+        diffweight = weight.clone().to(devices.device)
+        diffweight -= backupweight.to(devices.device)
+        if self.conv2d:
+            if self.conv2d and self.kernel_size != (1, 1):
+                diffweight = diffweight.flatten(start_dim=1)
+            else:
+                diffweight = diffweight.squeeze()
+        self.U, self.S, self.Vh = torch.svd_lowrank(diffweight.to(device=devices.device, dtype=torch.float), self.maxrank, 2)
+        # del diffweight
+        self.U = self.U.to(device=devices.cpu, dtype=torch.bfloat16)
+        self.S = self.S.to(device=devices.cpu, dtype=torch.bfloat16)
+        self.Vh = self.Vh.t().to(device=devices.cpu, dtype=torch.bfloat16)  # svd_lowrank outputs a transposed matrix
+
+    def findrank(self):
+        if self.rank_ratio < 1:
+            S_squared = self.S.pow(2)
+            S_fro_sq = float(torch.sum(S_squared))
+            sum_S_squared = torch.cumsum(S_squared, dim=0) / S_fro_sq
+            index = int(torch.searchsorted(sum_S_squared, self.rank_ratio ** 2)) + 1
+            index = max(1, min(index, len(self.S) - 1))
+            self.rank = index
+            if self.maxrank > 0:
+                self.rank = min(self.rank, self.maxrank)
+        else:
+            self.rank = min(self.in_size, self.out_size, self.maxrank)
+
+    def makeweights(self):
+        self.findrank()
+        up = self.U[:, :self.rank] @ torch.diag(self.S[:self.rank])
+        down = self.Vh[:self.rank, :]
+        if self.conv2d and self.kernel_size is not None:
+            up = up.reshape(self.out_size, self.rank, 1, 1)
+            down = down.reshape(self.rank, self.in_size, self.kernel_size[0], self.kernel_size[1]) # pylint: disable=unsubscriptable-object
+        return_dict = {f'{self.network_name}.lora_up.weight': up.contiguous(),
+                       f'{self.network_name}.lora_down.weight': down.contiguous(),
+                       f'{self.network_name}.alpha': torch.tensor(down.shape[0]),
+                       }
+        return return_dict
+
+
+def loaded_lora():
+    if not shared.sd_loaded:
+        return ""
+    loaded = set()
+    if hasattr(shared.sd_model, 'unet'):
+        for _name, module in shared.sd_model.unet.named_modules():
+            current = getattr(module, "network_current_names", None)
+            if current is not None:
+                current = [item[0] for item in current]
+                loaded.update(current)
+    return list(loaded)
+
+
+def loaded_lora_str():
+    return ", ".join(loaded_lora())
+
+
+def make_meta(fn, maxrank, rank_ratio):
+    meta = {
+        "model_spec.sai_model_spec": "1.0.0",
+        "model_spec.title": os.path.splitext(os.path.basename(fn))[0],
+        "model_spec.author": "SD.Next",
+        "model_spec.implementation": "https://github.com/vladmandic/automatic",
+        "model_spec.date": datetime.datetime.now().astimezone().replace(microsecond=0).isoformat(),
+        "model_spec.base_model": shared.opts.sd_model_checkpoint,
+        "model_spec.dtype": str(devices.dtype),
+        "model_spec.base_lora": json.dumps(loaded_lora()),
+        "model_spec.config": f"maxrank={maxrank} rank_ratio={rank_ratio}",
+    }
+    if shared.sd_model_type == "sdxl":
+        meta["model_spec.architecture"] = "stable-diffusion-xl-v1-base/lora" # sai standard
+        meta["ss_base_model_version"] = "sdxl_base_v1-0" # kohya standard
+    elif shared.sd_model_type == "sd":
+        meta["model_spec.architecture"] = "stable-diffusion-v1/lora"
+        meta["ss_base_model_version"] = "sd_v1"
+    elif shared.sd_model_type == "f1":
+        meta["model_spec.architecture"] = "flux-1-dev/lora"
+        meta["ss_base_model_version"] = "flux1"
+    elif shared.sd_model_type == "sc":
+        meta["model_spec.architecture"] = "stable-cascade-v1-prior/lora"
+    return meta
+
+
+def make_lora(fn, maxrank, auto_rank, rank_ratio, modules, overwrite):
+    if not shared.sd_loaded or not shared.native:
+        msg = "LoRA extract: model not loaded"
+        shared.log.warning(msg)
+        yield msg
+        return
+    if loaded_lora() == "":
+        msg = "LoRA extract: no LoRA detected"
+        shared.log.warning(msg)
+        yield msg
+        return
+    if not fn:
+        msg = "LoRA extract: target filename required"
+        shared.log.warning(msg)
+        yield msg
+        return
+    t0 = time.time()
+    maxrank = int(maxrank)
+    rank_ratio = 1 if not auto_rank else rank_ratio
+    shared.log.debug(f'LoRA extract: modules={modules} maxrank={maxrank} auto={auto_rank} ratio={rank_ratio} fn="{fn}"')
+    shared.state.begin('LoRA extract')
+
+    with p.Progress(p.TextColumn('[cyan]LoRA extract'), p.BarColumn(), p.TaskProgressColumn(), p.TimeRemainingColumn(), p.TimeElapsedColumn(), p.TextColumn('[cyan]{task.description}'), console=shared.console) as progress:
+
+        if 'te' in modules and getattr(shared.sd_model, 'text_encoder', None) is not None:
+            modules = shared.sd_model.text_encoder.named_modules()
+            task = progress.add_task(description="te1 decompose", total=len(list(modules)))
+            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_"
+                module.svdhandler = SVDHandler(maxrank, rank_ratio)
+                module.svdhandler.network_name = prefix + name.replace(".", "_")
+                with devices.inference_context():
+                    module.svdhandler.decompose(module.weight, weights_backup)
+            progress.remove_task(task)
+        t1 = time.time()
+
+        if 'te' in modules and getattr(shared.sd_model, 'text_encoder_2', None) is not None:
+            modules = shared.sd_model.text_encoder_2.named_modules()
+            task = progress.add_task(description="te2 decompose", total=len(list(modules)))
+            for name, module in shared.sd_model.text_encoder_2.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
+                module.svdhandler = SVDHandler(maxrank, rank_ratio)
+                module.svdhandler.network_name = "lora_te2_" + name.replace(".", "_")
+                with devices.inference_context():
+                    module.svdhandler.decompose(module.weight, weights_backup)
+            progress.remove_task(task)
+        t2 = time.time()
+
+        if 'unet' in modules and getattr(shared.sd_model, 'unet', None) is not None:
+            modules = shared.sd_model.unet.named_modules()
+            task = progress.add_task(description="unet decompose", total=len(list(modules)))
+            for name, module in shared.sd_model.unet.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
+                module.svdhandler = SVDHandler(maxrank, rank_ratio)
+                module.svdhandler.network_name = "lora_unet_" + name.replace(".", "_")
+                with devices.inference_context():
+                    module.svdhandler.decompose(module.weight, weights_backup)
+            progress.remove_task(task)
+        t3 = time.time()
+
+        # TODO: Handle quant for Flux
+        # if 'te' in modules and getattr(shared.sd_model, 'transformer', None) is not None:
+        #     for name, module in shared.sd_model.transformer.named_modules():
+        #         if "norm" in name and "linear" not in name:
+        #             continue
+        #         weights_backup = getattr(module, "network_weights_backup", None)
+        #         if weights_backup is None:
+        #             continue
+        #         module.svdhandler = SVDHandler()
+        #         module.svdhandler.network_name = "lora_transformer_" + name.replace(".", "_")
+        #         module.svdhandler.decompose(module.weight, weights_backup)
+        #         module.svdhandler.findrank(rank, rank_ratio)
+
+        lora_state_dict = {}
+        for sub in ['text_encoder', 'text_encoder_2', 'unet', 'transformer']:
+            submodel = getattr(shared.sd_model, sub, None)
+            if submodel is not None:
+                modules = submodel.named_modules()
+                task = progress.add_task(description=f"{sub} exctract", total=len(list(modules)))
+                for _name, module in submodel.named_modules():
+                    progress.update(task, advance=1)
+                    if not hasattr(module, "svdhandler"):
+                        continue
+                    lora_state_dict.update(module.svdhandler.makeweights())
+                    del module.svdhandler
+                progress.remove_task(task)
+        t4 = time.time()
+
+    if not os.path.isabs(fn):
+        fn = os.path.join(shared.cmd_opts.lora_dir, fn)
+    if not fn.endswith('.safetensors'):
+        fn += '.safetensors'
+    if os.path.exists(fn):
+        if overwrite:
+            os.remove(fn)
+        else:
+            msg = f'LoRA extract: fn="{fn}" file exists'
+            shared.log.warning(msg)
+            yield msg
+            return
+
+    shared.state.end()
+    meta = make_meta(fn, maxrank, rank_ratio)
+    shared.log.debug(f'LoRA metadata: {meta}')
+    try:
+        save_file(tensors=lora_state_dict, metadata=meta, filename=fn)
+    except Exception as e:
+        msg = f'LoRA extract error: fn="{fn}" {e}'
+        shared.log.error(msg)
+        yield msg
+        return
+    t5 = time.time()
+    shared.log.debug(f'LoRA extract: time={t5-t0:.2f} te1={t1-t0:.2f} te2={t2-t1:.2f} unet={t3-t2:.2f} save={t5-t4:.2f}')
+    keys = list(lora_state_dict.keys())
+    msg = f'LoRA extract: fn="{fn}" keys={len(keys)}'
+    shared.log.info(msg)
+    yield msg
+
+
+def create_ui():
+    def gr_show(visible=True):
+        return {"visible": visible, "__type__": "update"}
+
+    with gr.Tab(label="Extract LoRA"):
+        with gr.Row():
+            loaded = gr.Textbox(placeholder="Press refresh to query loaded LoRA", label="Loaded LoRA", interactive=False)
+            create_refresh_button(loaded, lambda: None, lambda: {'value': loaded_lora_str()}, "testid")
+        with gr.Group():
+            with gr.Row():
+                modules = gr.CheckboxGroup(label="Modules to extract", value=['unet'], choices=['te', 'unet'])
+            with gr.Row():
+                auto_rank = gr.Checkbox(value=False, label="Automatically determine rank")
+                rank_ratio = gr.Slider(label="Autorank ratio", value=1, minimum=0, maximum=1, step=0.05, visible=False)
+                rank = gr.Slider(label="Maximum rank", value=32, minimum=1, maximum=256)
+        with gr.Row():
+            filename = gr.Textbox(label="LoRA target filename")
+            overwrite = gr.Checkbox(value=False, label="Overwrite existing file")
+        with gr.Row():
+            extract = gr.Button(value="Extract LoRA", variant='primary')
+            status = gr.HTML(value="", show_label=False)
+
+        auto_rank.change(fn=lambda x: gr_show(x), inputs=[auto_rank], outputs=[rank_ratio])
+        extract.click(
+            fn=wrap_gradio_gpu_call(make_lora, extra_outputs=[]),
+            inputs=[filename, rank, auto_rank, rank_ratio, modules, overwrite],
+            outputs=[status]
+        )

modules/lora/lyco_helpers.py

@@ -0,0 +1,66 @@
+import torch
+
+
+def make_weight_cp(t, wa, wb):
+    temp = torch.einsum('i j k l, j r -> i r k l', t, wb)
+    return torch.einsum('i j k l, i r -> r j k l', temp, wa)
+
+
+def rebuild_conventional(up, down, shape, dyn_dim=None):
+    up = up.reshape(up.size(0), -1)
+    down = down.reshape(down.size(0), -1)
+    if dyn_dim is not None:
+        up = up[:, :dyn_dim]
+        down = down[:dyn_dim, :]
+    return (up @ down).reshape(shape).to(up.dtype)
+
+
+def rebuild_cp_decomposition(up, down, mid):
+    up = up.reshape(up.size(0), -1)
+    down = down.reshape(down.size(0), -1)
+    return torch.einsum('n m k l, i n, m j -> i j k l', mid, up, down).to(up.dtype)
+
+
+# copied from https://github.com/KohakuBlueleaf/LyCORIS/blob/dev/lycoris/modules/lokr.py
+def factorization(dimension: int, factor:int=-1) -> tuple[int, int]:
+    """
+    return a tuple of two value of input dimension decomposed by the number closest to factor
+    second value is higher or equal than first value.
+
+    In LoRA with Kroneckor Product, first value is a value for weight scale.
+    secon value is a value for weight.
+
+    Becuase of non-commutative property, A⊗B ≠ B⊗A. Meaning of two matrices is slightly different.
+
+    examples
+    factor
+        -1               2                4               8               16               ...
+    127 -> 1, 127   127 -> 1, 127    127 -> 1, 127   127 -> 1, 127   127 -> 1, 127
+    128 -> 8, 16    128 -> 2, 64     128 -> 4, 32    128 -> 8, 16    128 -> 8, 16
+    250 -> 10, 25   250 -> 2, 125    250 -> 2, 125   250 -> 5, 50    250 -> 10, 25
+    360 -> 8, 45    360 -> 2, 180    360 -> 4, 90    360 -> 8, 45    360 -> 12, 30
+    512 -> 16, 32   512 -> 2, 256    512 -> 4, 128   512 -> 8, 64    512 -> 16, 32
+    1024 -> 32, 32  1024 -> 2, 512   1024 -> 4, 256  1024 -> 8, 128  1024 -> 16, 64
+    """
+
+    if factor > 0 and (dimension % factor) == 0:
+        m = factor
+        n = dimension // factor
+        if m > n:
+            n, m = m, n
+        return m, n
+    if factor < 0:
+        factor = dimension
+    m, n = 1, dimension
+    length = m + n
+    while m<n:
+        new_m = m + 1
+        while dimension%new_m != 0:
+            new_m += 1
+        new_n = dimension // new_m
+        if new_m + new_n > length or new_m>factor:
+            break
+        m, n = new_m, new_n
+    if m > n:
+        n, m = m, n
+    return m, n

modules/lora/network.py

@@ -0,0 +1,189 @@
+import os
+import enum
+from typing import Union
+from collections import namedtuple
+
+from modules import sd_models, hashes, shared
+
+
+NetworkWeights = namedtuple('NetworkWeights', ['network_key', 'sd_key', 'w', 'sd_module'])
+metadata_tags_order = {"ss_sd_model_name": 1, "ss_resolution": 2, "ss_clip_skip": 3, "ss_num_train_images": 10, "ss_tag_frequency": 20}
+
+
+class SdVersion(enum.Enum):
+    Unknown = 1
+    SD1 = 2
+    SD2 = 3
+    SD3 = 3
+    SDXL = 4
+    SC = 5
+    F1 = 6
+
+
+class NetworkOnDisk:
+    def __init__(self, name, filename):
+        self.shorthash = None
+        self.hash = None
+        self.name = name
+        self.filename = filename
+        if filename.startswith(shared.cmd_opts.lora_dir):
+            self.fullname = os.path.splitext(filename[len(shared.cmd_opts.lora_dir):].strip("/"))[0]
+        else:
+            self.fullname = name
+        self.metadata = {}
+        self.is_safetensors = os.path.splitext(filename)[1].lower() == ".safetensors"
+        if self.is_safetensors:
+            self.metadata = sd_models.read_metadata_from_safetensors(filename)
+        if self.metadata:
+            m = {}
+            for k, v in sorted(self.metadata.items(), key=lambda x: metadata_tags_order.get(x[0], 999)):
+                m[k] = v
+            self.metadata = m
+        self.alias = self.metadata.get('ss_output_name', self.name)
+        sha256 = hashes.sha256_from_cache(self.filename, "lora/" + self.name) or hashes.sha256_from_cache(self.filename, "lora/" + self.name, use_addnet_hash=True) or self.metadata.get('sshs_model_hash')
+        self.set_hash(sha256)
+        self.sd_version = self.detect_version()
+
+    def detect_version(self):
+        base = str(self.metadata.get('ss_base_model_version', "")).lower()
+        arch = str(self.metadata.get('modelspec.architecture', "")).lower()
+        if base.startswith("sd_v1"):
+            return 'sd1'
+        if base.startswith("sdxl"):
+            return 'xl'
+        if base.startswith("stable_cascade"):
+            return 'sc'
+        if base.startswith("sd3"):
+            return 'sd3'
+        if base.startswith("flux"):
+            return 'f1'
+
+        if arch.startswith("stable-diffusion-v1"):
+            return 'sd1'
+        if arch.startswith("stable-diffusion-xl"):
+            return 'xl'
+        if arch.startswith("stable-cascade"):
+            return 'sc'
+        if arch.startswith("flux"):
+            return 'f1'
+
+        if "v1-5" in str(self.metadata.get('ss_sd_model_name', "")):
+            return 'sd1'
+        if str(self.metadata.get('ss_v2', "")) == "True":
+            return 'sd2'
+        if 'flux' in self.name.lower():
+            return 'f1'
+        if 'xl' in self.name.lower():
+            return 'xl'
+
+        return ''
+
+    def set_hash(self, v):
+        self.hash = v or ''
+        self.shorthash = self.hash[0:8]
+
+    def read_hash(self):
+        if not self.hash:
+            self.set_hash(hashes.sha256(self.filename, "lora/" + self.name, use_addnet_hash=self.is_safetensors) or '')
+
+    def get_alias(self):
+        import modules.lora.networks as networks
+        return self.name if shared.opts.lora_preferred_name == "filename" or self.alias.lower() in networks.forbidden_network_aliases else self.alias
+
+
+class Network:  # LoraModule
+    def __init__(self, name, network_on_disk: NetworkOnDisk):
+        self.name = name
+        self.network_on_disk = network_on_disk
+        self.te_multiplier = 1.0
+        self.unet_multiplier = [1.0] * 3
+        self.dyn_dim = None
+        self.modules = {}
+        self.bundle_embeddings = {}
+        self.mtime = None
+        self.mentioned_name = None
+        self.tags = None
+        """the text that was used to add the network to prompt - can be either name or an alias"""
+
+
+class ModuleType:
+    def create_module(self, net: Network, weights: NetworkWeights) -> Union[Network, None]: # pylint: disable=W0613
+        return None
+
+
+class NetworkModule:
+    def __init__(self, net: Network, weights: NetworkWeights):
+        self.network = net
+        self.network_key = weights.network_key
+        self.sd_key = weights.sd_key
+        self.sd_module = weights.sd_module
+        if hasattr(self.sd_module, 'weight'):
+            self.shape = self.sd_module.weight.shape
+        self.dim = None
+        self.bias = weights.w.get("bias")
+        self.alpha = weights.w["alpha"].item() if "alpha" in weights.w else None
+        self.scale = weights.w["scale"].item() if "scale" in weights.w else None
+        self.dora_scale = weights.w.get("dora_scale", None)
+        self.dora_norm_dims = len(self.shape) - 1
+
+    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]:
+            return self.network.te_multiplier
+        if "down_blocks" in self.sd_key:
+            return unet_multiplier[0]
+        if "mid_block" in self.sd_key:
+            return unet_multiplier[1]
+        if "up_blocks" in self.sd_key:
+            return unet_multiplier[2]
+        else:
+            return unet_multiplier[0]
+
+    def calc_scale(self):
+        if self.scale is not None:
+            return self.scale
+        if self.dim is not None and self.alpha is not None:
+            return self.alpha / self.dim
+        return 1.0
+
+    def apply_weight_decompose(self, updown, orig_weight):
+        # Match the device/dtype
+        orig_weight = orig_weight.to(updown.dtype)
+        dora_scale = self.dora_scale.to(device=orig_weight.device, dtype=updown.dtype)
+        updown = updown.to(orig_weight.device)
+
+        merged_scale1 = updown + orig_weight
+        merged_scale1_norm = (
+            merged_scale1.transpose(0, 1)
+            .reshape(merged_scale1.shape[1], -1)
+            .norm(dim=1, keepdim=True)
+            .reshape(merged_scale1.shape[1], *[1] * self.dora_norm_dims)
+            .transpose(0, 1)
+        )
+
+        dora_merged = (
+                merged_scale1 * (dora_scale / merged_scale1_norm)
+        )
+        final_updown = dora_merged - orig_weight
+        return final_updown
+
+    def finalize_updown(self, updown, orig_weight, output_shape, ex_bias=None):
+        if self.bias is not None:
+            updown = updown.reshape(self.bias.shape)
+            updown += self.bias.to(orig_weight.device, dtype=orig_weight.dtype)
+            updown = updown.reshape(output_shape)
+        if len(output_shape) == 4:
+            updown = updown.reshape(output_shape)
+        if orig_weight.size().numel() == updown.size().numel():
+            updown = updown.reshape(orig_weight.shape)
+        if ex_bias is not None:
+            ex_bias = ex_bias * self.multiplier()
+        if self.dora_scale is not None:
+            updown = self.apply_weight_decompose(updown, orig_weight)
+        return updown * self.calc_scale() * self.multiplier(), ex_bias
+
+    def calc_updown(self, target):
+        raise NotImplementedError
+
+    def forward(self, x, y):
+        raise NotImplementedError

modules/lora/network_full.py

@@ -0,0 +1,26 @@
+import modules.lora.network as network
+
+
+class ModuleTypeFull(network.ModuleType):
+    def create_module(self, net: network.Network, weights: network.NetworkWeights):
+        if all(x in weights.w for x in ["diff"]):
+            return NetworkModuleFull(net, weights)
+        return None
+
+
+class NetworkModuleFull(network.NetworkModule): # pylint: disable=abstract-method
+    def __init__(self,  net: network.Network, weights: network.NetworkWeights):
+        super().__init__(net, weights)
+
+        self.weight = weights.w.get("diff")
+        self.ex_bias = weights.w.get("diff_b")
+
+    def calc_updown(self, target):
+        output_shape = self.weight.shape
+        updown = self.weight.to(target.device, dtype=target.dtype)
+        if self.ex_bias is not None:
+            ex_bias = self.ex_bias.to(target.device, dtype=target.dtype)
+        else:
+            ex_bias = None
+
+        return self.finalize_updown(updown, target, output_shape, ex_bias)

modules/lora/network_glora.py

@@ -0,0 +1,30 @@
+import modules.lora.network as network
+
+
+class ModuleTypeGLora(network.ModuleType):
+    def create_module(self, net: network.Network, weights: network.NetworkWeights):
+        if all(x in weights.w for x in ["a1.weight", "a2.weight", "alpha", "b1.weight", "b2.weight"]):
+            return NetworkModuleGLora(net, weights)
+        return None
+
+# adapted from https://github.com/KohakuBlueleaf/LyCORIS
+class NetworkModuleGLora(network.NetworkModule): # pylint: disable=abstract-method
+    def __init__(self,  net: network.Network, weights: network.NetworkWeights):
+        super().__init__(net, weights)
+
+        if hasattr(self.sd_module, 'weight'):
+            self.shape = self.sd_module.weight.shape
+
+        self.w1a = weights.w["a1.weight"]
+        self.w1b = weights.w["b1.weight"]
+        self.w2a = weights.w["a2.weight"]
+        self.w2b = weights.w["b2.weight"]
+
+    def calc_updown(self, target): # pylint: disable=arguments-differ
+        w1a = self.w1a.to(target.device, dtype=target.dtype)
+        w1b = self.w1b.to(target.device, dtype=target.dtype)
+        w2a = self.w2a.to(target.device, dtype=target.dtype)
+        w2b = self.w2b.to(target.device, dtype=target.dtype)
+        output_shape = [w1a.size(0), w1b.size(1)]
+        updown = (w2b @ w1b) + ((target @ w2a) @ w1a)
+        return self.finalize_updown(updown, target, output_shape)

modules/lora/network_hada.py

@@ -0,0 +1,46 @@
+import modules.lora.lyco_helpers as lyco_helpers
+import modules.lora.network as network
+
+
+class ModuleTypeHada(network.ModuleType):
+    def create_module(self, net: network.Network, weights: network.NetworkWeights):
+        if all(x in weights.w for x in ["hada_w1_a", "hada_w1_b", "hada_w2_a", "hada_w2_b"]):
+            return NetworkModuleHada(net, weights)
+        return None
+
+
+class NetworkModuleHada(network.NetworkModule): # pylint: disable=abstract-method
+    def __init__(self,  net: network.Network, weights: network.NetworkWeights):
+        super().__init__(net, weights)
+        if hasattr(self.sd_module, 'weight'):
+            self.shape = self.sd_module.weight.shape
+        self.w1a = weights.w["hada_w1_a"]
+        self.w1b = weights.w["hada_w1_b"]
+        self.dim = self.w1b.shape[0]
+        self.w2a = weights.w["hada_w2_a"]
+        self.w2b = weights.w["hada_w2_b"]
+        self.t1 = weights.w.get("hada_t1")
+        self.t2 = weights.w.get("hada_t2")
+
+    def calc_updown(self, target):
+        w1a = self.w1a.to(target.device, dtype=target.dtype)
+        w1b = self.w1b.to(target.device, dtype=target.dtype)
+        w2a = self.w2a.to(target.device, dtype=target.dtype)
+        w2b = self.w2b.to(target.device, dtype=target.dtype)
+        output_shape = [w1a.size(0), w1b.size(1)]
+        if self.t1 is not None:
+            output_shape = [w1a.size(1), w1b.size(1)]
+            t1 = self.t1.to(target.device, dtype=target.dtype)
+            updown1 = lyco_helpers.make_weight_cp(t1, w1a, w1b)
+            output_shape += t1.shape[2:]
+        else:
+            if len(w1b.shape) == 4:
+                output_shape += w1b.shape[2:]
+            updown1 = lyco_helpers.rebuild_conventional(w1a, w1b, output_shape)
+        if self.t2 is not None:
+            t2 = self.t2.to(target.device, dtype=target.dtype)
+            updown2 = lyco_helpers.make_weight_cp(t2, w2a, w2b)
+        else:
+            updown2 = lyco_helpers.rebuild_conventional(w2a, w2b, output_shape)
+        updown = updown1 * updown2
+        return self.finalize_updown(updown, target, output_shape)

modules/lora/network_ia3.py

@@ -0,0 +1,24 @@
+import modules.lora.network as network
+
+class ModuleTypeIa3(network.ModuleType):
+    def create_module(self, net: network.Network, weights: network.NetworkWeights):
+        if all(x in weights.w for x in ["weight"]):
+            return NetworkModuleIa3(net, weights)
+        return None
+
+
+class NetworkModuleIa3(network.NetworkModule): # pylint: disable=abstract-method
+    def __init__(self,  net: network.Network, weights: network.NetworkWeights):
+        super().__init__(net, weights)
+        self.w = weights.w["weight"]
+        self.on_input = weights.w["on_input"].item()
+
+    def calc_updown(self, target):
+        w = self.w.to(target.device, dtype=target.dtype)
+        output_shape = [w.size(0), target.size(1)]
+        if self.on_input:
+            output_shape.reverse()
+        else:
+            w = w.reshape(-1, 1)
+        updown = target * w
+        return self.finalize_updown(updown, target, output_shape)

modules/lora/network_lokr.py

@@ -0,0 +1,57 @@
+import torch
+import modules.lora.lyco_helpers as lyco_helpers
+import modules.lora.network as network
+
+
+class ModuleTypeLokr(network.ModuleType):
+    def create_module(self, net: network.Network, weights: network.NetworkWeights):
+        has_1 = "lokr_w1" in weights.w or ("lokr_w1_a" in weights.w and "lokr_w1_b" in weights.w)
+        has_2 = "lokr_w2" in weights.w or ("lokr_w2_a" in weights.w and "lokr_w2_b" in weights.w)
+        if has_1 and has_2:
+            return NetworkModuleLokr(net, weights)
+        return None
+
+
+def make_kron(orig_shape, w1, w2):
+    if len(w2.shape) == 4:
+        w1 = w1.unsqueeze(2).unsqueeze(2)
+    w2 = w2.contiguous()
+    return torch.kron(w1, w2).reshape(orig_shape)
+
+
+class NetworkModuleLokr(network.NetworkModule): # pylint: disable=abstract-method
+    def __init__(self,  net: network.Network, weights: network.NetworkWeights):
+        super().__init__(net, weights)
+        self.w1 = weights.w.get("lokr_w1")
+        self.w1a = weights.w.get("lokr_w1_a")
+        self.w1b = weights.w.get("lokr_w1_b")
+        self.dim = self.w1b.shape[0] if self.w1b is not None else self.dim
+        self.w2 = weights.w.get("lokr_w2")
+        self.w2a = weights.w.get("lokr_w2_a")
+        self.w2b = weights.w.get("lokr_w2_b")
+        self.dim = self.w2b.shape[0] if self.w2b is not None else self.dim
+        self.t2 = weights.w.get("lokr_t2")
+
+    def calc_updown(self, target):
+        if self.w1 is not None:
+            w1 = self.w1.to(target.device, dtype=target.dtype)
+        else:
+            w1a = self.w1a.to(target.device, dtype=target.dtype)
+            w1b = self.w1b.to(target.device, dtype=target.dtype)
+            w1 = w1a @ w1b
+        if self.w2 is not None:
+            w2 = self.w2.to(target.device, dtype=target.dtype)
+        elif self.t2 is None:
+            w2a = self.w2a.to(target.device, dtype=target.dtype)
+            w2b = self.w2b.to(target.device, dtype=target.dtype)
+            w2 = w2a @ w2b
+        else:
+            t2 = self.t2.to(target.device, dtype=target.dtype)
+            w2a = self.w2a.to(target.device, dtype=target.dtype)
+            w2b = self.w2b.to(target.device, dtype=target.dtype)
+            w2 = lyco_helpers.make_weight_cp(t2, w2a, w2b)
+        output_shape = [w1.size(0) * w2.size(0), w1.size(1) * w2.size(1)]
+        if len(target.shape) == 4:
+            output_shape = target.shape
+        updown = make_kron(output_shape, w1, w2)
+        return self.finalize_updown(updown, target, output_shape)

modules/lora/network_lora.py

@@ -0,0 +1,78 @@
+import torch
+import diffusers.models.lora as diffusers_lora
+import modules.lora.lyco_helpers as lyco_helpers
+import modules.lora.network as network
+from modules import devices
+
+
+class ModuleTypeLora(network.ModuleType):
+    def create_module(self, net: network.Network, weights: network.NetworkWeights):
+        if all(x in weights.w for x in ["lora_up.weight", "lora_down.weight"]):
+            return NetworkModuleLora(net, weights)
+        return None
+
+
+class NetworkModuleLora(network.NetworkModule):
+
+    def __init__(self,  net: network.Network, weights: network.NetworkWeights):
+        super().__init__(net, weights)
+        self.up_model = self.create_module(weights.w, "lora_up.weight")
+        self.down_model = self.create_module(weights.w, "lora_down.weight")
+        self.mid_model = self.create_module(weights.w, "lora_mid.weight", none_ok=True)
+        self.dim = weights.w["lora_down.weight"].shape[0]
+
+    def create_module(self, weights, key, none_ok=False):
+        from modules.shared import opts
+        weight = weights.get(key)
+        if weight is None and none_ok:
+            return None
+        linear_modules = [torch.nn.Linear, torch.nn.modules.linear.NonDynamicallyQuantizableLinear, torch.nn.MultiheadAttention, diffusers_lora.LoRACompatibleLinear]
+        is_linear = type(self.sd_module) in linear_modules or self.sd_module.__class__.__name__ in {"NNCFLinear", "QLinear", "Linear4bit"}
+        is_conv = type(self.sd_module) in [torch.nn.Conv2d, diffusers_lora.LoRACompatibleConv] or self.sd_module.__class__.__name__ in {"NNCFConv2d", "QConv2d"}
+        if is_linear:
+            weight = weight.reshape(weight.shape[0], -1)
+            module = torch.nn.Linear(weight.shape[1], weight.shape[0], bias=False)
+        elif is_conv and key == "lora_down.weight" or key == "dyn_up":
+            if len(weight.shape) == 2:
+                weight = weight.reshape(weight.shape[0], -1, 1, 1)
+            if weight.shape[2] != 1 or weight.shape[3] != 1:
+                module = torch.nn.Conv2d(weight.shape[1], weight.shape[0], self.sd_module.kernel_size, self.sd_module.stride, self.sd_module.padding, bias=False)
+            else:
+                module = torch.nn.Conv2d(weight.shape[1], weight.shape[0], (1, 1), bias=False)
+        elif is_conv and key == "lora_mid.weight":
+            module = torch.nn.Conv2d(weight.shape[1], weight.shape[0], self.sd_module.kernel_size, self.sd_module.stride, self.sd_module.padding, bias=False)
+        elif is_conv and key == "lora_up.weight" or key == "dyn_down":
+            module = torch.nn.Conv2d(weight.shape[1], weight.shape[0], (1, 1), bias=False)
+        else:
+            raise AssertionError(f'Lora unsupported: layer={self.network_key} type={type(self.sd_module).__name__}')
+        with torch.no_grad():
+            if weight.shape != module.weight.shape:
+                weight = weight.reshape(module.weight.shape)
+            module.weight.copy_(weight)
+        if opts.lora_load_gpu:
+            module = module.to(device=devices.device, dtype=devices.dtype)
+        module.weight.requires_grad_(False)
+        return module
+
+    def calc_updown(self, target): # pylint: disable=W0237
+        target_dtype = target.dtype if target.dtype != torch.uint8 else self.up_model.weight.dtype
+        up = self.up_model.weight.to(target.device, dtype=target_dtype)
+        down = self.down_model.weight.to(target.device, dtype=target_dtype)
+        output_shape = [up.size(0), down.size(1)]
+        if self.mid_model is not None:
+            # cp-decomposition
+            mid = self.mid_model.weight.to(target.device, dtype=target_dtype)
+            updown = lyco_helpers.rebuild_cp_decomposition(up, down, mid)
+            output_shape += mid.shape[2:]
+        else:
+            if len(down.shape) == 4:
+                output_shape += down.shape[2:]
+            updown = lyco_helpers.rebuild_conventional(up, down, output_shape, self.network.dyn_dim)
+        return self.finalize_updown(updown, target, output_shape)
+
+    def forward(self, x, y):
+        self.up_model.to(device=devices.device)
+        self.down_model.to(device=devices.device)
+        if hasattr(y, "scale"):
+            return y(scale=1) + self.up_model(self.down_model(x)) * self.multiplier() * self.calc_scale()
+        return y + self.up_model(self.down_model(x)) * self.multiplier() * self.calc_scale()

modules/lora/network_norm.py

@@ -0,0 +1,24 @@
+import modules.lora.network as network
+
+
+class ModuleTypeNorm(network.ModuleType):
+    def create_module(self, net: network.Network, weights: network.NetworkWeights):
+        if all(x in weights.w for x in ["w_norm", "b_norm"]):
+            return NetworkModuleNorm(net, weights)
+        return None
+
+
+class NetworkModuleNorm(network.NetworkModule): # pylint: disable=abstract-method
+    def __init__(self,  net: network.Network, weights: network.NetworkWeights):
+        super().__init__(net, weights)
+        self.w_norm = weights.w.get("w_norm")
+        self.b_norm = weights.w.get("b_norm")
+
+    def calc_updown(self, target):
+        output_shape = self.w_norm.shape
+        updown = self.w_norm.to(target.device, dtype=target.dtype)
+        if self.b_norm is not None:
+            ex_bias = self.b_norm.to(target.device, dtype=target.dtype)
+        else:
+            ex_bias = None
+        return self.finalize_updown(updown, target, output_shape, ex_bias)

modules/lora/network_oft.py

@@ -0,0 +1,82 @@
+import torch
+from einops import rearrange
+import modules.lora.network as network
+from modules.lora.lyco_helpers import factorization
+
+
+class ModuleTypeOFT(network.ModuleType):
+    def create_module(self, net: network.Network, weights: network.NetworkWeights):
+        if all(x in weights.w for x in ["oft_blocks"]) or all(x in weights.w for x in ["oft_diag"]):
+            return NetworkModuleOFT(net, weights)
+        return None
+
+
+# Supports both kohya-ss' implementation of COFT  https://github.com/kohya-ss/sd-scripts/blob/main/networks/oft.py
+# and KohakuBlueleaf's implementation of OFT/COFT https://github.com/KohakuBlueleaf/LyCORIS/blob/dev/lycoris/modules/diag_oft.py
+class NetworkModuleOFT(network.NetworkModule): # pylint: disable=abstract-method
+    def __init__(self,  net: network.Network, weights: network.NetworkWeights):
+        super().__init__(net, weights)
+        self.lin_module = None
+        self.org_module: list[torch.Module] = [self.sd_module]
+        self.scale = 1.0
+
+        # kohya-ss
+        if "oft_blocks" in weights.w.keys():
+            self.is_kohya = True
+            self.oft_blocks = weights.w["oft_blocks"] # (num_blocks, block_size, block_size)
+            self.alpha = weights.w["alpha"] # alpha is constraint
+            self.dim = self.oft_blocks.shape[0] # lora dim
+        # LyCORIS
+        elif "oft_diag" in weights.w.keys():
+            self.is_kohya = False
+            self.oft_blocks = weights.w["oft_diag"]
+            # self.alpha is unused
+            self.dim = self.oft_blocks.shape[1] # (num_blocks, block_size, block_size)
+
+        is_linear = type(self.sd_module) in [torch.nn.Linear, torch.nn.modules.linear.NonDynamicallyQuantizableLinear]
+        is_conv = type(self.sd_module) in [torch.nn.Conv2d]
+        is_other_linear = type(self.sd_module) in [torch.nn.MultiheadAttention] # unsupported
+
+        if is_linear:
+            self.out_dim = self.sd_module.out_features
+        elif is_conv:
+            self.out_dim = self.sd_module.out_channels
+        elif is_other_linear:
+            self.out_dim = self.sd_module.embed_dim
+
+        if self.is_kohya:
+            self.constraint = self.alpha * self.out_dim
+            self.num_blocks = self.dim
+            self.block_size = self.out_dim // self.dim
+        else:
+            self.constraint = None
+            self.block_size, self.num_blocks = factorization(self.out_dim, self.dim)
+
+    def calc_updown(self, target):
+        oft_blocks = self.oft_blocks.to(target.device, dtype=target.dtype)
+        eye = torch.eye(self.block_size, device=target.device)
+        constraint = self.constraint.to(target.device)
+
+        if self.is_kohya:
+            block_Q = oft_blocks - oft_blocks.transpose(1, 2) # ensure skew-symmetric orthogonal matrix
+            norm_Q = torch.norm(block_Q.flatten()).to(target.device)
+            new_norm_Q = torch.clamp(norm_Q, max=constraint)
+            block_Q = block_Q * ((new_norm_Q + 1e-8) / (norm_Q + 1e-8))
+            mat1 = eye + block_Q
+            mat2 = (eye - block_Q).float().inverse()
+            oft_blocks = torch.matmul(mat1, mat2)
+
+        R = oft_blocks.to(target.device, dtype=target.dtype)
+
+        # This errors out for MultiheadAttention, might need to be handled up-stream
+        merged_weight = rearrange(target, '(k n) ... -> k n ...', k=self.num_blocks, n=self.block_size)
+        merged_weight = torch.einsum(
+            'k n m, k n ... -> k m ...',
+            R,
+            merged_weight
+        )
+        merged_weight = rearrange(merged_weight, 'k m ... -> (k m) ...')
+
+        updown = merged_weight.to(target.device, dtype=target.dtype) - target
+        output_shape = target.shape
+        return self.finalize_updown(updown, target, output_shape)

modules/lora/network_overrides.py

@@ -0,0 +1,49 @@
+from modules import shared
+
+
+maybe_diffusers = [ # forced if lora_maybe_diffusers is enabled
+    'aaebf6360f7d', # sd15-lcm
+    '3d18b05e4f56', # sdxl-lcm
+    'b71dcb732467', # sdxl-tcd
+    '813ea5fb1c67', # sdxl-turbo
+    # not really needed, but just in case
+    '5a48ac366664', # hyper-sd15-1step
+    'ee0ff23dcc42', # hyper-sd15-2step
+    'e476eb1da5df', # hyper-sd15-4step
+    'ecb844c3f3b0', # hyper-sd15-8step
+    '1ab289133ebb', # hyper-sd15-8step-cfg
+    '4f494295edb1', # hyper-sdxl-8step
+    'ca14a8c621f8', # hyper-sdxl-8step-cfg
+    '1c88f7295856', # hyper-sdxl-4step
+    'fdd5dcd1d88a', # hyper-sdxl-2step
+    '8cca3706050b', # hyper-sdxl-1step
+]
+
+force_diffusers = [ # forced always
+    '816d0eed49fd', # flash-sdxl
+    'c2ec22757b46', # flash-sd15
+]
+
+force_models = [ # forced always
+    'sc',
+    # 'sd3',
+    'kandinsky',
+    'hunyuandit',
+    'auraflow',
+]
+
+force_classes = [ # forced always
+]
+
+
+def check_override(shorthash=''):
+    force = False
+    force = force or (shared.sd_model_type in force_models)
+    force = force or (shared.sd_model.__class__.__name__ in force_classes)
+    if len(shorthash) < 4:
+        return force
+    force = force or (any(x.startswith(shorthash) for x in maybe_diffusers) if shared.opts.lora_maybe_diffusers else False)
+    force = force or any(x.startswith(shorthash) for x in force_diffusers)
+    if force and shared.opts.lora_maybe_diffusers:
+        shared.log.debug('LoRA override: force diffusers')
+    return force

modules/lora/networks.py

@@ -0,0 +1,502 @@
+from typing import Union, List
+import os
+import re
+import time
+import concurrent
+from contextlib import nullcontext
+import torch
+import diffusers.models.lora
+import rich.progress as rp
+
+import modules.lora.network as network
+import modules.lora.network_lora as network_lora
+import modules.lora.network_hada as network_hada
+import modules.lora.network_ia3 as network_ia3
+import modules.lora.network_oft as network_oft
+import modules.lora.network_lokr as network_lokr
+import modules.lora.network_full as network_full
+import modules.lora.network_norm as network_norm
+import modules.lora.network_glora as network_glora
+import modules.lora.network_overrides as network_overrides
+import modules.lora.lora_convert as lora_convert
+from modules.lora.extra_networks_lora import ExtraNetworkLora
+from modules import shared, devices, sd_models, sd_models_compile, errors, files_cache, model_quant
+
+
+debug = os.environ.get('SD_LORA_DEBUG', None) is not None
+extra_network_lora = ExtraNetworkLora()
+available_networks = {}
+available_network_aliases = {}
+loaded_networks: List[network.Network] = []
+timer = { 'list': 0, 'load': 0, 'backup': 0, 'calc': 0, 'apply': 0, 'move': 0, 'restore': 0, 'deactivate': 0 }
+bnb = None
+lora_cache = {}
+diffuser_loaded = []
+diffuser_scales = []
+available_network_hash_lookup = {}
+forbidden_network_aliases = {}
+re_network_name = re.compile(r"(.*)\s*\([0-9a-fA-F]+\)")
+module_types = [
+    network_lora.ModuleTypeLora(),
+    network_hada.ModuleTypeHada(),
+    network_ia3.ModuleTypeIa3(),
+    network_oft.ModuleTypeOFT(),
+    network_lokr.ModuleTypeLokr(),
+    network_full.ModuleTypeFull(),
+    network_norm.ModuleTypeNorm(),
+    network_glora.ModuleTypeGLora(),
+]
+
+
+def total_time():
+    return sum(timer.values())
+
+
+def get_timers():
+    t = { 'total': round(sum(timer.values()), 2) }
+    for k, v in timer.items():
+        if v > 0.1:
+            t[k] = round(v, 2)
+    return t
+
+
+# section: load networks from disk
+
+def load_diffusers(name, network_on_disk, lora_scale=shared.opts.extra_networks_default_multiplier) -> Union[network.Network, None]:
+    name = name.replace(".", "_")
+    shared.log.debug(f'Load network: 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 shared.native:
+        return None
+    if not hasattr(shared.sd_model, 'load_lora_weights'):
+        shared.log.error(f'Load network: type=LoRA class={shared.sd_model.__class__} does not implement load lora')
+        return None
+    try:
+        shared.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'Load network: type=LoRA name="{name}" diffusers unsupported format')
+            else:
+                shared.log.error(f'Load network: type=LoRA name="{name}" {e}')
+            if debug:
+                errors.display(e, "LoRA")
+            return None
+    if name not in diffuser_loaded:
+        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)
+    return net
+
+
+def load_safetensors(name, network_on_disk) -> Union[network.Network, None]:
+    if not shared.sd_loaded:
+        return None
+
+    cached = lora_cache.get(name, None)
+    if debug:
+        shared.log.debug(f'Load network: type=LoRA name="{name}" file="{network_on_disk.filename}" type=lora {"cached" if cached else ""}')
+    if cached is not None:
+        return cached
+    net = network.Network(name, network_on_disk)
+    net.mtime = os.path.getmtime(network_on_disk.filename)
+    sd = sd_models.read_state_dict(network_on_disk.filename, what='network')
+    if shared.sd_model_type == 'f1':  # if kohya flux lora, convert state_dict
+        sd = lora_convert._convert_kohya_flux_lora_to_diffusers(sd) or sd  # pylint: disable=protected-access
+    if shared.sd_model_type == 'sd3':  # if kohya flux lora, convert state_dict
+        try:
+            sd = lora_convert._convert_kohya_sd3_lora_to_diffusers(sd) or sd  # pylint: disable=protected-access
+        except ValueError:  # EAFP for diffusers PEFT keys
+            pass
+    lora_convert.assign_network_names_to_compvis_modules(shared.sd_model)
+    keys_failed_to_match = {}
+    matched_networks = {}
+    bundle_embeddings = {}
+    convert = lora_convert.KeyConvert()
+    for key_network, weight in sd.items():
+        parts = key_network.split('.')
+        if parts[0] == "bundle_emb":
+            emb_name, vec_name = parts[1], key_network.split(".", 2)[-1]
+            emb_dict = bundle_embeddings.get(emb_name, {})
+            emb_dict[vec_name] = weight
+            bundle_embeddings[emb_name] = emb_dict
+            continue
+        if len(parts) > 5: # messy handler for diffusers peft lora
+            key_network_without_network_parts = '_'.join(parts[:-2])
+            if not key_network_without_network_parts.startswith('lora_'):
+                key_network_without_network_parts = 'lora_' + key_network_without_network_parts
+            network_part = '.'.join(parts[-2:]).replace('lora_A', 'lora_down').replace('lora_B', 'lora_up')
+        else:
+            key_network_without_network_parts, network_part = key_network.split(".", 1)
+        key, sd_module = convert(key_network_without_network_parts)
+        if sd_module is None:
+            keys_failed_to_match[key_network] = key
+            continue
+        if key not in matched_networks:
+            matched_networks[key] = network.NetworkWeights(network_key=key_network, sd_key=key, w={}, sd_module=sd_module)
+        matched_networks[key].w[network_part] = weight
+    network_types = []
+    for key, weights in matched_networks.items():
+        net_module = None
+        for nettype in module_types:
+            net_module = nettype.create_module(net, weights)
+            if net_module is not None:
+                network_types.append(nettype.__class__.__name__)
+                break
+        if net_module is None:
+            shared.log.error(f'LoRA unhandled: name={name} key={key} weights={weights.w.keys()}')
+        else:
+            net.modules[key] = net_module
+    if len(keys_failed_to_match) > 0:
+        shared.log.warning(f'LoRA name="{name}" type={set(network_types)} unmatched={len(keys_failed_to_match)} matched={len(matched_networks)}')
+        if debug:
+            shared.log.debug(f'LoRA name="{name}" unmatched={keys_failed_to_match}')
+    else:
+        shared.log.debug(f'LoRA name="{name}" type={set(network_types)} keys={len(matched_networks)}')
+    if len(matched_networks) == 0:
+        return None
+    lora_cache[name] = net
+    net.bundle_embeddings = bundle_embeddings
+    return net
+
+
+def maybe_recompile_model(names, te_multipliers):
+    recompile_model = False
+    if shared.compiled_model_state is not None and shared.compiled_model_state.is_compiled:
+        if len(names) == len(shared.compiled_model_state.lora_model):
+            for i, name in enumerate(names):
+                if shared.compiled_model_state.lora_model[
+                    i] != f"{name}:{te_multipliers[i] if te_multipliers else shared.opts.extra_networks_default_multiplier}":
+                    recompile_model = True
+                    shared.compiled_model_state.lora_model = []
+                    break
+            if not recompile_model:
+                if len(loaded_networks) > 0 and debug:
+                    shared.log.debug('Model Compile: Skipping LoRa loading')
+                return recompile_model
+        else:
+            recompile_model = True
+            shared.compiled_model_state.lora_model = []
+    if recompile_model:
+        backup_cuda_compile = shared.opts.cuda_compile
+        sd_models.unload_model_weights(op='model')
+        shared.opts.cuda_compile = []
+        sd_models.reload_model_weights(op='model')
+        shared.opts.cuda_compile = backup_cuda_compile
+    return recompile_model
+
+
+def list_available_networks():
+    t0 = time.time()
+    available_networks.clear()
+    available_network_aliases.clear()
+    forbidden_network_aliases.clear()
+    available_network_hash_lookup.clear()
+    forbidden_network_aliases.update({"none": 1, "Addams": 1})
+    if not os.path.exists(shared.cmd_opts.lora_dir):
+        shared.log.warning(f'LoRA directory not found: path="{shared.cmd_opts.lora_dir}"')
+
+    def add_network(filename):
+        if not os.path.isfile(filename):
+            return
+        name = os.path.splitext(os.path.basename(filename))[0]
+        name = name.replace('.', '_')
+        try:
+            entry = network.NetworkOnDisk(name, filename)
+            available_networks[entry.name] = entry
+            if entry.alias in available_network_aliases:
+                forbidden_network_aliases[entry.alias.lower()] = 1
+            if shared.opts.lora_preferred_name == 'filename':
+                available_network_aliases[entry.name] = entry
+            else:
+                available_network_aliases[entry.alias] = entry
+            if entry.shorthash:
+                available_network_hash_lookup[entry.shorthash] = entry
+        except OSError as e:  # should catch FileNotFoundError and PermissionError etc.
+            shared.log.error(f'LoRA: filename="{filename}" {e}')
+
+    candidates = list(files_cache.list_files(shared.cmd_opts.lora_dir, ext_filter=[".pt", ".ckpt", ".safetensors"]))
+    with concurrent.futures.ThreadPoolExecutor(max_workers=shared.max_workers) as executor:
+        for fn in candidates:
+            executor.submit(add_network, fn)
+    t1 = time.time()
+    timer['list'] = t1 - t0
+    shared.log.info(f'Available LoRAs: path="{shared.cmd_opts.lora_dir}" items={len(available_networks)} folders={len(forbidden_network_aliases)} time={t1 - t0:.2f}')
+
+
+def network_load(names, te_multipliers=None, unet_multipliers=None, dyn_dims=None):
+    timer['list'] = 0
+    networks_on_disk: list[network.NetworkOnDisk] = [available_network_aliases.get(name, None) for name in names]
+    if any(x is None for x in networks_on_disk):
+        list_available_networks()
+        networks_on_disk: list[network.NetworkOnDisk] = [available_network_aliases.get(name, None) for name in names]
+    failed_to_load_networks = []
+    recompile_model = maybe_recompile_model(names, te_multipliers)
+
+    loaded_networks.clear()
+    diffuser_loaded.clear()
+    diffuser_scales.clear()
+    t0 = time.time()
+
+    for i, (network_on_disk, name) in enumerate(zip(networks_on_disk, names)):
+        net = None
+        if network_on_disk is not None:
+            shorthash = getattr(network_on_disk, 'shorthash', '').lower()
+            if debug:
+                shared.log.debug(f'Load network: type=LoRA name="{name}" file="{network_on_disk.filename}" hash="{shorthash}"')
+            try:
+                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}")
+                if shared.opts.lora_force_diffusers or network_overrides.check_override(shorthash): # OpenVINO only works with Diffusers LoRa loading
+                    net = load_diffusers(name, network_on_disk, lora_scale=te_multipliers[i] if te_multipliers else shared.opts.extra_networks_default_multiplier)
+                else:
+                    net = load_safetensors(name, network_on_disk)
+                if net is not None:
+                    net.mentioned_name = name
+                    network_on_disk.read_hash()
+            except Exception as e:
+                shared.log.error(f'Load network: type=LoRA file="{network_on_disk.filename}" {e}')
+                if debug:
+                    errors.display(e, 'LoRA')
+                continue
+        if net is None:
+            failed_to_load_networks.append(name)
+            shared.log.error(f'Load network: type=LoRA name="{name}" detected={network_on_disk.sd_version if network_on_disk is not None else None} failed')
+            continue
+        shared.sd_model.embedding_db.load_diffusers_embedding(None, net.bundle_embeddings)
+        net.te_multiplier = te_multipliers[i] if te_multipliers else shared.opts.extra_networks_default_multiplier
+        net.unet_multiplier = unet_multipliers[i] if unet_multipliers else shared.opts.extra_networks_default_multiplier
+        net.dyn_dim = dyn_dims[i] if dyn_dims else shared.opts.extra_networks_default_multiplier
+        loaded_networks.append(net)
+
+    while len(lora_cache) > shared.opts.lora_in_memory_limit:
+        name = next(iter(lora_cache))
+        lora_cache.pop(name, None)
+
+    if len(diffuser_loaded) > 0:
+        shared.log.debug(f'Load network: type=LoRA loaded={diffuser_loaded} available={shared.sd_model.get_list_adapters()} active={shared.sd_model.get_active_adapters()} scales={diffuser_scales}')
+        try:
+            shared.sd_model.set_adapters(adapter_names=diffuser_loaded, adapter_weights=diffuser_scales)
+            if shared.opts.lora_fuse_diffusers:
+                shared.sd_model.fuse_lora(adapter_names=diffuser_loaded, lora_scale=1.0, fuse_unet=True, fuse_text_encoder=True) # fuse uses fixed scale since later apply does the scaling
+                shared.sd_model.unload_lora_weights()
+        except Exception as e:
+            shared.log.error(f'Load network: type=LoRA {e}')
+            if debug:
+                errors.display(e, 'LoRA')
+
+    if len(loaded_networks) > 0 and debug:
+        shared.log.debug(f'Load network: type=LoRA loaded={len(loaded_networks)} cache={list(lora_cache)}')
+
+    if recompile_model:
+        shared.log.info("Load network: type=LoRA recompiling model")
+        backup_lora_model = shared.compiled_model_state.lora_model
+        if 'Model' in shared.opts.cuda_compile:
+            shared.sd_model = sd_models_compile.compile_diffusers(shared.sd_model)
+
+        shared.compiled_model_state.lora_model = backup_lora_model
+
+    if len(loaded_networks) > 0:
+        devices.torch_gc()
+
+    t1 = time.time()
+    timer['load'] = t1 - t0
+
+
+# section: process loaded networks
+
+def network_backup_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.GroupNorm, torch.nn.LayerNorm, diffusers.models.lora.LoRACompatibleLinear, diffusers.models.lora.LoRACompatibleConv], weight, network_layer_name, wanted_names):
+    global bnb # pylint: disable=W0603
+    backup_size = 0
+    if len(loaded_networks) > 0 and network_layer_name is not None and any([net.modules.get(network_layer_name, None) for net in loaded_networks]): # noqa: C419 # pylint: disable=R1729
+        t0 = time.time()
+
+        weights_backup = getattr(self, "network_weights_backup", None)
+        if weights_backup is None and wanted_names != (): # pylint: disable=C1803
+            self.network_weights_backup = None
+            if shared.opts.lora_fuse_diffusers:
+                weights_backup = True
+            elif getattr(weight, "quant_type", None) in ['nf4', 'fp4']:
+                if bnb is None:
+                    bnb = model_quant.load_bnb('Load network: type=LoRA', silent=True)
+                if bnb is not None:
+                    with devices.inference_context():
+                        weights_backup = bnb.functional.dequantize_4bit(weight, quant_state=weight.quant_state, quant_type=weight.quant_type, blocksize=weight.blocksize,)
+                        self.quant_state = weight.quant_state
+                        self.quant_type = weight.quant_type
+                        self.blocksize = weight.blocksize
+                else:
+                    weights_backup = weight.clone()
+            else:
+                weights_backup = weight.clone()
+            if shared.opts.lora_offload_backup and weights_backup is not None and isinstance(weights_backup, torch.Tensor):
+                weights_backup = weights_backup.to(devices.cpu)
+            self.network_weights_backup = weights_backup
+        bias_backup = getattr(self, "network_bias_backup", None)
+        if bias_backup is None:
+            if getattr(self, 'bias', None) is not None:
+                if shared.opts.lora_fuse_diffusers:
+                    bias_backup = True
+                else:
+                    bias_backup = self.bias.clone()
+            else:
+                bias_backup = None
+            if shared.opts.lora_offload_backup and bias_backup is not None and isinstance(bias_backup, torch.Tensor):
+                bias_backup = bias_backup.to(devices.cpu)
+            self.network_bias_backup = bias_backup
+
+        if getattr(self, 'network_weights_backup', None) is not None:
+            backup_size += self.network_weights_backup.numel() * self.network_weights_backup.element_size() if isinstance(self.network_weights_backup, torch.Tensor) else 0
+        if getattr(self, 'network_bias_backup', None) is not None:
+            backup_size += self.network_bias_backup.numel() * self.network_bias_backup.element_size() if isinstance(self.network_bias_backup, torch.Tensor) else 0
+        t1 = time.time()
+        timer['backup'] += t1 - t0
+    return backup_size
+
+
+def network_calc_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.GroupNorm, torch.nn.LayerNorm, diffusers.models.lora.LoRACompatibleLinear, diffusers.models.lora.LoRACompatibleConv], weight, network_layer_name):
+    if shared.opts.diffusers_offload_mode == "none":
+        self.to(devices.device, non_blocking=True)
+    batch_updown = None
+    batch_ex_bias = None
+    for net in loaded_networks:
+        module = net.modules.get(network_layer_name, None)
+        if module is not None and hasattr(self, 'weight'):
+            try:
+                t0 = time.time()
+                updown, ex_bias = module.calc_updown(weight)
+                t1 = time.time()
+                if batch_updown is not None and updown is not None:
+                    batch_updown += updown
+                else:
+                    batch_updown = updown
+                if batch_ex_bias is not None and ex_bias is not None:
+                    batch_ex_bias += ex_bias
+                else:
+                    batch_ex_bias = ex_bias
+                timer['calc'] += t1 - t0
+                if shared.opts.diffusers_offload_mode != "none":
+                    t0 = time.time()
+                    if batch_updown is not None:
+                        batch_updown = batch_updown.to(devices.cpu, non_blocking=True)
+                    if batch_ex_bias is not None:
+                        batch_ex_bias = batch_ex_bias.to(devices.cpu, non_blocking=True)
+                    if devices.backend == "ipex":
+                        # using non_blocking=True here causes NaNs on Intel
+                        torch.xpu.synchronize(devices.device)
+                    t1 = time.time()
+                    timer['move'] += t1 - t0
+            except RuntimeError as e:
+                extra_network_lora.errors[net.name] = extra_network_lora.errors.get(net.name, 0) + 1
+                if debug:
+                    module_name = net.modules.get(network_layer_name, None)
+                    shared.log.error(f'LoRA apply weight name="{net.name}" module="{module_name}" layer="{network_layer_name}" {e}')
+                    errors.display(e, 'LoRA')
+                    raise RuntimeError('LoRA apply weight') from e
+            continue
+        if module is None:
+            continue
+        shared.log.warning(f'LoRA network="{net.name}" layer="{network_layer_name}" unsupported operation')
+        extra_network_lora.errors[net.name] = extra_network_lora.errors.get(net.name, 0) + 1
+    return batch_updown, batch_ex_bias
+
+
+def network_apply_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.GroupNorm, torch.nn.LayerNorm, diffusers.models.lora.LoRACompatibleLinear, diffusers.models.lora.LoRACompatibleConv], updown, ex_bias):
+    t0 = time.time()
+    weights_backup = getattr(self, "network_weights_backup", None)
+    bias_backup = getattr(self, "network_bias_backup", None)
+    if weights_backup is None and bias_backup is None:
+        return None, None
+    if weights_backup is not None:
+        if isinstance(weights_backup, bool):
+            weights_backup = self.weight
+        else:
+            self.weight = None
+        if updown is not None and len(weights_backup.shape) == 4 and weights_backup.shape[1] == 9: # inpainting model. zero pad updown to make channel[1]  4 to 9
+            updown = torch.nn.functional.pad(updown, (0, 0, 0, 0, 0, 5))  # pylint: disable=not-callable
+        if updown is not None:
+            new_weight = weights_backup.to(devices.device, non_blocking=True) + updown.to(devices.device, non_blocking=True)
+            if getattr(self, "quant_type", None) in ['nf4', 'fp4'] and bnb is not None:
+                self.weight = bnb.nn.Params4bit(new_weight, quant_state=self.quant_state, quant_type=self.quant_type, blocksize=self.blocksize)
+            else:
+                self.weight = torch.nn.Parameter(new_weight, requires_grad=False)
+            del new_weight
+        else:
+            self.weight = torch.nn.Parameter(weights_backup, requires_grad=False)
+        if hasattr(self, "qweight") and hasattr(self, "freeze"):
+            self.freeze()
+    if bias_backup is not None:
+        if isinstance(bias_backup, bool):
+            bias_backup = self.bias
+        else:
+            self.bias = None
+        if ex_bias is not None:
+            new_weight = bias_backup.to(devices.device, non_blocking=True) + ex_bias.to(devices.device, non_blocking=True)
+            self.bias = torch.nn.Parameter(new_weight, requires_grad=False)
+            del new_weight
+        else:
+            self.bias = torch.nn.Parameter(bias_backup, requires_grad=False)
+    else:
+        self.bias = None
+    t1 = time.time()
+    timer['apply'] += t1 - t0
+    return self.weight.device, self.weight.dtype
+
+
+def network_deactivate():
+    pass
+
+def network_activate():
+    timer['backup'] = 0
+    timer['calc'] = 0
+    timer['apply'] = 0
+    timer['move'] = 0
+    sd_model = getattr(shared.sd_model, "pipe", shared.sd_model)  # wrapped model compatiblility
+    if shared.opts.diffusers_offload_mode == "sequential":
+        sd_models.disable_offload(sd_model)
+        sd_models.move_model(sd_model, device=devices.cpu)
+    modules = []
+    for component_name in ['text_encoder','text_encoder_2', 'unet', 'transformer']:
+        component = getattr(sd_model, component_name, None)
+        if component is not None and hasattr(component, 'named_modules'):
+            modules += list(component.named_modules())
+    if len(loaded_networks) > 0:
+        pbar = rp.Progress(rp.TextColumn('[cyan]Apply network: type=LoRA'), rp.BarColumn(), rp.TaskProgressColumn(), rp.TimeRemainingColumn(), rp.TimeElapsedColumn(), rp.TextColumn('[cyan]{task.description}'), console=shared.console)
+        task = pbar.add_task(description='' , total=len(modules))
+    else:
+        task = None
+        pbar = nullcontext()
+    with devices.inference_context(), pbar:
+        wanted_names = tuple((x.name, x.te_multiplier, x.unet_multiplier, x.dyn_dim) for x in loaded_networks) if len(loaded_networks) > 0 else ()
+        applied = 0
+        backup_size = 0
+        weights_devices = []
+        weights_dtypes = []
+        for _, module in modules:
+            network_layer_name = getattr(module, 'network_layer_name', None)
+            current_names = getattr(module, "network_current_names", ())
+            if shared.state.interrupted or network_layer_name is None or current_names == wanted_names:
+                if task is not None:
+                    pbar.update(task, advance=1, description=f'networks={len(loaded_networks)} skip')
+                continue
+            weight = getattr(module, 'weight', None)
+            weight = weight.to(devices.device, non_blocking=True) if weight is not None else None
+            backup_size += network_backup_weights(module, weight, network_layer_name, wanted_names)
+            batch_updown, batch_ex_bias = network_calc_weights(module, weight, network_layer_name)
+            weights_device, weights_dtype = network_apply_weights(module, batch_updown, batch_ex_bias)
+            weights_devices.append(weights_device)
+            weights_dtypes.append(weights_dtype)
+            if batch_updown is not None or batch_ex_bias is not None:
+                applied += 1
+            del weight, batch_updown, batch_ex_bias
+            module.network_current_names = wanted_names
+            if task is not None:
+                pbar.update(task, advance=1, description=f'networks={len(loaded_networks)} modules={len(modules)} apply={applied} backup={backup_size}')
+    weights_devices, weights_dtypes = list(set([x for x in weights_devices if x is not None])), list(set([x for x in weights_dtypes if x is not None])) # noqa: C403 # pylint: disable=R1718
+    if debug and len(loaded_networks) > 0:
+        shared.log.debug(f'Load network: type=LoRA networks={len(loaded_networks)} modules={len(modules)} apply={applied} device={weights_devices} dtype={weights_dtypes} backup={backup_size} fuse={shared.opts.lora_fuse_diffusers} time={get_timers()}')
+    modules.clear()
+    if shared.opts.diffusers_offload_mode == "sequential":
+        sd_models.set_diffuser_offload(sd_model, op="model")

modules/memstats.py

@@ -5,11 +5,12 @@ from modules import shared, errors
 
 fail_once = False
 
+def gb(val: float):
+    return round(val / 1024 / 1024 / 1024, 2)
+
+
 def memory_stats():
     global fail_once # pylint: disable=global-statement
-    def gb(val: float):
-        return round(val / 1024 / 1024 / 1024, 2)
-
     mem = {}
     try:
         process = psutil.Process(os.getpid())
@@ -38,3 +39,14 @@ def memory_stats():
     except Exception:
         pass
     return mem
+
+
+def ram_stats():
+    try:
+        process = psutil.Process(os.getpid())
+        res = process.memory_info()
+        ram_total = 100 * res.rss / process.memory_percent()
+        ram = { 'used': gb(res.rss), 'total': gb(ram_total) }
+        return ram
+    except Exception:
+        return { 'used': 0, 'total': 0 }

modules/model_flux.py

@@ -223,10 +223,8 @@ def load_flux(checkpoint_info, diffusers_load_config): # triggered by opts.sd_ch
     if shared.opts.sd_unet != 'None':
         try:
             debug(f'Load model: type=FLUX unet="{shared.opts.sd_unet}"')
-            _transformer = load_transformer(sd_unet.unet_dict[shared.opts.sd_unet])
-            if _transformer is not None:
-                transformer = _transformer
-            else:
+            transformer = load_transformer(sd_unet.unet_dict[shared.opts.sd_unet])
+            if transformer is None:
                 shared.opts.sd_unet = 'None'
                 sd_unet.failed_unet.append(shared.opts.sd_unet)
         except Exception as e:
@@ -334,6 +332,8 @@ def load_flux(checkpoint_info, diffusers_load_config): # triggered by opts.sd_ch
     text_encoder_1 = None
     text_encoder_2 = None
     vae = None
+    for k in kwargs.keys():
+        kwargs[k] = None
     devices.torch_gc()
 
     return pipe

modules/model_stablecascade.py

@@ -330,14 +330,12 @@ class StableCascadeDecoderPipelineFixed(diffusers.StableCascadeDecoderPipeline):
             elif output_type == "pil":
                 images = images.permute(0, 2, 3, 1).cpu().float().numpy()  # float() as bfloat16-> numpy doesnt work
                 images = self.numpy_to_pil(images)
+            if shared.opts.diffusers_offload_mode == "balanced":
+                shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model)
         else:
             images = latents
 
-        # Offload all models
-        if shared.opts.diffusers_offload_mode == "balanced":
-            shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model)
-        else:
-            self.maybe_free_model_hooks()
+        self.maybe_free_model_hooks()
 
         if not return_dict:
             return images

modules/modelloader.py

@@ -267,7 +267,7 @@ def download_diffusers_model(hub_id: str, cache_dir: str = None, download_config
 
 def load_diffusers_models(clear=True):
     excluded_models = []
-    t0 = time.time()
+    # t0 = time.time()
     place = shared.opts.diffusers_dir
     if place is None or len(place) == 0 or not os.path.isdir(place):
         place = os.path.join(models_path, 'Diffusers')
@@ -316,7 +316,7 @@ def load_diffusers_models(clear=True):
                 debug(f'Error analyzing diffusers model: "{folder}" {e}')
     except Exception as e:
         shared.log.error(f"Error listing diffusers: {place} {e}")
-    shared.log.debug(f'Scanning diffusers cache: folder="{place}" items={len(list(diffuser_repos))} time={time.time()-t0:.2f}')
+    # shared.log.debug(f'Scanning diffusers cache: folder="{place}" items={len(list(diffuser_repos))} time={time.time()-t0:.2f}')
     return diffuser_repos
 
 

modules/processing.py

@@ -472,5 +472,16 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
     if p.scripts is not None and isinstance(p.scripts, scripts.ScriptRunner) and not (shared.state.interrupted or shared.state.skipped):
         p.scripts.postprocess(p, processed)
     timer.process.record('post')
-    shared.log.info(f'Processed: images={len(output_images)} its={(p.steps * len(output_images)) / (t1 - t0):.2f} time={t1-t0:.2f} timers={timer.process.dct(min_time=0.02)} memory={memstats.memory_stats()}')
+    if not p.disable_extra_networks:
+        shared.log.info(f'Processed: images={len(output_images)} its={(p.steps * len(output_images)) / (t1 - t0):.2f} time={t1-t0:.2f} timers={timer.process.dct(min_time=0.02)} memory={memstats.memory_stats()}')
+
+    if shared.cmd_opts.malloc:
+        import tracemalloc
+        snapshot = tracemalloc.take_snapshot()
+        stats = snapshot.statistics('lineno')
+        shared.log.debug('Profile malloc:')
+        for stat in stats[:20]:
+            frame = stat.traceback[0]
+            shared.log.debug(f'  file="{frame.filename}":{frame.lineno} size={stat.size}')
+    devices.torch_gc(force=True)
     return processed

modules/processing_args.py

@@ -12,7 +12,8 @@ from modules.processing_helpers import resize_hires, fix_prompts, calculate_base
 from modules.api import helpers
 
 
-debug = shared.log.trace if os.environ.get('SD_DIFFUSERS_DEBUG', None) is not None else lambda *args, **kwargs: None
+debug_enabled = os.environ.get('SD_DIFFUSERS_DEBUG', None)
+debug_log = shared.log.trace if os.environ.get('SD_DIFFUSERS_DEBUG', None) is not None else lambda *args, **kwargs: None
 
 
 def task_specific_kwargs(p, model):
@@ -93,7 +94,8 @@ def task_specific_kwargs(p, model):
             'target_subject_category': getattr(p, 'prompt', '').split()[-1],
             'output_type': 'pil',
         }
-    debug(f'Diffusers task specific args: {task_args}')
+    if debug_enabled:
+        debug_log(f'Diffusers task specific args: {task_args}')
     return task_args
 
 
@@ -108,7 +110,8 @@ def set_pipeline_args(p, model, prompts: list, negative_prompts: list, prompts_2
     signature = inspect.signature(type(model).__call__, follow_wrapped=True)
     possible = list(signature.parameters)
 
-    debug(f'Diffusers pipeline possible: {possible}')
+    if debug_enabled:
+        debug_log(f'Diffusers pipeline possible: {possible}')
     prompts, negative_prompts, prompts_2, negative_prompts_2 = fix_prompts(prompts, negative_prompts, prompts_2, negative_prompts_2)
     steps = kwargs.get("num_inference_steps", None) or len(getattr(p, 'timesteps', ['1']))
     clip_skip = kwargs.pop("clip_skip", 1)
@@ -159,6 +162,8 @@ def set_pipeline_args(p, model, prompts: list, negative_prompts: list, prompts_2
                 args['negative_prompt'] = negative_prompts[0]
             else:
                 args['negative_prompt'] = negative_prompts
+    if prompt_parser_diffusers.embedder is not None and not prompt_parser_diffusers.embedder.scheduled_prompt: # not scheduled so we dont need it anymore
+        prompt_parser_diffusers.embedder = None
 
     if 'clip_skip' in possible and parser == 'fixed':
         if clip_skip == 1:
@@ -248,14 +253,16 @@ def set_pipeline_args(p, model, prompts: list, negative_prompts: list, prompts_2
         if arg in possible:
             args[arg] = task_kwargs[arg]
     task_args = getattr(p, 'task_args', {})
-    debug(f'Diffusers task args: {task_args}')
+    if debug_enabled:
+        debug_log(f'Diffusers task args: {task_args}')
     for k, v in task_args.items():
         if k in possible:
             args[k] = v
         else:
-            debug(f'Diffusers unknown task args: {k}={v}')
+            debug_log(f'Diffusers unknown task args: {k}={v}')
     cross_attention_args = getattr(p, 'cross_attention_kwargs', {})
-    debug(f'Diffusers cross-attention args: {cross_attention_args}')
+    if debug_enabled:
+        debug_log(f'Diffusers cross-attention args: {cross_attention_args}')
     for k, v in cross_attention_args.items():
         if args.get('cross_attention_kwargs', None) is None:
             args['cross_attention_kwargs'] = {}
@@ -273,7 +280,7 @@ def set_pipeline_args(p, model, prompts: list, negative_prompts: list, prompts_2
 
     # handle implicit controlnet
     if 'control_image' in possible and 'control_image' not in args and 'image' in args:
-        debug('Diffusers: set control image')
+        debug_log('Diffusers: set control image')
         args['control_image'] = args['image']
 
     sd_hijack_hypertile.hypertile_set(p, hr=len(getattr(p, 'init_images', [])) > 0)
@@ -309,5 +316,6 @@ def set_pipeline_args(p, model, prompts: list, negative_prompts: list, prompts_2
     if shared.cmd_opts.profile:
         t1 = time.time()
         shared.log.debug(f'Profile: pipeline args: {t1-t0:.2f}')
-    debug(f'Diffusers pipeline args: {args}')
+    if debug_enabled:
+        debug_log(f'Diffusers pipeline args: {args}')
     return args

modules/processing_callbacks.py

@@ -5,6 +5,7 @@ import torch
 import numpy as np
 from modules import shared, processing_correction, extra_networks, timer, prompt_parser_diffusers
 
+
 p = None
 debug = os.environ.get('SD_CALLBACK_DEBUG', None) is not None
 debug_callback = shared.log.trace if debug else lambda *args, **kwargs: None
@@ -14,6 +15,7 @@ def set_callbacks_p(processing):
     global p # pylint: disable=global-statement
     p = processing
 
+
 def prompt_callback(step, kwargs):
     if prompt_parser_diffusers.embedder is None or 'prompt_embeds' not in kwargs:
         return kwargs
@@ -28,6 +30,7 @@ def prompt_callback(step, kwargs):
         debug_callback(f"Callback: {e}")
     return kwargs
 
+
 def diffusers_callback_legacy(step: int, timestep: int, latents: typing.Union[torch.FloatTensor, np.ndarray]):
     if p is None:
         return
@@ -63,7 +66,7 @@ def diffusers_callback(pipe, step: int = 0, timestep: int = 0, kwargs: dict = {}
             if shared.state.interrupted or shared.state.skipped:
                 raise AssertionError('Interrupted...')
             time.sleep(0.1)
-    if hasattr(p, "stepwise_lora"):
+    if hasattr(p, "stepwise_lora") and shared.native:
         extra_networks.activate(p, p.extra_network_data, step=step)
     if latents is None:
         return kwargs

modules/processing_diffusers.py

@@ -8,6 +8,7 @@ from modules import shared, devices, processing, sd_models, errors, sd_hijack_hy
 from modules.processing_helpers import resize_hires, calculate_base_steps, calculate_hires_steps, calculate_refiner_steps, save_intermediate, update_sampler, is_txt2img, is_refiner_enabled
 from modules.processing_args import set_pipeline_args
 from modules.onnx_impl import preprocess_pipeline as preprocess_onnx_pipeline, check_parameters_changed as olive_check_parameters_changed
+from modules.lora import networks
 
 
 debug = shared.log.trace if os.environ.get('SD_DIFFUSERS_DEBUG', None) is not None else lambda *args, **kwargs: None
@@ -82,6 +83,8 @@ def process_base(p: processing.StableDiffusionProcessing):
     try:
         t0 = time.time()
         sd_models_compile.check_deepcache(enable=True)
+        if shared.opts.diffusers_offload_mode == "balanced":
+            shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model)
         sd_models.move_model(shared.sd_model, devices.device)
         if hasattr(shared.sd_model, 'unet'):
             sd_models.move_model(shared.sd_model.unet, devices.device)
@@ -198,11 +201,6 @@ def process_hires(p: processing.StableDiffusionProcessing, output):
                 if hasattr(shared.sd_model, "vae") and output.images is not None and len(output.images) > 0:
                     output.images = processing_vae.vae_decode(latents=output.images, model=shared.sd_model, full_quality=p.full_quality, output_type='pil', width=p.hr_upscale_to_x, height=p.hr_upscale_to_y) # controlnet cannnot deal with latent input
                     p.task_args['image'] = output.images # replace so hires uses new output
-            sd_models.move_model(shared.sd_model, devices.device)
-            if hasattr(shared.sd_model, 'unet'):
-                sd_models.move_model(shared.sd_model.unet, devices.device)
-            if hasattr(shared.sd_model, 'transformer'):
-                sd_models.move_model(shared.sd_model.transformer, devices.device)
             update_sampler(p, shared.sd_model, second_pass=True)
             orig_denoise = p.denoising_strength
             p.denoising_strength = strength
@@ -226,6 +224,11 @@ def process_hires(p: processing.StableDiffusionProcessing, output):
             shared.state.job = 'HiRes'
             shared.state.sampling_steps = hires_args.get('prior_num_inference_steps', None) or p.steps or hires_args.get('num_inference_steps', None)
             try:
+                sd_models.move_model(shared.sd_model, devices.device)
+                if hasattr(shared.sd_model, 'unet'):
+                    sd_models.move_model(shared.sd_model.unet, devices.device)
+                if hasattr(shared.sd_model, 'transformer'):
+                    sd_models.move_model(shared.sd_model.transformer, devices.device)
                 sd_models_compile.check_deepcache(enable=True)
                 output = shared.sd_model(**hires_args) # pylint: disable=not-callable
                 if isinstance(output, dict):
@@ -404,6 +407,9 @@ def process_diffusers(p: processing.StableDiffusionProcessing):
         shared.sd_model = orig_pipeline
         return results
 
+    if shared.opts.diffusers_offload_mode == "balanced":
+        shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model)
+
     # sanitize init_images
     if hasattr(p, 'init_images') and getattr(p, 'init_images', None) is None:
         del p.init_images
@@ -426,7 +432,6 @@ def process_diffusers(p: processing.StableDiffusionProcessing):
     if p.negative_prompts is None or len(p.negative_prompts) == 0:
         p.negative_prompts = p.all_negative_prompts[p.iteration * p.batch_size:(p.iteration+1) * p.batch_size]
 
-    sd_models.move_model(shared.sd_model, devices.device)
     sd_models_compile.openvino_recompile_model(p, hires=False, refiner=False) # recompile if a parameter changes
 
     if 'base' not in p.skip:
@@ -454,7 +459,13 @@ def process_diffusers(p: processing.StableDiffusionProcessing):
     results = process_decode(p, output)
 
     timer.process.record('decode')
+    timer.process.add('lora', networks.total_time())
+
     shared.sd_model = orig_pipeline
+
+    if shared.opts.diffusers_offload_mode == "balanced":
+        shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model)
+
     if p.state == '':
         global last_p # pylint: disable=global-statement
         last_p = p

modules/processing_helpers.py

@@ -1,4 +1,5 @@
 import os
+import time
 import math
 import random
 import warnings
@@ -9,7 +10,7 @@ import cv2
 from PIL import Image
 from skimage import exposure
 from blendmodes.blend import blendLayers, BlendType
-from modules import shared, devices, images, sd_models, sd_samplers, sd_hijack_hypertile, processing_vae
+from modules import shared, devices, images, sd_models, sd_samplers, sd_hijack_hypertile, processing_vae, timer
 
 
 debug = shared.log.trace if os.environ.get('SD_PROCESS_DEBUG', None) is not None else lambda *args, **kwargs: None
@@ -352,6 +353,7 @@ def img2img_image_conditioning(p, source_image, latent_image, image_mask=None):
 
 
 def validate_sample(tensor):
+    t0 = time.time()
     if not isinstance(tensor, np.ndarray) and not isinstance(tensor, torch.Tensor):
         return tensor
     dtype = tensor.dtype
@@ -366,17 +368,18 @@ def validate_sample(tensor):
     sample = 255.0 * np.moveaxis(sample, 0, 2) if not shared.native else 255.0 * sample
     with warnings.catch_warnings(record=True) as w:
         cast = sample.astype(np.uint8)
-    minimum, maximum, mean = np.min(cast), np.max(cast), np.mean(cast)
-    if len(w) > 0 or minimum == maximum:
+    if len(w) > 0:
         nans = np.isnan(sample).sum()
         cast = np.nan_to_num(sample)
         cast = cast.astype(np.uint8)
         vae = shared.sd_model.vae.dtype if hasattr(shared.sd_model, 'vae') else None
         upcast = getattr(shared.sd_model.vae.config, 'force_upcast', None) if hasattr(shared.sd_model, 'vae') and hasattr(shared.sd_model.vae, 'config') else None
-        shared.log.error(f'Decode: sample={sample.shape} invalid={nans} mean={mean} dtype={dtype} vae={vae} upcast={upcast} failed to validate')
+        shared.log.error(f'Decode: sample={sample.shape} invalid={nans} dtype={dtype} vae={vae} upcast={upcast} failed to validate')
         if upcast is not None and not upcast:
             setattr(shared.sd_model.vae.config, 'force_upcast', True) # noqa: B010
             shared.log.warning('Decode: upcast=True set, retry operation')
+    t1 = time.time()
+    timer.process.add('validate', t1 - t0)
     return cast
 
 

modules/processing_vae.py

@@ -117,7 +117,7 @@ def full_vae_decode(latents, model):
             model.vae.orig_dtype = model.vae.dtype
             model.vae = model.vae.to(dtype=torch.float32)
         latents = latents.to(torch.float32)
-    latents = latents.to(devices.device)
+    latents = latents.to(devices.device, non_blocking=True)
     if getattr(model.vae, "post_quant_conv", None) is not None:
         latents = latents.to(next(iter(model.vae.post_quant_conv.parameters())).dtype)
 

modules/prompt_parser_diffusers.py

@@ -16,6 +16,7 @@ orig_encode_token_ids_to_embeddings = EmbeddingsProvider._encode_token_ids_to_em
 token_dict = None # used by helper get_tokens
 token_type = None # used by helper get_tokens
 cache = OrderedDict()
+last_attention = None
 embedder = None
 
 
@@ -38,8 +39,8 @@ def prepare_model(pipe = None):
         pipe = pipe.pipe
     if not hasattr(pipe, "text_encoder"):
         return None
-    if shared.opts.diffusers_offload_mode == "balanced":
-        pipe = sd_models.apply_balanced_offload(pipe)
+    # if shared.opts.diffusers_offload_mode == "balanced":
+    #    pipe = sd_models.apply_balanced_offload(pipe)
     elif hasattr(pipe, "maybe_free_model_hooks"):
         pipe.maybe_free_model_hooks()
         devices.torch_gc()
@@ -52,7 +53,7 @@ class PromptEmbedder:
         self.prompts = prompts
         self.negative_prompts = negative_prompts
         self.batchsize = len(self.prompts)
-        self.attention = None
+        self.attention = last_attention
         self.allsame = self.compare_prompts()  # collapses batched prompts to single prompt if possible
         self.steps = steps
         self.clip_skip = clip_skip
@@ -78,6 +79,8 @@ class PromptEmbedder:
                 self.scheduled_encode(pipe, batchidx)
             else:
                 self.encode(pipe, prompt, negative_prompt, batchidx)
+        # if shared.opts.diffusers_offload_mode == "balanced":
+        #    pipe = sd_models.apply_balanced_offload(pipe)
         self.checkcache(p)
         debug(f"Prompt encode: time={(time.time() - t0):.3f}")
 
@@ -113,6 +116,7 @@ class PromptEmbedder:
                 debug(f"Prompt cache: add={key}")
                 while len(cache) > int(shared.opts.sd_textencoder_cache_size):
                     cache.popitem(last=False)
+                return True
         if item:
             self.__dict__.update(cache[key])
             cache.move_to_end(key)
@@ -161,7 +165,9 @@ class PromptEmbedder:
             self.negative_pooleds[batchidx].append(self.negative_pooleds[batchidx][idx])
 
     def encode(self, pipe, positive_prompt, negative_prompt, batchidx):
+        global last_attention # pylint: disable=global-statement
         self.attention = shared.opts.prompt_attention
+        last_attention = self.attention
         if self.attention == "xhinker":
             prompt_embed, positive_pooled, negative_embed, negative_pooled = get_xhinker_text_embeddings(pipe, positive_prompt, negative_prompt, self.clip_skip)
         else:
@@ -178,7 +184,6 @@ class PromptEmbedder:
         if debug_enabled:
             get_tokens(pipe, 'positive', positive_prompt)
             get_tokens(pipe, 'negative', negative_prompt)
-        pipe = prepare_model()
 
     def __call__(self, key, step=0):
         batch = getattr(self, key)
@@ -194,8 +199,6 @@ class PromptEmbedder:
 
 
 def compel_hijack(self, token_ids: torch.Tensor, attention_mask: typing.Optional[torch.Tensor] = None) -> torch.Tensor:
-    if not devices.same_device(self.text_encoder.device, devices.device):
-        sd_models.move_model(self.text_encoder, devices.device)
     needs_hidden_states = self.returned_embeddings_type != 1
     text_encoder_output = self.text_encoder(token_ids, attention_mask, output_hidden_states=needs_hidden_states, return_dict=True)
 
@@ -372,25 +375,31 @@ def prepare_embedding_providers(pipe, clip_skip) -> list[EmbeddingsProvider]:
         embedding_type = -(clip_skip + 1)
     else:
         embedding_type = clip_skip
+    embedding_args = {
+        'truncate': False,
+        'returned_embeddings_type': embedding_type,
+        'device': device,
+        'dtype_for_device_getter': lambda device: devices.dtype,
+    }
     if getattr(pipe, "prior_pipe", None) is not None and getattr(pipe.prior_pipe, "tokenizer", None) is not None and getattr(pipe.prior_pipe, "text_encoder", None) is not None:
-        provider = EmbeddingsProvider(padding_attention_mask_value=0, tokenizer=pipe.prior_pipe.tokenizer, text_encoder=pipe.prior_pipe.text_encoder, truncate=False, returned_embeddings_type=embedding_type, device=device)
+        provider = EmbeddingsProvider(padding_attention_mask_value=0, tokenizer=pipe.prior_pipe.tokenizer, text_encoder=pipe.prior_pipe.text_encoder, **embedding_args)
         embeddings_providers.append(provider)
-        no_mask_provider = EmbeddingsProvider(padding_attention_mask_value=1 if "sote" in pipe.sd_checkpoint_info.name.lower() else 0, tokenizer=pipe.prior_pipe.tokenizer, text_encoder=pipe.prior_pipe.text_encoder, truncate=False, returned_embeddings_type=embedding_type, device=device)
+        no_mask_provider = EmbeddingsProvider(padding_attention_mask_value=1 if "sote" in pipe.sd_checkpoint_info.name.lower() else 0, tokenizer=pipe.prior_pipe.tokenizer, text_encoder=pipe.prior_pipe.text_encoder, **embedding_args)
         embeddings_providers.append(no_mask_provider)
     elif getattr(pipe, "tokenizer", None) is not None and getattr(pipe, "text_encoder", None) is not None:
-        if not devices.same_device(pipe.text_encoder.device, devices.device):
-            sd_models.move_model(pipe.text_encoder, devices.device)
-        provider = EmbeddingsProvider(tokenizer=pipe.tokenizer, text_encoder=pipe.text_encoder, truncate=False, returned_embeddings_type=embedding_type, device=device)
+        if pipe.text_encoder.__class__.__name__.startswith('CLIP'):
+            sd_models.move_model(pipe.text_encoder, devices.device, force=True)
+        provider = EmbeddingsProvider(tokenizer=pipe.tokenizer, text_encoder=pipe.text_encoder, **embedding_args)
         embeddings_providers.append(provider)
     if getattr(pipe, "tokenizer_2", None) is not None and getattr(pipe, "text_encoder_2", None) is not None:
-        if not devices.same_device(pipe.text_encoder_2.device, devices.device):
-            sd_models.move_model(pipe.text_encoder_2, devices.device)
-        provider = EmbeddingsProvider(tokenizer=pipe.tokenizer_2, text_encoder=pipe.text_encoder_2, truncate=False, returned_embeddings_type=embedding_type, device=device)
+        if pipe.text_encoder_2.__class__.__name__.startswith('CLIP'):
+            sd_models.move_model(pipe.text_encoder_2, devices.device, force=True)
+        provider = EmbeddingsProvider(tokenizer=pipe.tokenizer_2, text_encoder=pipe.text_encoder_2, **embedding_args)
         embeddings_providers.append(provider)
     if getattr(pipe, "tokenizer_3", None) is not None and getattr(pipe, "text_encoder_3", None) is not None:
-        if not devices.same_device(pipe.text_encoder_3.device, devices.device):
-            sd_models.move_model(pipe.text_encoder_3, devices.device)
-        provider = EmbeddingsProvider(tokenizer=pipe.tokenizer_3, text_encoder=pipe.text_encoder_3, truncate=False, returned_embeddings_type=embedding_type, device=device)
+        if pipe.text_encoder_3.__class__.__name__.startswith('CLIP'):
+            sd_models.move_model(pipe.text_encoder_3, devices.device, force=True)
+        provider = EmbeddingsProvider(tokenizer=pipe.tokenizer_3, text_encoder=pipe.text_encoder_3, **embedding_args)
         embeddings_providers.append(provider)
     return embeddings_providers
 

modules/schedulers/scheduler_bdia.py

@@ -0,0 +1,551 @@
+# Copyright 2024 Stanford University Team and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# DISCLAIMER: This code is strongly influenced by https://github.com/pesser/pytorch_diffusion
+# and https://github.com/hojonathanho/diffusion
+
+import math
+from dataclasses import dataclass
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.utils import BaseOutput
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.schedulers.scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin
+
+
+@dataclass
+# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->DDIM
+class DDIMSchedulerOutput(BaseOutput):
+    """
+    Output class for the scheduler's `step` function output.
+
+    Args:
+        prev_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images):
+            Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the
+            denoising loop.
+        pred_original_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images):
+            The predicted denoised sample `(x_{0})` based on the model output from the current timestep.
+            `pred_original_sample` can be used to preview progress or for guidance.
+    """
+
+    prev_sample: torch.Tensor
+    pred_original_sample: Optional[torch.Tensor] = None
+
+
+# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar
+def betas_for_alpha_bar(
+    num_diffusion_timesteps,
+    max_beta=0.999,
+    alpha_transform_type="cosine",
+):
+    """
+    Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of
+    (1-beta) over time from t = [0,1].
+
+    Contains a function alpha_bar that takes an argument t and transforms it to the cumulative product of (1-beta) up
+    to that part of the diffusion process.
+
+
+    Args:
+        num_diffusion_timesteps (`int`): the number of betas to produce.
+        max_beta (`float`): the maximum beta to use; use values lower than 1 to
+                     prevent singularities.
+        alpha_transform_type (`str`, *optional*, default to `cosine`): the type of noise schedule for alpha_bar.
+                     Choose from `cosine` or `exp`
+
+    Returns:
+        betas (`np.ndarray`): the betas used by the scheduler to step the model outputs
+    """
+    if alpha_transform_type == "cosine":
+
+        def alpha_bar_fn(t):
+            return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2
+
+    elif alpha_transform_type == "exp":
+
+        def alpha_bar_fn(t):
+            return math.exp(t * -12.0)
+
+    else:
+        raise ValueError(f"Unsupported alpha_transform_type: {alpha_transform_type}")
+
+    betas = []
+    for i in range(num_diffusion_timesteps):
+        t1 = i / num_diffusion_timesteps
+        t2 = (i + 1) / num_diffusion_timesteps
+        betas.append(min(1 - alpha_bar_fn(t2) / alpha_bar_fn(t1), max_beta))
+    return torch.tensor(betas, dtype=torch.float32)
+
+
+def rescale_zero_terminal_snr(betas):
+    """
+    Rescales betas to have zero terminal SNR Based on https://arxiv.org/pdf/2305.08891.pdf (Algorithm 1)
+
+
+    Args:
+        betas (`torch.Tensor`):
+            the betas that the scheduler is being initialized with.
+
+    Returns:
+        `torch.Tensor`: rescaled betas with zero terminal SNR
+    """
+    # Convert betas to alphas_bar_sqrt
+    alphas = 1.0 - betas
+    alphas_cumprod = torch.cumprod(alphas, dim=0)
+    alphas_bar_sqrt = alphas_cumprod.sqrt()
+
+    # Store old values.
+    alphas_bar_sqrt_0 = alphas_bar_sqrt[0].clone()
+    alphas_bar_sqrt_T = alphas_bar_sqrt[-1].clone()
+
+    # Shift so the last timestep is zero.
+    alphas_bar_sqrt -= alphas_bar_sqrt_T
+
+    # Scale so the first timestep is back to the old value.
+    alphas_bar_sqrt *= alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 - alphas_bar_sqrt_T)
+
+    # Convert alphas_bar_sqrt to betas
+    alphas_bar = alphas_bar_sqrt**2  # Revert sqrt
+    alphas = alphas_bar[1:] / alphas_bar[:-1]  # Revert cumprod
+    alphas = torch.cat([alphas_bar[0:1], alphas])
+    betas = 1 - alphas
+
+    return betas
+
+class BDIA_DDIMScheduler(SchedulerMixin, ConfigMixin):
+    """
+    `DDIMScheduler` extends the denoising procedure introduced in denoising diffusion probabilistic models (DDPMs) with
+    non-Markovian guidance.
+
+    This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic
+    methods the library implements for all schedulers such as loading and saving.
+
+    Args:
+        num_train_timesteps (`int`, defaults to 1000):
+            The number of diffusion steps to train the model.
+        beta_start (`float`, defaults to 0.0001):
+            The starting `beta` value of inference.
+        beta_end (`float`, defaults to 0.02):
+            The final `beta` value.
+        beta_schedule (`str`, defaults to `"linear"`):
+            The beta schedule, a mapping from a beta range to a sequence of betas for stepping the model. Choose from
+            `linear`, `scaled_linear`, or `squaredcos_cap_v2`.
+        trained_betas (`np.ndarray`, *optional*):
+            Pass an array of betas directly to the constructor to bypass `beta_start` and `beta_end`.
+        clip_sample (`bool`, defaults to `True`):
+            Clip the predicted sample for numerical stability.
+        clip_sample_range (`float`, defaults to 1.0):
+            The maximum magnitude for sample clipping. Valid only when `clip_sample=True`.
+        set_alpha_to_one (`bool`, defaults to `True`):
+            Each diffusion step uses the alphas product value at that step and at the previous one. For the final step
+            there is no previous alpha. When this option is `True` the previous alpha product is fixed to `1`,
+            otherwise it uses the alpha value at step 0.
+        steps_offset (`int`, defaults to 0):
+            An offset added to the inference steps, as required by some model families.
+        prediction_type (`str`, defaults to `epsilon`, *optional*):
+            Prediction type of the scheduler function; can be `epsilon` (predicts the noise of the diffusion process),
+            `sample` (directly predicts the noisy sample`) or `v_prediction` (see section 2.4 of [Imagen
+            Video](https://imagen.research.google/video/paper.pdf) paper).
+        thresholding (`bool`, defaults to `False`):
+            Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such
+            as Stable Diffusion.
+        dynamic_thresholding_ratio (`float`, defaults to 0.995):
+            The ratio for the dynamic thresholding method. Valid only when `thresholding=True`.
+        sample_max_value (`float`, defaults to 1.0):
+            The threshold value for dynamic thresholding. Valid only when `thresholding=True`.
+        timestep_spacing (`str`, defaults to `"leading"`):
+            The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and
+            Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information.
+        rescale_betas_zero_snr (`bool`, defaults to `False`):
+            Whether to rescale the betas to have zero terminal SNR. This enables the model to generate very bright and
+            dark samples instead of limiting it to samples with medium brightness. Loosely related to
+            [`--offset_noise`](https://github.com/huggingface/diffusers/blob/74fd735eb073eb1d774b1ab4154a0876eb82f055/examples/dreambooth/train_dreambooth.py#L506).
+    """
+
+    _compatibles = [e.name for e in KarrasDiffusionSchedulers]
+    order = 1
+
+    @register_to_config
+    def __init__(
+        self,
+        num_train_timesteps: int = 1000,
+        beta_start: float = 0.0001,
+        beta_end: float = 0.02,
+        beta_schedule: str = "linear",
+        trained_betas: Optional[Union[np.ndarray, List[float]]] = None,
+        clip_sample: bool = True,
+        set_alpha_to_one: bool = True, #was True
+        steps_offset: int = 0,
+        prediction_type: str = "epsilon",
+        thresholding: bool = False,
+        dynamic_thresholding_ratio: float = 0.995,
+        clip_sample_range: float = 1.0,
+        sample_max_value: float = 1.0,
+        timestep_spacing: str = "leading", #leading
+        rescale_betas_zero_snr: bool = False,
+        gamma: float = 1.0,
+
+    ):
+        if trained_betas is not None:
+            self.betas = torch.tensor(trained_betas, dtype=torch.float32)
+        elif beta_schedule == "linear":
+            self.betas = torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32)
+        elif beta_schedule == "scaled_linear":
+            # this schedule is very specific to the latent diffusion model.
+            self.betas = torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float32) ** 2
+        elif beta_schedule == "squaredcos_cap_v2":
+            # Glide cosine schedule
+            self.betas = betas_for_alpha_bar(num_train_timesteps)
+        else:
+            raise NotImplementedError(f"{beta_schedule} is not implemented for {self.__class__}")
+
+        # Rescale for zero SNR
+        if rescale_betas_zero_snr:
+            self.betas = rescale_zero_terminal_snr(self.betas)
+
+        self.alphas = 1.0 - self.betas #may have to add something for last step
+        
+        self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
+        # At every step in ddim, we are looking into the previous alphas_cumprod
+        # For the final step, there is no previous alphas_cumprod because we are already at 0
+        # `set_alpha_to_one` decides whether we set this parameter simply to one or
+        # whether we use the final alpha of the "non-previous" one.
+        self.final_alpha_cumprod = torch.tensor(1.0) if set_alpha_to_one else self.alphas_cumprod[0]
+
+        # standard deviation of the initial noise distribution
+        self.init_noise_sigma = 1.0
+
+        # setable values
+        self.num_inference_steps = None
+        self.timesteps = torch.from_numpy(np.arange(0, num_train_timesteps)[::-1].copy().astype(np.int64))
+        self.next_sample = []
+        self.BDIA = False
+
+
+    def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor:
+        """
+        Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
+        current timestep.
+
+        Args:
+            sample (`torch.Tensor`):
+                The input sample.
+            timestep (`int`, *optional*):
+                The current timestep in the diffusion chain.
+
+        Returns:
+            `torch.Tensor`:
+                A scaled input sample.
+        """
+        return sample
+
+    def _get_variance(self, timestep, prev_timestep):
+        alpha_prod_t = self.alphas_cumprod[timestep]
+        alpha_prod_t_prev = self.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.final_alpha_cumprod
+        beta_prod_t = 1 - alpha_prod_t
+        beta_prod_t_prev = 1 - alpha_prod_t_prev
+
+        variance = (beta_prod_t_prev / beta_prod_t) * (1 - alpha_prod_t / alpha_prod_t_prev)
+
+        return variance
+
+    # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample
+    def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor:
+        """
+        "Dynamic thresholding: At each sampling step we set s to a certain percentile absolute pixel value in xt0 (the
+        prediction of x_0 at timestep t), and if s > 1, then we threshold xt0 to the range [-s, s] and then divide by
+        s. Dynamic thresholding pushes saturated pixels (those near -1 and 1) inwards, thereby actively preventing
+        pixels from saturation at each step. We find that dynamic thresholding results in significantly better
+        photorealism as well as better image-text alignment, especially when using very large guidance weights."
+
+        https://arxiv.org/abs/2205.11487
+        """
+        dtype = sample.dtype
+        batch_size, channels, *remaining_dims = sample.shape
+
+        if dtype not in (torch.float32, torch.float64):
+            sample = sample.float()  # upcast for quantile calculation, and clamp not implemented for cpu half
+
+        # Flatten sample for doing quantile calculation along each image
+        sample = sample.reshape(batch_size, channels * np.prod(remaining_dims))
+
+        abs_sample = sample.abs()  # "a certain percentile absolute pixel value"
+
+        s = torch.quantile(abs_sample, self.config.dynamic_thresholding_ratio, dim=1)
+        s = torch.clamp(
+            s, min=1, max=self.config.sample_max_value
+        )  # When clamped to min=1, equivalent to standard clipping to [-1, 1]
+        s = s.unsqueeze(1)  # (batch_size, 1) because clamp will broadcast along dim=0
+        sample = torch.clamp(sample, -s, s) / s  # "we threshold xt0 to the range [-s, s] and then divide by s"
+
+        sample = sample.reshape(batch_size, channels, *remaining_dims)
+        sample = sample.to(dtype)
+
+        return sample
+
+    def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device] = None):
+        """
+        Sets the discrete timesteps used for the diffusion chain (to be run before inference).
+
+        Args:
+            num_inference_steps (`int`):
+                The number of diffusion steps used when generating samples with a pre-trained model.
+        """
+
+        if num_inference_steps > self.config.num_train_timesteps:
+            raise ValueError(
+                f"`num_inference_steps`: {num_inference_steps} cannot be larger than `self.config.train_timesteps`:"
+                f" {self.config.num_train_timesteps} as the unet model trained with this scheduler can only handle"
+                f" maximal {self.config.num_train_timesteps} timesteps."
+            )
+
+        self.num_inference_steps = num_inference_steps
+
+        # "linspace", "leading", "trailing" corresponds to annotation of Table 2. of https://arxiv.org/abs/2305.08891
+        if self.config.timestep_spacing == "linspace":
+            timesteps = (
+                np.linspace(0, self.config.num_train_timesteps - 1, num_inference_steps)
+                .round()[::-1]
+                .copy()
+                .astype(np.int64)
+            )
+        elif self.config.timestep_spacing == "leading":
+            step_ratio = self.config.num_train_timesteps // self.num_inference_steps
+            # creates integer timesteps by multiplying by ratio
+            # casting to int to avoid issues when num_inference_step is power of 3
+            timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(np.int64)
+            timesteps += self.config.steps_offset
+        elif self.config.timestep_spacing == "trailing":
+            step_ratio = self.config.num_train_timesteps / self.num_inference_steps
+            # creates integer timesteps by multiplying by ratio
+            # casting to int to avoid issues when num_inference_step is power of 3
+            timesteps = np.round(np.arange(self.config.num_train_timesteps, 0, -step_ratio)).astype(np.int64)
+            timesteps -= 1
+        else:
+            raise ValueError(
+                f"{self.config.timestep_spacing} is not supported. Please make sure to choose one of 'leading' or 'trailing'."
+            )
+
+        self.timesteps = torch.from_numpy(timesteps).to(device)
+
+    def step(
+        self,
+        model_output: torch.Tensor,
+        timestep: int,
+        sample: torch.Tensor,
+        eta: float = 0.0,
+        use_clipped_model_output: bool = False,
+        generator=None,
+        variance_noise: Optional[torch.Tensor] = None,
+        return_dict: bool = True,
+        debug: bool = False,
+    ) -> Union[DDIMSchedulerOutput, Tuple]:
+        """
+        Predict the sample from the previous timestep by reversing the SDE.
+        
+        Args:
+            model_output (torch.Tensor): Direct output from learned diffusion model
+            timestep (int): Current discrete timestep in the diffusion chain
+            sample (torch.Tensor): Current instance of sample created by diffusion process
+            eta (float): Weight of noise for added noise in diffusion step
+            use_clipped_model_output (bool): Whether to use clipped model output
+            generator (torch.Generator, optional): Random number generator
+            variance_noise (torch.Tensor, optional): Pre-generated noise for variance
+            return_dict (bool): Whether to return as DDIMSchedulerOutput or tuple
+            debug (bool): Whether to print debug information
+        """
+        if self.num_inference_steps is None:
+            raise ValueError("Number of inference steps is 'None', run 'set_timesteps' first")
+
+        # Calculate timesteps
+        step_size = self.config.num_train_timesteps // self.num_inference_steps
+        prev_timestep = timestep - step_size
+        next_timestep = timestep + step_size
+
+        if debug:
+            print("\n=== Timestep Information ===")
+            print(f"Current timestep: {timestep}")
+            print(f"Previous timestep: {prev_timestep}")
+            print(f"Next timestep: {next_timestep}")
+            print(f"Step size: {step_size}")
+
+        # Pre-compute alpha and variance values
+        alpha_prod_t = self.alphas_cumprod[timestep]
+        alpha_prod_t_prev = self.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.final_alpha_cumprod
+        variance = self._get_variance(timestep, prev_timestep)
+        std_dev_t = eta * variance ** 0.5
+
+        # Compute required values
+        alpha_i = alpha_prod_t ** 0.5
+        alpha_i_minus_1 = alpha_prod_t_prev ** 0.5
+        sigma_i = (1 - alpha_prod_t) ** 0.5
+        sigma_i_minus_1 = (1 - alpha_prod_t_prev - std_dev_t**2) ** 0.5
+
+        if debug:
+            print("\n=== Alpha Values ===")
+            print(f"alpha_i: {alpha_i}")
+            print(f"alpha_i_minus_1: {alpha_i_minus_1}")
+            print(f"sigma_i: {sigma_i}")
+            print(f"sigma_i_minus_1: {sigma_i_minus_1}")
+
+        # Predict original sample based on prediction type
+        if self.config.prediction_type == "epsilon":
+            pred_original_sample = (sample - sigma_i * model_output) / alpha_i
+            pred_epsilon = model_output
+            if debug:
+                print("\nPrediction type: epsilon")
+        elif self.config.prediction_type == "sample":
+            pred_original_sample = model_output
+            pred_epsilon = (sample - alpha_i * pred_original_sample) / sigma_i
+            if debug:
+                print("\nPrediction type: sample")
+        elif self.config.prediction_type == "v_prediction":
+            pred_original_sample = alpha_i * sample - sigma_i * model_output
+            pred_epsilon = alpha_i * model_output + sigma_i * sample
+            if debug:
+                print("\nPrediction type: v_prediction")
+        else:
+            raise ValueError(
+                f"prediction_type {self.config.prediction_type} must be one of `epsilon`, `sample`, or `v_prediction`"
+            )
+
+        # Apply thresholding or clipping if configured
+        if self.config.thresholding:
+            if debug:
+                print("\nApplying thresholding")
+            pred_original_sample = self._threshold_sample(pred_original_sample)
+        elif self.config.clip_sample:
+            if debug:
+                print("\nApplying clipping")
+            pred_original_sample = pred_original_sample.clamp(
+                -self.config.clip_sample_range, self.config.clip_sample_range
+            )
+
+        # Recompute pred_epsilon if using clipped model output
+        if use_clipped_model_output:
+            if debug:
+                print("\nUsing clipped model output")
+            pred_epsilon = (sample - alpha_i * pred_original_sample) / sigma_i
+
+        # Compute DDIM step
+        ddim_step = alpha_i_minus_1 * pred_original_sample + sigma_i_minus_1 * pred_epsilon
+
+        # Handle initial DDIM step or BDIA steps
+        if len(self.next_sample) == 0:
+            if debug:
+                print("\nFirst iteration (DDIM)")
+            self.update_next_sample_BDIA(sample)
+            self.update_next_sample_BDIA(ddim_step)
+        else:
+            if debug:
+                print("\nBDIA step")
+            # BDIA implementation
+            alpha_prod_t_next = self.alphas_cumprod[next_timestep]
+            alpha_i_plus_1 = alpha_prod_t_next ** 0.5
+            sigma_i_plus_1 = (1 - alpha_prod_t_next) ** 0.5
+            
+            if debug:
+                print(f"alpha_i_plus_1: {alpha_i_plus_1}")
+                print(f"sigma_i_plus_1: {sigma_i_plus_1}")
+            
+            a = alpha_i_plus_1 * pred_original_sample + sigma_i_plus_1 * pred_epsilon
+            bdia_step = (
+                self.config.gamma * self.next_sample[-2] + 
+                ddim_step - 
+                (self.config.gamma * a)
+            )
+            self.update_next_sample_BDIA(bdia_step)
+
+        prev_sample = self.next_sample[-1]
+
+        # Apply variance noise if eta > 0
+        if eta > 0:
+            if debug:
+                print(f"\nApplying variance noise with eta: {eta}")
+            
+            if variance_noise is not None and generator is not None:
+                raise ValueError(
+                    "Cannot pass both generator and variance_noise. Use either `generator` or `variance_noise`."
+                )
+
+            if variance_noise is None:
+                variance_noise = randn_tensor(
+                    model_output.shape,
+                    generator=generator,
+                    device=model_output.device,
+                    dtype=model_output.dtype
+                )
+            prev_sample = prev_sample + std_dev_t * variance_noise
+
+        if not return_dict:
+            return (prev_sample,)
+
+        return DDIMSchedulerOutput(prev_sample=prev_sample, pred_original_sample=pred_original_sample)
+    
+    def add_noise(
+        self,
+        original_samples: torch.Tensor,
+        noise: torch.Tensor,
+        timesteps: torch.IntTensor,
+    ) -> torch.Tensor:
+        # Make sure alphas_cumprod and timestep have same device and dtype as original_samples
+        # Move the self.alphas_cumprod to device to avoid redundant CPU to GPU data movement
+        # for the subsequent add_noise calls
+        self.alphas_cumprod = self.alphas_cumprod.to(device=original_samples.device)
+        alphas_cumprod = self.alphas_cumprod.to(dtype=original_samples.dtype)
+        timesteps = timesteps.to(original_samples.device)
+
+        sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5
+        sqrt_alpha_prod = sqrt_alpha_prod.flatten()
+        while len(sqrt_alpha_prod.shape) < len(original_samples.shape):
+            sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1)
+
+        sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5
+        sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten()
+        while len(sqrt_one_minus_alpha_prod.shape) < len(original_samples.shape):
+            sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1)
+
+        noisy_samples = sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise
+        return noisy_samples
+
+    # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.get_velocity
+    def get_velocity(self, sample: torch.Tensor, noise: torch.Tensor, timesteps: torch.IntTensor) -> torch.Tensor:
+        # Make sure alphas_cumprod and timestep have same device and dtype as sample
+        self.alphas_cumprod = self.alphas_cumprod.to(device=sample.device)
+        alphas_cumprod = self.alphas_cumprod.to(dtype=sample.dtype)
+        timesteps = timesteps.to(sample.device)
+
+        sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5
+        sqrt_alpha_prod = sqrt_alpha_prod.flatten()
+        while len(sqrt_alpha_prod.shape) < len(sample.shape):
+            sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1)
+
+        sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5
+        sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten()
+        while len(sqrt_one_minus_alpha_prod.shape) < len(sample.shape):
+            sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1)
+
+        velocity = sqrt_alpha_prod * noise - sqrt_one_minus_alpha_prod * sample
+        return velocity
+    
+    def update_next_sample_BDIA(self, new_value):
+        self.next_sample.append(new_value.clone())
+
+
+    def __len__(self):
+        return self.config.num_train_timesteps

modules/sd_checkpoint.py

@@ -123,13 +123,17 @@ def list_models():
     checkpoint_aliases.clear()
     ext_filter = [".safetensors"] if shared.opts.sd_disable_ckpt or shared.native else [".ckpt", ".safetensors"]
     model_list = list(modelloader.load_models(model_path=model_path, model_url=None, command_path=shared.opts.ckpt_dir, ext_filter=ext_filter, download_name=None, ext_blacklist=[".vae.ckpt", ".vae.safetensors"]))
+    safetensors_list = []
     for filename in sorted(model_list, key=str.lower):
         checkpoint_info = CheckpointInfo(filename)
+        safetensors_list.append(checkpoint_info)
         if checkpoint_info.name is not None:
             checkpoint_info.register()
+    diffusers_list = []
     if shared.native:
         for repo in modelloader.load_diffusers_models(clear=True):
             checkpoint_info = CheckpointInfo(repo['name'], sha=repo['hash'])
+            diffusers_list.append(checkpoint_info)
             if checkpoint_info.name is not None:
                 checkpoint_info.register()
     if shared.cmd_opts.ckpt is not None:
@@ -143,7 +147,7 @@ def list_models():
                 shared.opts.data['sd_model_checkpoint'] = checkpoint_info.title
     elif shared.cmd_opts.ckpt != shared.default_sd_model_file and shared.cmd_opts.ckpt is not None:
         shared.log.warning(f'Load model: path="{shared.cmd_opts.ckpt}" not found')
-    shared.log.info(f'Available Models: path="{shared.opts.ckpt_dir}" items={len(checkpoints_list)} time={time.time()-t0:.2f}')
+    shared.log.info(f'Available Models: items={len(checkpoints_list)} safetensors="{shared.opts.ckpt_dir}":{len(safetensors_list)} diffusers="{shared.opts.diffusers_dir}":{len(diffusers_list)} time={time.time()-t0:.2f}')
     checkpoints_list = dict(sorted(checkpoints_list.items(), key=lambda cp: cp[1].filename))
 
 def update_model_hashes():
@@ -249,6 +253,8 @@ def select_checkpoint(op='model'):
         model_checkpoint = shared.opts.data.get('sd_model_refiner', None)
     else:
         model_checkpoint = shared.opts.sd_model_checkpoint
+    if len(model_checkpoint) < 3:
+        return None
     if model_checkpoint is None or model_checkpoint == 'None':
         return None
     checkpoint_info = get_closet_checkpoint_match(model_checkpoint)
@@ -275,6 +281,12 @@ def select_checkpoint(op='model'):
     return checkpoint_info
 
 
+def init_metadata():
+    global sd_metadata # pylint: disable=global-statement
+    if sd_metadata is None:
+        sd_metadata = shared.readfile(sd_metadata_file, lock=True) if os.path.isfile(sd_metadata_file) else {}
+
+
 def read_metadata_from_safetensors(filename):
     global sd_metadata # pylint: disable=global-statement
     if sd_metadata is None:

modules/sd_models.py

@@ -13,6 +13,7 @@ import diffusers.loaders.single_file_utils
 from rich import progress # pylint: disable=redefined-builtin
 import torch
 import safetensors.torch
+import accelerate
 from omegaconf import OmegaConf
 from ldm.util import instantiate_from_config
 from modules import paths, shared, shared_state, modelloader, devices, script_callbacks, sd_vae, sd_unet, errors, sd_models_config, sd_models_compile, sd_hijack_accelerate, sd_detect
@@ -310,6 +311,7 @@ def set_accelerate(sd_model):
 
 
 def set_diffuser_offload(sd_model, op: str = 'model'):
+    t0 = time.time()
     if not shared.native:
         shared.log.warning('Attempting to use offload with backend=original')
         return
@@ -359,73 +361,106 @@ def set_diffuser_offload(sd_model, op: str = 'model'):
             shared.log.error(f'Setting {op}: offload={shared.opts.diffusers_offload_mode} {e}')
     if shared.opts.diffusers_offload_mode == "balanced":
         try:
-            shared.log.debug(f'Setting {op}: offload={shared.opts.diffusers_offload_mode} threshold={shared.opts.diffusers_offload_max_gpu_memory} limit={shared.opts.cuda_mem_fraction}')
+            shared.log.debug(f'Setting {op}: offload={shared.opts.diffusers_offload_mode} watermarks low={shared.opts.diffusers_offload_min_gpu_memory} high={shared.opts.diffusers_offload_max_gpu_memory} limit={shared.opts.cuda_mem_fraction:.2f}')
             sd_model = apply_balanced_offload(sd_model)
         except Exception as e:
             shared.log.error(f'Setting {op}: offload={shared.opts.diffusers_offload_mode} {e}')
+    process_timer.add('offload', time.time() - t0)
+
+
+class OffloadHook(accelerate.hooks.ModelHook):
+    def __init__(self):
+        if shared.opts.diffusers_offload_max_gpu_memory > 1:
+            shared.opts.diffusers_offload_max_gpu_memory = 0.75
+        if shared.opts.diffusers_offload_max_cpu_memory > 1:
+            shared.opts.diffusers_offload_max_cpu_memory = 0.75
+        self.min_watermark = shared.opts.diffusers_offload_min_gpu_memory
+        self.max_watermark = shared.opts.diffusers_offload_max_gpu_memory
+        self.cpu_watermark = shared.opts.diffusers_offload_max_cpu_memory
+        self.gpu = int(shared.gpu_memory * shared.opts.diffusers_offload_max_gpu_memory * 1024*1024*1024)
+        self.cpu = int(shared.cpu_memory * shared.opts.diffusers_offload_max_cpu_memory * 1024*1024*1024)
+        gpu_dict = { "min": self.min_watermark, "max": self.max_watermark, "bytes": self.gpu }
+        cpu_dict = { "max": self.cpu_watermark, "bytes": self.cpu }
+        shared.log.info(f'Init offload: type=balanced gpu={gpu_dict} cpu={cpu_dict}')
+        super().__init__()
+
+    def init_hook(self, module):
+        return module
+
+    def pre_forward(self, module, *args, **kwargs):
+        if devices.normalize_device(module.device) != devices.normalize_device(devices.device):
+            device_index = torch.device(devices.device).index
+            if device_index is None:
+                device_index = 0
+            max_memory = { device_index: self.gpu, "cpu": self.cpu }
+            device_map = getattr(module, "balanced_offload_device_map", None)
+            if device_map is None or max_memory != getattr(module, "balanced_offload_max_memory", None):
+                device_map = accelerate.infer_auto_device_map(module, max_memory=max_memory)
+            offload_dir = getattr(module, "offload_dir", os.path.join(shared.opts.accelerate_offload_path, module.__class__.__name__))
+            module = accelerate.dispatch_model(module, device_map=device_map, offload_dir=offload_dir)
+            module._hf_hook.execution_device = torch.device(devices.device) # pylint: disable=protected-access
+            module.balanced_offload_device_map = device_map
+            module.balanced_offload_max_memory = max_memory
+        return args, kwargs
+
+    def post_forward(self, module, output):
+        return output
+
+    def detach_hook(self, module):
+        return module
+
+
+offload_hook_instance = None
 
 
 def apply_balanced_offload(sd_model):
-    from accelerate import infer_auto_device_map, dispatch_model
-    from accelerate.hooks import add_hook_to_module, remove_hook_from_module, ModelHook
+    global offload_hook_instance # pylint: disable=global-statement
+    if offload_hook_instance is None or offload_hook_instance.min_watermark != shared.opts.diffusers_offload_min_gpu_memory or offload_hook_instance.max_watermark != shared.opts.diffusers_offload_max_gpu_memory:
+        offload_hook_instance = OffloadHook()
+    t0 = time.time()
     excluded = ['OmniGenPipeline']
     if sd_model.__class__.__name__ in excluded:
         return sd_model
-
-    class dispatch_from_cpu_hook(ModelHook):
-        def init_hook(self, module):
-            return module
-
-        def pre_forward(self, module, *args, **kwargs):
-            if devices.normalize_device(module.device) != devices.normalize_device(devices.device):
-                device_index = torch.device(devices.device).index
-                if device_index is None:
-                    device_index = 0
-                max_memory = {
-                    device_index: f"{shared.opts.diffusers_offload_max_gpu_memory}GiB",
-                    "cpu": f"{shared.opts.diffusers_offload_max_cpu_memory}GiB",
-                }
-                device_map = infer_auto_device_map(module, max_memory=max_memory)
-                module = remove_hook_from_module(module, recurse=True)
-                offload_dir = getattr(module, "offload_dir", os.path.join(shared.opts.accelerate_offload_path, module.__class__.__name__))
-                module = dispatch_model(module, device_map=device_map, offload_dir=offload_dir)
-                module = add_hook_to_module(module, dispatch_from_cpu_hook(), append=True)
-                module._hf_hook.execution_device = torch.device(devices.device) # pylint: disable=protected-access
-            return args, kwargs
-
-        def post_forward(self, module, output):
-            return output
-
-        def detach_hook(self, module):
-            return module
+    fn = f'{sys._getframe(2).f_code.co_name}:{sys._getframe(1).f_code.co_name}' # pylint: disable=protected-access
+    debug_move(f'Apply offload: type=balanced fn={fn}')
+    checkpoint_name = sd_model.sd_checkpoint_info.name if getattr(sd_model, "sd_checkpoint_info", None) is not None else None
+    if checkpoint_name is None:
+        checkpoint_name = sd_model.__class__.__name__
 
     def apply_balanced_offload_to_module(pipe):
+        used_gpu, used_ram = devices.torch_gc(fast=True)
         if hasattr(pipe, "pipe"):
             apply_balanced_offload_to_module(pipe.pipe)
         if hasattr(pipe, "_internal_dict"):
             keys = pipe._internal_dict.keys() # pylint: disable=protected-access
         else:
-            keys = get_signature(shared.sd_model).keys()
+            keys = get_signature(pipe).keys()
         for module_name in keys: # pylint: disable=protected-access
             module = getattr(pipe, module_name, None)
             if isinstance(module, torch.nn.Module):
-                checkpoint_name = pipe.sd_checkpoint_info.name if getattr(pipe, "sd_checkpoint_info", None) is not None else None
-                if checkpoint_name is None:
-                    checkpoint_name = pipe.__class__.__name__
-                offload_dir = os.path.join(shared.opts.accelerate_offload_path, checkpoint_name, module_name)
-                module = remove_hook_from_module(module, recurse=True)
+                network_layer_name = getattr(module, "network_layer_name", None)
+                device_map = getattr(module, "balanced_offload_device_map", None)
+                max_memory = getattr(module, "balanced_offload_max_memory", None)
+                module = accelerate.hooks.remove_hook_from_module(module, recurse=True)
                 try:
-                    module = module.to("cpu")
-                    module.offload_dir = offload_dir
-                    network_layer_name = getattr(module, "network_layer_name", None)
-                    module = add_hook_to_module(module, dispatch_from_cpu_hook(), append=True)
-                    module._hf_hook.execution_device = torch.device(devices.device) # pylint: disable=protected-access
-                    if network_layer_name:
-                        module.network_layer_name = network_layer_name
+                    do_offload = used_gpu > 100 * shared.opts.diffusers_offload_min_gpu_memory
+                    debug_move(f'Balanced offload: gpu={used_gpu} ram={used_ram} current={module.device} dtype={module.dtype} op={"move" if do_offload else "skip"} component={module.__class__.__name__}')
+                    if do_offload:
+                        module = module.to(devices.cpu)
+                        used_gpu, used_ram = devices.torch_gc(fast=True, force=True)
                 except Exception as e:
                     if 'bitsandbytes' not in str(e):
                         shared.log.error(f'Balanced offload: module={module_name} {e}')
-                devices.torch_gc(fast=True)
+                    if os.environ.get('SD_MOVE_DEBUG', None):
+                        errors.display(e, f'Balanced offload: module={module_name}')
+                module.offload_dir = os.path.join(shared.opts.accelerate_offload_path, checkpoint_name, module_name)
+                module = accelerate.hooks.add_hook_to_module(module, offload_hook_instance, append=True)
+                module._hf_hook.execution_device = torch.device(devices.device) # pylint: disable=protected-access
+                if network_layer_name:
+                    module.network_layer_name = network_layer_name
+                if device_map and max_memory:
+                    module.balanced_offload_device_map = device_map
+                    module.balanced_offload_max_memory = max_memory
 
     apply_balanced_offload_to_module(sd_model)
     if hasattr(sd_model, "pipe"):
@@ -435,6 +470,8 @@ def apply_balanced_offload(sd_model):
     if hasattr(sd_model, "decoder_pipe"):
         apply_balanced_offload_to_module(sd_model.decoder_pipe)
     set_accelerate(sd_model)
+    devices.torch_gc(fast=True)
+    process_timer.add('offload', time.time() - t0)
     return sd_model
 
 
@@ -479,13 +516,13 @@ def move_model(model, device=None, force=False):
                             shared.log.error(f'Model move execution device: device={device} {e}')
     if getattr(model, 'has_accelerate', False) and not force:
         return
-    if hasattr(model, "device") and devices.normalize_device(model.device) == devices.normalize_device(device):
+    if hasattr(model, "device") and devices.normalize_device(model.device) == devices.normalize_device(device) and not force:
         return
     try:
         t0 = time.time()
         try:
             if hasattr(model, 'to'):
-                model.to(device)
+                model.to(device, non_blocking=True)
             if hasattr(model, "prior_pipe"):
                 model.prior_pipe.to(device)
         except Exception as e0:
@@ -515,7 +552,7 @@ def move_model(model, device=None, force=False):
     if 'move' not in process_timer.records:
         process_timer.records['move'] = 0
     process_timer.records['move'] += t1 - t0
-    if os.environ.get('SD_MOVE_DEBUG', None) or (t1-t0) > 1:
+    if os.environ.get('SD_MOVE_DEBUG', None) or (t1-t0) > 2:
         shared.log.debug(f'Model move: device={device} class={model.__class__.__name__} accelerate={getattr(model, "has_accelerate", False)} fn={fn} time={t1-t0:.2f}') # pylint: disable=protected-access
     devices.torch_gc()
 
@@ -1448,10 +1485,17 @@ def disable_offload(sd_model):
     from accelerate.hooks import remove_hook_from_module
     if not getattr(sd_model, 'has_accelerate', False):
         return
-    if hasattr(sd_model, 'components'):
-        for _name, model in sd_model.components.items():
-            if isinstance(model, torch.nn.Module):
-                remove_hook_from_module(model, recurse=True)
+    if hasattr(sd_model, "_internal_dict"):
+        keys = sd_model._internal_dict.keys() # pylint: disable=protected-access
+    else:
+        keys = get_signature(sd_model).keys()
+    for module_name in keys: # pylint: disable=protected-access
+        module = getattr(sd_model, module_name, None)
+        if isinstance(module, torch.nn.Module):
+            network_layer_name = getattr(module, "network_layer_name", None)
+            module = remove_hook_from_module(module, recurse=True)
+            if network_layer_name:
+                module.network_layer_name = network_layer_name
     sd_model.has_accelerate = False
 
 

modules/sd_samplers_diffusers.py

@@ -52,6 +52,7 @@ try:
     from modules.schedulers.scheduler_dc import DCSolverMultistepScheduler # pylint: disable=ungrouped-imports
     from modules.schedulers.scheduler_vdm import VDMScheduler # pylint: disable=ungrouped-imports
     from modules.schedulers.scheduler_dpm_flowmatch import FlowMatchDPMSolverMultistepScheduler # pylint: disable=ungrouped-imports
+    from modules.schedulers.scheduler_bdia import BDIA_DDIMScheduler # pylint: disable=ungrouped-imports
 except Exception as e:
     shared.log.error(f'Diffusers import error: version={diffusers.__version__} error: {e}')
     if os.environ.get('SD_SAMPLER_DEBUG', None) is not None:
@@ -97,6 +98,7 @@ config = {
     'VDM Solver': { 'clip_sample_range': 2.0, },
     'LCM': { 'beta_start': 0.00085, 'beta_end': 0.012, 'beta_schedule': "scaled_linear", 'set_alpha_to_one': True, 'rescale_betas_zero_snr': False, 'thresholding': False, 'timestep_spacing': 'linspace' },
     'TCD': { 'set_alpha_to_one': True, 'rescale_betas_zero_snr': False, 'beta_schedule': 'scaled_linear' },
+    'BDIA DDIM': { 'clip_sample': False, 'set_alpha_to_one': True, 'steps_offset': 0, 'clip_sample_range': 1.0, 'sample_max_value': 1.0, 'timestep_spacing': 'leading', 'rescale_betas_zero_snr': False, 'thresholding': False, 'gamma': 1.0 },
 
     'PNDM': { 'skip_prk_steps': False, 'set_alpha_to_one': False, 'steps_offset': 0, 'timestep_spacing': 'linspace' },
     'IPNDM': { },
@@ -142,6 +144,7 @@ samplers_data_diffusers = [
     sd_samplers_common.SamplerData('SA Solver', lambda model: DiffusionSampler('SA Solver', SASolverScheduler, model), [], {}),
     sd_samplers_common.SamplerData('DC Solver', lambda model: DiffusionSampler('DC Solver', DCSolverMultistepScheduler, model), [], {}),
     sd_samplers_common.SamplerData('VDM Solver', lambda model: DiffusionSampler('VDM Solver', VDMScheduler, model), [], {}),
+    sd_samplers_common.SamplerData('BDIA DDIM', lambda model: DiffusionSampler('BDIA DDIM g=0', BDIA_DDIMScheduler, model), [], {}),
 
     sd_samplers_common.SamplerData('PNDM', lambda model: DiffusionSampler('PNDM', PNDMScheduler, model), [], {}),
     sd_samplers_common.SamplerData('IPNDM', lambda model: DiffusionSampler('IPNDM', IPNDMScheduler, model), [], {}),

modules/shared.py

@@ -20,7 +20,7 @@ from modules import errors, devices, shared_items, shared_state, cmd_args, theme
 from modules.paths import models_path, script_path, data_path, sd_configs_path, sd_default_config, sd_model_file, default_sd_model_file, extensions_dir, extensions_builtin_dir # pylint: disable=W0611
 from modules.dml import memory_providers, default_memory_provider, directml_do_hijack
 from modules.onnx_impl import initialize_onnx, execution_providers
-from modules.memstats import memory_stats
+from modules.memstats import memory_stats, ram_stats # pylint: disable=unused-import
 from modules.ui_components import DropdownEditable
 import modules.interrogate
 import modules.memmon
@@ -132,7 +132,8 @@ def readfile(filename, silent=False, lock=False):
         #    data = json.loads(data)
         t1 = time.time()
         if not silent:
-            log.debug(f'Read: file="{filename}" json={len(data)} bytes={os.path.getsize(filename)} time={t1-t0:.3f}')
+            fn = f'{sys._getframe(2).f_code.co_name}:{sys._getframe(1).f_code.co_name}' # pylint: disable=protected-access
+            log.debug(f'Read: file="{filename}" json={len(data)} bytes={os.path.getsize(filename)} time={t1-t0:.3f} fn={fn}')
     except FileNotFoundError as err:
         log.debug(f'Reading failed: {filename} {err}')
     except Exception as err:
@@ -363,7 +364,7 @@ def list_samplers():
 
 def temp_disable_extensions():
     disable_safe = ['sd-webui-controlnet', 'multidiffusion-upscaler-for-automatic1111', 'a1111-sd-webui-lycoris', 'sd-webui-agent-scheduler', 'clip-interrogator-ext', 'stable-diffusion-webui-rembg', 'sd-extension-chainner', 'stable-diffusion-webui-images-browser']
-    disable_diffusers = ['sd-webui-controlnet', 'multidiffusion-upscaler-for-automatic1111', 'a1111-sd-webui-lycoris', 'sd-webui-animatediff']
+    disable_diffusers = ['sd-webui-controlnet', 'multidiffusion-upscaler-for-automatic1111', 'a1111-sd-webui-lycoris', 'sd-webui-animatediff', 'Lora']
     disable_themes = ['sd-webui-lobe-theme', 'cozy-nest', 'sdnext-modernui']
     disable_original = []
     disabled = []
@@ -559,8 +560,9 @@ options_templates.update(options_section(('diffusers', "Diffusers Settings"), {
     "diffusers_extract_ema": OptionInfo(False, "Use model EMA weights when possible"),
     "diffusers_generator_device": OptionInfo("GPU", "Generator device", gr.Radio, {"choices": ["GPU", "CPU", "Unset"]}),
     "diffusers_offload_mode": OptionInfo(startup_offload_mode, "Model offload mode", gr.Radio, {"choices": ['none', 'balanced', 'model', 'sequential']}),
-    "diffusers_offload_max_gpu_memory": OptionInfo(round(gpu_memory * 0.75, 1), "Max GPU memory before balanced offload", gr.Slider, {"minimum": 0, "maximum": gpu_memory, "step": 0.01, "visible": True }),
-    "diffusers_offload_max_cpu_memory": OptionInfo(round(cpu_memory * 0.75, 1), "Max CPU memory before balanced offload", gr.Slider, {"minimum": 0, "maximum": cpu_memory, "step": 0.01, "visible": False }),
+    "diffusers_offload_min_gpu_memory": OptionInfo(0.25, "Balanced offload GPU low watermark", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01 }),
+    "diffusers_offload_max_gpu_memory": OptionInfo(0.70, "Balanced offload GPU high watermark", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01 }),
+    "diffusers_offload_max_cpu_memory": OptionInfo(0.75, "Balanced offload CPU high watermark", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01 }),
     "diffusers_vae_upcast": OptionInfo("default", "VAE upcasting", gr.Radio, {"choices": ['default', 'true', 'false']}),
     "diffusers_vae_slicing": OptionInfo(True, "VAE slicing"),
     "diffusers_vae_tiling": OptionInfo(cmd_opts.lowvram or cmd_opts.medvram, "VAE tiling"),
@@ -901,15 +903,16 @@ options_templates.update(options_section(('extra_networks', "Networks"), {
     "wildcards_enabled": OptionInfo(True, "Enable file wildcards support"),
     "extra_networks_lora_sep": OptionInfo("<h2>LoRA</h2>", "", gr.HTML),
     "extra_networks_default_multiplier": OptionInfo(1.0, "Default strength", gr.Slider, {"minimum": 0.0, "maximum": 2.0, "step": 0.01}),
-    "lora_preferred_name": OptionInfo("filename", "LoRA preferred name", gr.Radio, {"choices": ["filename", "alias"]}),
+    "lora_preferred_name": OptionInfo("filename", "LoRA preferred name", gr.Radio, {"choices": ["filename", "alias"], "visible": False}),
     "lora_add_hashes_to_infotext": OptionInfo(False, "LoRA add hash info"),
+    "lora_fuse_diffusers": OptionInfo(False if not cmd_opts.use_openvino else True, "LoRA fuse directly to model"),
+    "lora_load_gpu": OptionInfo(True if not (cmd_opts.lowvram or cmd_opts.medvram) else False, "LoRA load directly to GPU"),
+    "lora_offload_backup": OptionInfo(True, "LoRA offload backup weights"),
     "lora_force_diffusers": OptionInfo(False if not cmd_opts.use_openvino else True, "LoRA force loading of all models using Diffusers"),
     "lora_maybe_diffusers": OptionInfo(False, "LoRA force loading of specific models using Diffusers"),
-    "lora_fuse_diffusers": OptionInfo(False if not cmd_opts.use_openvino else True, "LoRA use fuse when possible"),
     "lora_apply_tags": OptionInfo(0, "LoRA auto-apply tags", gr.Slider, {"minimum": -1, "maximum": 32, "step": 1}),
     "lora_in_memory_limit": OptionInfo(0, "LoRA memory cache", gr.Slider, {"minimum": 0, "maximum": 24, "step": 1}),
     "lora_quant": OptionInfo("NF4","LoRA precision in quantized models", gr.Radio, {"choices": ["NF4", "FP4"]}),
-    "lora_load_gpu": OptionInfo(True if not cmd_opts.lowvram else False, "Load LoRA directly to GPU"),
 }))
 
 options_templates.update(options_section((None, "Internal options"), {

modules/ui_control.py

@@ -29,15 +29,20 @@ def return_stats(t: float = None):
         elapsed_m = int(elapsed // 60)
         elapsed_s = elapsed % 60
         elapsed_text = f"Time: {elapsed_m}m {elapsed_s:.2f}s |" if elapsed_m > 0 else f"Time: {elapsed_s:.2f}s |"
-    summary = timer.process.summary(min_time=0.1, total=False).replace('=', ' ')
-    vram_html = ''
+    summary = timer.process.summary(min_time=0.25, total=False).replace('=', ' ')
+    gpu = ''
+    cpu = ''
     if not shared.mem_mon.disabled:
         vram = {k: -(v//-(1024*1024)) for k, v in shared.mem_mon.read().items()}
-        used = round(100 * vram['used'] / (vram['total'] + 0.001))
-        if vram.get('active_peak', 0) > 0:
-            vram_html += f"| GPU {max(vram['active_peak'], vram['reserved_peak'])} MB {used}%"
-            vram_html += f" | retries {vram['retries']} oom {vram['oom']}" if vram.get('retries', 0) > 0 or vram.get('oom', 0) > 0 else ''
-    return f"<div class='performance'><p>{elapsed_text} {summary} {vram_html}</p></div>"
+        peak = max(vram['active_peak'], vram['reserved_peak'], vram['used'])
+        used = round(100.0 * peak / vram['total']) if vram['total'] > 0 else 0
+        if used > 0:
+            gpu += f"| GPU {peak} MB {used}%"
+            gpu += f" | retries {vram['retries']} oom {vram['oom']}" if vram.get('retries', 0) > 0 or vram.get('oom', 0) > 0 else ''
+        ram = shared.ram_stats()
+        if ram['used'] > 0:
+            cpu += f"| RAM {ram['used']} GB {round(100.0 * ram['used'] / ram['total'])}%"
+    return f"<div class='performance'><p>Time: {elapsed_text} | {summary} {gpu} {cpu}</p></div>"
 
 
 def return_controls(res, t: float = None):

modules/ui_extra_networks.py

@@ -460,17 +460,20 @@ def register_page(page: ExtraNetworksPage):
 
 
 def register_pages():
-    from modules.ui_extra_networks_textual_inversion import ExtraNetworksPageTextualInversion
-    from modules.ui_extra_networks_checkpoints import ExtraNetworksPageCheckpoints
-    from modules.ui_extra_networks_styles import ExtraNetworksPageStyles
-    from modules.ui_extra_networks_vae import ExtraNetworksPageVAEs
-    from modules.ui_extra_networks_history import ExtraNetworksPageHistory
     debug('EN register-pages')
+    from modules.ui_extra_networks_checkpoints import ExtraNetworksPageCheckpoints
+    from modules.ui_extra_networks_vae import ExtraNetworksPageVAEs
+    from modules.ui_extra_networks_styles import ExtraNetworksPageStyles
+    from modules.ui_extra_networks_history import ExtraNetworksPageHistory
+    from modules.ui_extra_networks_textual_inversion import ExtraNetworksPageTextualInversion
     register_page(ExtraNetworksPageCheckpoints())
-    register_page(ExtraNetworksPageStyles())
-    register_page(ExtraNetworksPageTextualInversion())
     register_page(ExtraNetworksPageVAEs())
+    register_page(ExtraNetworksPageStyles())
     register_page(ExtraNetworksPageHistory())
+    register_page(ExtraNetworksPageTextualInversion())
+    if shared.native:
+        from modules.ui_extra_networks_lora import ExtraNetworksPageLora
+        register_page(ExtraNetworksPageLora())
     if shared.opts.hypernetwork_enabled:
         from modules.ui_extra_networks_hypernets import ExtraNetworksPageHypernetworks
         register_page(ExtraNetworksPageHypernetworks())

modules/ui_extra_networks_lora.py

@@ -0,0 +1,123 @@
+import os
+import json
+import concurrent
+import modules.lora.networks as networks
+from modules import shared, ui_extra_networks
+
+
+debug = os.environ.get('SD_LORA_DEBUG', None) is not None
+
+
+class ExtraNetworksPageLora(ui_extra_networks.ExtraNetworksPage):
+    def __init__(self):
+        super().__init__('Lora')
+        self.list_time = 0
+
+    def refresh(self):
+        networks.list_available_networks()
+
+    @staticmethod
+    def get_tags(l, info):
+        tags = {}
+        try:
+            if l.metadata is not None:
+                modelspec_tags = l.metadata.get('modelspec.tags', {})
+                possible_tags = l.metadata.get('ss_tag_frequency', {}) # tags from model metedata
+                if isinstance(possible_tags, str):
+                    possible_tags = {}
+                if isinstance(modelspec_tags, str):
+                    modelspec_tags = {}
+                if len(list(modelspec_tags)) > 0:
+                    possible_tags.update(modelspec_tags)
+                for k, v in possible_tags.items():
+                    words = k.split('_', 1) if '_' in k else [v, k]
+                    words = [str(w).replace('.json', '') for w in words]
+                    if words[0] == '{}':
+                        words[0] = 0
+                    tag = ' '.join(words[1:]).lower()
+                    tags[tag] = words[0]
+
+            def find_version():
+                found_versions = []
+                current_hash = l.hash[:8].upper()
+                all_versions = info.get('modelVersions', [])
+                for v in info.get('modelVersions', []):
+                    for f in v.get('files', []):
+                        if any(h.startswith(current_hash) for h in f.get('hashes', {}).values()):
+                            found_versions.append(v)
+                if len(found_versions) == 0:
+                    found_versions = all_versions
+                return found_versions
+
+            for v in find_version():  # trigger words from info json
+                possible_tags = v.get('trainedWords', [])
+                if isinstance(possible_tags, list):
+                    for tag_str in possible_tags:
+                        for tag in tag_str.split(','):
+                            tag = tag.strip().lower()
+                            if tag not in tags:
+                                tags[tag] = 0
+
+            possible_tags = info.get('tags', []) # tags from info json
+            if not isinstance(possible_tags, list):
+                possible_tags = list(possible_tags.values())
+            for tag in possible_tags:
+                tag = tag.strip().lower()
+                if tag not in tags:
+                    tags[tag] = 0
+        except Exception:
+            pass
+        bad_chars = [';', ':', '<', ">", "*", '?', '\'', '\"', '(', ')', '[', ']', '{', '}', '\\', '/']
+        clean_tags = {}
+        for k, v in tags.items():
+            tag = ''.join(i for i in k if i not in bad_chars).strip()
+            clean_tags[tag] = v
+
+        clean_tags.pop('img', None)
+        clean_tags.pop('dataset', None)
+        return clean_tags
+
+    def create_item(self, name):
+        l = networks.available_networks.get(name)
+        if l is None:
+            shared.log.warning(f'Networks: type=lora registered={len(list(networks.available_networks))} file="{name}" not registered')
+            return None
+        try:
+            # path, _ext = os.path.splitext(l.filename)
+            name = os.path.splitext(os.path.relpath(l.filename, shared.cmd_opts.lora_dir))[0]
+            item = {
+                "type": 'Lora',
+                "name": name,
+                "filename": l.filename,
+                "hash": l.shorthash,
+                "prompt": json.dumps(f" <lora:{l.get_alias()}:{shared.opts.extra_networks_default_multiplier}>"),
+                "metadata": json.dumps(l.metadata, indent=4) if l.metadata else None,
+                "mtime": os.path.getmtime(l.filename),
+                "size": os.path.getsize(l.filename),
+                "version": l.sd_version,
+            }
+            info = self.find_info(l.filename)
+            item["info"] = info
+            item["description"] = self.find_description(l.filename, info) # use existing info instead of double-read
+            item["tags"] = self.get_tags(l, info)
+            return item
+        except Exception as e:
+            shared.log.error(f'Networks: type=lora file="{name}" {e}')
+            if debug:
+                from modules import errors
+                errors.display(e, 'Lora')
+            return None
+
+    def list_items(self):
+        items = []
+        with concurrent.futures.ThreadPoolExecutor(max_workers=shared.max_workers) as executor:
+            future_items = {executor.submit(self.create_item, net): net for net in networks.available_networks}
+            for future in concurrent.futures.as_completed(future_items):
+                item = future.result()
+                if item is not None:
+                    items.append(item)
+        self.update_all_previews(items)
+        return items
+
+    def allowed_directories_for_previews(self):
+        return [shared.cmd_opts.lora_dir]

modules/ui_models.py

@@ -8,7 +8,7 @@ from modules import sd_models, sd_vae, extras
 from modules.ui_components import ToolButton
 from modules.ui_common import create_refresh_button
 from modules.call_queue import wrap_gradio_gpu_call
-from modules.shared import opts, log, req, readfile, max_workers
+from modules.shared import opts, log, req, readfile, max_workers, native
 import modules.ui_symbols
 import modules.errors
 import modules.hashes
@@ -794,6 +794,10 @@ def create_ui():
                 civit_results4.select(fn=civit_update_select, inputs=[civit_results4], outputs=[models_outcome, civit_update_download_btn])
                 civit_update_download_btn.click(fn=civit_update_download, inputs=[], outputs=[models_outcome])
 
+            if native:
+                from modules.lora.lora_extract import create_ui as lora_extract_ui
+                lora_extract_ui()
+
             for ui in extra_ui:
                 if callable(ui):
                     ui()

scripts/flux_tools.py

@@ -100,7 +100,7 @@ class Script(scripts.Script):
 
         if tool == 'Depth':
             # pipe = FluxControlPipeline.from_pretrained("black-forest-labs/FLUX.1-Depth-dev", torch_dtype=torch.bfloat16, revision="refs/pr/1").to("cuda")
-            install('git+https://github.com/asomoza/image_gen_aux.git', 'image_gen_aux')
+            install('git+https://github.com/huggingface/image_gen_aux.git', 'image_gen_aux')
             if shared.sd_model.__class__.__name__ != 'FluxControlPipeline' or 'Depth' not in shared.opts.sd_model_checkpoint:
                 shared.opts.data["sd_model_checkpoint"] = "black-forest-labs/FLUX.1-Depth-dev"
                 sd_models.reload_model_weights(op='model', revision="refs/pr/1")

scripts/xyz_grid.py

@@ -12,6 +12,7 @@ import gradio as gr
 from scripts.xyz_grid_shared import str_permutations, list_to_csv_string, re_range # pylint: disable=no-name-in-module
 from scripts.xyz_grid_classes import axis_options, AxisOption, SharedSettingsStackHelper # pylint: disable=no-name-in-module
 from scripts.xyz_grid_draw import draw_xyz_grid # pylint: disable=no-name-in-module
+from scripts.xyz_grid_shared import apply_field, apply_task_args, apply_setting, apply_prompt, apply_order, apply_sampler, apply_hr_sampler_name, confirm_samplers, apply_checkpoint, apply_refiner, apply_unet, apply_dict, apply_clip_skip, apply_vae, list_lora, apply_lora, apply_lora_strength, apply_te, apply_styles, apply_upscaler, apply_context, apply_detailer, apply_override, apply_processing, apply_options, apply_seed, format_value_add_label, format_value, format_value_join_list, do_nothing, format_nothing # pylint: disable=no-name-in-module, unused-import
 from modules import shared, errors, scripts, images, processing
 from modules.ui_components import ToolButton
 import modules.ui_symbols as symbols

scripts/xyz_grid_on.py

@@ -413,6 +413,7 @@ class Script(scripts.Script):
 
         p.do_not_save_grid = True
         p.do_not_save_samples = True
+        p.disable_extra_networks = True
         active = False
         cache = processed
         return processed

scripts/xyz_grid_shared.py

@@ -192,7 +192,7 @@ def apply_vae(p, x, xs):
 
 def list_lora():
     import sys
-    lora = [v for k, v in sys.modules.items() if k == 'networks'][0]
+    lora = [v for k, v in sys.modules.items() if k == 'networks' or k == 'modules.lora.networks'][0]
     loras = [v.fullname for v in lora.available_networks.values()]
     return ['None'] + loras
 

webui.py

@@ -8,6 +8,7 @@ import logging
 import importlib
 import contextlib
 from threading import Thread
+import modules.hashes
 import modules.loader
 import torch # pylint: disable=wrong-import-order
 from modules import timer, errors, paths # pylint: disable=unused-import
@@ -18,6 +19,7 @@ from modules import extra_networks, ui_extra_networks # pylint: disable=ungroupe
 from modules.paths import create_paths
 from modules.call_queue import queue_lock, wrap_queued_call, wrap_gradio_gpu_call # pylint: disable=unused-import
 import modules.devices
+import modules.sd_checkpoint
 import modules.sd_samplers
 import modules.lowvram
 import modules.scripts
@@ -77,6 +79,9 @@ def check_rollback_vae():
 
 def initialize():
     log.debug('Initializing')
+
+    modules.sd_checkpoint.init_metadata()
+    modules.hashes.init_cache()
     check_rollback_vae()
 
     modules.sd_samplers.list_samplers()
@@ -89,15 +94,20 @@ def initialize():
     timer.startup.record("unet")
 
     modules.model_te.refresh_te_list()
-    timer.startup.record("unet")
-
-    extensions.list_extensions()
-    timer.startup.record("extensions")
+    timer.startup.record("te")
 
     modelloader.cleanup_models()
     modules.sd_models.setup_model()
     timer.startup.record("models")
 
+    if shared.native:
+        import modules.lora.networks as lora_networks
+        lora_networks.list_available_networks()
+        timer.startup.record("lora")
+
+    shared.prompt_styles.reload()
+    timer.startup.record("styles")
+
     import modules.postprocess.codeformer_model as codeformer
     codeformer.setup_model(shared.opts.codeformer_models_path)
     sys.modules["modules.codeformer_model"] = codeformer
@@ -107,6 +117,9 @@ def initialize():
     yolo.initialize()
     timer.startup.record("detailer")
 
+    extensions.list_extensions()
+    timer.startup.record("extensions")
+
     log.info('Load extensions')
     t_timer, t_total = modules.scripts.load_scripts()
     timer.startup.record("extensions")
@@ -116,8 +129,9 @@ def initialize():
     modelloader.load_upscalers()
     timer.startup.record("upscalers")
 
-    shared.reload_hypernetworks()
-    shared.prompt_styles.reload()
+    if shared.opts.hypernetwork_enabled:
+        shared.reload_hypernetworks()
+        timer.startup.record("hypernetworks")
 
     ui_extra_networks.initialize()
     ui_extra_networks.register_pages()