From 27442697a2a258d7e5a736322677dec02d30dcac Mon Sep 17 00:00:00 2001
From: AI-Casanova <54461896+AI-Casanova@users.noreply.github.com>
Date: Fri, 18 Oct 2024 22:29:04 -0500
Subject: [PATCH 1/7] WIP Lora Extract

---
 extensions-builtin/Lora/lora_extract.py       | 192 ++++++++++++++++++
 extensions-builtin/Lora/networks.py           |  30 +--
 .../Lora/scripts/lora_script.py               |   2 +
 3 files changed, 212 insertions(+), 12 deletions(-)
 create mode 100644 extensions-builtin/Lora/lora_extract.py

diff --git a/extensions-builtin/Lora/lora_extract.py b/extensions-builtin/Lora/lora_extract.py
new file mode 100644
index 000000000..9f20d101d
--- /dev/null
+++ b/extensions-builtin/Lora/lora_extract.py
@@ -0,0 +1,192 @@
+import torch
+from os import path
+from safetensors.torch import save_file
+import gradio as gr
+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):
+        self.network_name = None
+        self.U = None
+        self.S = None
+        self.Vh = None
+        self.rank = 0
+        self.maxrank = 0
+        self.out_size = None
+        self.in_size = None
+        self.kernel_size = None
+        self.conv2d = 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.linalg.svd(diffweight.to(device=devices.device, dtype=torch.float))
+        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.to(device=devices.cpu, dtype=torch.bfloat16)
+
+    def findrank(self, maxrank, rankratio):
+        if rankratio < 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, rankratio ** 2)) + 1
+            index = max(1, min(index, len(self.S) - 1))
+            self.rank = index
+            if maxrank > 0:
+                self.rank = min(self.rank, maxrank)
+        elif maxrank == 0:
+            self.rank = min(self.in_size, self.out_size)
+        else:
+            self.rank = min(self.in_size, self.out_size, maxrank)
+
+    def makeweights(self, rankoverride=None):
+        if rankoverride:
+            self.rank = min(self.in_size, self.out_size, rankoverride)
+        up = self.U[:, :self.rank] @ torch.diag(self.S[:self.rank])
+        down = self.Vh[:self.rank, :]
+        if self.conv2d:
+            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])
+        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 ", ".join(list(loaded))
+
+
+def make_lora(basename, rank, auto_rank, rank_ratio, constant_rank):
+    if not shared.sd_loaded or not shared.native or loaded_lora() == "":
+        return
+    rank = int(rank)
+    rank_ratio = 1 if not auto_rank else rank_ratio
+    constant_rank = False if not auto_rank else constant_rank
+    rank_overide = 0 if constant_rank else None
+
+    if hasattr(shared.sd_model, 'text_encoder') and shared.sd_model.text_encoder is not None:
+        for name, module in shared.sd_model.text_encoder.named_modules():
+            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()
+            module.svdhandler.network_name = prefix + name.replace(".", "_")
+            with devices.inference_context():
+                module.svdhandler.decompose(module.weight, weights_backup)
+            module.svdhandler.findrank(rank, rank_ratio)
+    print("TE1 done")
+    if hasattr(shared.sd_model, 'text_encoder_2'):
+        for name, module in shared.sd_model.text_encoder_2.named_modules():
+            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()
+            module.svdhandler.network_name = "lora_te2_" + name.replace(".", "_")
+            with devices.inference_context():
+                module.svdhandler.decompose(module.weight, weights_backup)
+            module.svdhandler.findrank(rank, rank_ratio)
+
+    print("TE2 done")
+    if hasattr(shared.sd_model, 'unet'):
+        for name, module in shared.sd_model.unet.named_modules():
+            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()
+            module.svdhandler.network_name = "lora_unet_" + name.replace(".", "_")
+            with devices.inference_context():
+                module.svdhandler.decompose(module.weight, weights_backup)
+            module.svdhandler.findrank(rank, rank_ratio)
+
+    # if hasattr(shared.sd_model, 'transformer'):  # TODO: Handle quant for Flux
+    #     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)
+
+    submodelname = ['text_encoder', 'text_encoder_2', 'unet', 'transformer']
+
+    if constant_rank:
+        for sub in submodelname:
+            submodel = getattr(shared.sd_model, sub, None)
+            if submodel is not None:
+                for name, module in submodel.named_modules():
+                    if not hasattr(module, "svdhandler"):
+                        continue
+                    rank_overide = max(rank_overide, module.svdhandler.rank)
+    print(f"rank_overide: {rank_overide}")
+    lora_state_dict = {}
+    for sub in submodelname:
+        submodel = getattr(shared.sd_model, sub, None)
+        if submodel is not None:
+            for name, module in submodel.named_modules():
+                if not hasattr(module, "svdhandler"):
+                    continue
+                lora_state_dict.update(module.svdhandler.makeweights(rank_overide))
+                del module.svdhandler
+
+    save_file(lora_state_dict, path.join(shared.cmd_opts.lora_dir, basename+".safetensors"))
+
+
+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(label="Loaded LoRA", interactive=False)
+            # create_refresh_button(loaded, lambda: None, gr.update(value=loaded_lora()), "testid")
+            create_refresh_button(loaded, lambda: None, lambda: {'value': loaded_lora()}, "testid")
+        with gr.Row():
+            rank = gr.Number(value=0, label="Optional max rank")
+        with gr.Row():
+            auto_rank = gr.Checkbox(value=False, label="Automatically determine rank")
+        with gr.Row(visible=False) as rank_options:
+            rank_ratio = gr.Slider(minimum=0, maximum=1, value=1, label="Autorank ratio", visible=True)
+            constant_rank = gr.Checkbox(value=False, label="Constant rank", visible=True)
+        with gr.Row():
+            basename = gr.Textbox(label="Base name for LoRa")
+        with gr.Row():
+            extract = gr.Button(value="Extract Lora", variant='primary')
+
+    auto_rank.change(fn=lambda x: gr_show(x), inputs=[auto_rank], outputs=[rank_options])
+    # extract.click(
+    #     fn=wrap_gradio_gpu_call(make_lora(basename, rank, auto_rank, rank_ratio, constant_rank),
+    #                             extra_outputs=None), _js='loraextract', inputs=[],
+    #     outputs=[])
+    extract.click(fn=make_lora, inputs=[basename, rank, auto_rank, rank_ratio, constant_rank], outputs=[])
+    # extract.click(fn= lambda: None)
diff --git a/extensions-builtin/Lora/networks.py b/extensions-builtin/Lora/networks.py
index f911e1b3e..3814fb50a 100644
--- a/extensions-builtin/Lora/networks.py
+++ b/extensions-builtin/Lora/networks.py
@@ -285,6 +285,8 @@ def network_restore_weights_from_backup(self: Union[torch.nn.Conv2d, torch.nn.Li
     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:
+        t1 = time.time()
+        timer['restore'] += t1 - t0
         return
     # if debug:
     #     shared.log.debug('LoRA restore weights')
@@ -319,18 +321,7 @@ def network_restore_weights_from_backup(self: Union[torch.nn.Conv2d, torch.nn.Li
     timer['restore'] += t1 - t0
 
 
-def network_apply_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.GroupNorm, torch.nn.LayerNorm, torch.nn.MultiheadAttention, diffusers.models.lora.LoRACompatibleLinear, diffusers.models.lora.LoRACompatibleConv]):
-    """
-    Applies the currently selected set of networks to the weights of torch layer self.
-    If weights already have this particular set of networks applied, does nothing.
-    If not, restores orginal weights from backup and alters weights according to networks.
-    """
-    network_layer_name = getattr(self, 'network_layer_name', None)
-    if network_layer_name is None:
-        return
-    t0 = time.time()
-    current_names = getattr(self, "network_current_names", ())
-    wanted_names = tuple((x.name, x.te_multiplier, x.unet_multiplier, x.dyn_dim) for x in loaded_networks)
+def maybe_backup_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.GroupNorm, torch.nn.LayerNorm, torch.nn.MultiheadAttention, diffusers.models.lora.LoRACompatibleLinear, diffusers.models.lora.LoRACompatibleConv], wanted_names, current_names):
     weights_backup = getattr(self, "network_weights_backup", None)
     if weights_backup is None and wanted_names != (): # pylint: disable=C1803
         if current_names != ():
@@ -360,6 +351,21 @@ def network_apply_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn
             bias_backup = None
         self.network_bias_backup = bias_backup
 
+
+def network_apply_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.GroupNorm, torch.nn.LayerNorm, torch.nn.MultiheadAttention, diffusers.models.lora.LoRACompatibleLinear, diffusers.models.lora.LoRACompatibleConv]):
+    """
+    Applies the currently selected set of networks to the weights of torch layer self.
+    If weights already have this particular set of networks applied, does nothing.
+    If not, restores orginal weights from backup and alters weights according to networks.
+    """
+    network_layer_name = getattr(self, 'network_layer_name', None)
+    if network_layer_name is None:
+        return
+    t0 = time.time()
+    current_names = getattr(self, "network_current_names", ())
+    wanted_names = tuple((x.name, x.te_multiplier, x.unet_multiplier, x.dyn_dim) for x in loaded_networks)
+    if any([net.modules.get(network_layer_name, None) for net in loaded_networks]):
+        maybe_backup_weights(self, wanted_names, current_names)
     if current_names != wanted_names:
         network_restore_weights_from_backup(self)
         for net in loaded_networks:
diff --git a/extensions-builtin/Lora/scripts/lora_script.py b/extensions-builtin/Lora/scripts/lora_script.py
index 9302a061e..ffbef47d9 100644
--- a/extensions-builtin/Lora/scripts/lora_script.py
+++ b/extensions-builtin/Lora/scripts/lora_script.py
@@ -1,6 +1,7 @@
 import re
 import networks
 import lora # pylint: disable=unused-import
+from lora_extract import create_ui
 from network import NetworkOnDisk
 from ui_extra_networks_lora import ExtraNetworksPageLora
 from extra_networks_lora import ExtraNetworkLora
@@ -14,6 +15,7 @@ def before_ui():
     ui_extra_networks.register_page(ExtraNetworksPageLora())
     networks.extra_network_lora = ExtraNetworkLora()
     extra_networks.register_extra_network(networks.extra_network_lora)
+    ui_models.extra_ui.append(create_ui)
 
 
 def create_lora_json(obj: NetworkOnDisk):

From 6a9b7bccd59d5303f74ce57f3405b0bcf459bddd Mon Sep 17 00:00:00 2001
From: AI-Casanova <54461896+AI-Casanova@users.noreply.github.com>
Date: Sat, 19 Oct 2024 22:50:34 -0500
Subject: [PATCH 2/7] finish Lora Extract

---
 extensions-builtin/Lora/lora_extract.py | 99 +++++++++++--------------
 1 file changed, 45 insertions(+), 54 deletions(-)

diff --git a/extensions-builtin/Lora/lora_extract.py b/extensions-builtin/Lora/lora_extract.py
index 9f20d101d..fc5e342a0 100644
--- a/extensions-builtin/Lora/lora_extract.py
+++ b/extensions-builtin/Lora/lora_extract.py
@@ -1,20 +1,21 @@
 import torch
+import time
 from os import path
 from safetensors.torch import save_file
 import gradio as gr
 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):
+    def __init__(self, maxrank=0, rank_ratio=1):
         self.network_name = None
         self.U = None
         self.S = None
         self.Vh = None
+        self.maxrank = maxrank
+        self.rank_ratio = rank_ratio
         self.rank = 0
-        self.maxrank = 0
         self.out_size = None
         self.in_size = None
         self.kernel_size = None
@@ -31,31 +32,28 @@ class SVDHandler:
                 diffweight = diffweight.flatten(start_dim=1)
             else:
                 diffweight = diffweight.squeeze()
-
-        self.U, self.S, self.Vh = torch.linalg.svd(diffweight.to(device=devices.device, dtype=torch.float))
-        del diffweight
+        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.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, maxrank, rankratio):
-        if rankratio < 1:
+    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, rankratio ** 2)) + 1
+            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 maxrank > 0:
-                self.rank = min(self.rank, maxrank)
-        elif maxrank == 0:
-            self.rank = min(self.in_size, self.out_size)
+            if self.maxrank > 0:
+                self.rank = min(self.rank, self.maxrank)
         else:
-            self.rank = min(self.in_size, self.out_size, maxrank)
+            self.rank = min(self.in_size, self.out_size, self.maxrank)
 
-    def makeweights(self, rankoverride=None):
-        if rankoverride:
-            self.rank = min(self.in_size, self.out_size, rankoverride)
+    def makeweights(self):
+        self.findrank()
         up = self.U[:, :self.rank] @ torch.diag(self.S[:self.rank])
         down = self.Vh[:self.rank, :]
         if self.conv2d:
@@ -81,13 +79,18 @@ def loaded_lora():
     return ", ".join(list(loaded))
 
 
-def make_lora(basename, rank, auto_rank, rank_ratio, constant_rank):
-    if not shared.sd_loaded or not shared.native or loaded_lora() == "":
+def make_lora(basename, maxrank, auto_rank, rank_ratio):
+    if not shared.sd_loaded or not shared.native:
         return
-    rank = int(rank)
+    if loaded_lora() == "":
+        shared.log.warning("Lora extract: No LoRA detected. Aborting...")
+        return
+    if not basename:
+        shared.log.warning("Lora extract: Base name required. Aborting...")
+        return
+    t0 = time.time()
+    maxrank = int(maxrank)
     rank_ratio = 1 if not auto_rank else rank_ratio
-    constant_rank = False if not auto_rank else constant_rank
-    rank_overide = 0 if constant_rank else None
 
     if hasattr(shared.sd_model, 'text_encoder') and shared.sd_model.text_encoder is not None:
         for name, module in shared.sd_model.text_encoder.named_modules():
@@ -95,34 +98,30 @@ def make_lora(basename, rank, auto_rank, rank_ratio, constant_rank):
             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()
+            module.svdhandler = SVDHandler(maxrank, rank_ratio)
             module.svdhandler.network_name = prefix + name.replace(".", "_")
             with devices.inference_context():
                 module.svdhandler.decompose(module.weight, weights_backup)
-            module.svdhandler.findrank(rank, rank_ratio)
-    print("TE1 done")
+
     if hasattr(shared.sd_model, 'text_encoder_2'):
         for name, module in shared.sd_model.text_encoder_2.named_modules():
             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()
+            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)
-            module.svdhandler.findrank(rank, rank_ratio)
 
-    print("TE2 done")
     if hasattr(shared.sd_model, 'unet'):
         for name, module in shared.sd_model.unet.named_modules():
             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()
+            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)
-            module.svdhandler.findrank(rank, rank_ratio)
 
     # if hasattr(shared.sd_model, 'transformer'):  # TODO: Handle quant for Flux
     #     for name, module in shared.sd_model.transformer.named_modules():
@@ -138,15 +137,6 @@ def make_lora(basename, rank, auto_rank, rank_ratio, constant_rank):
 
     submodelname = ['text_encoder', 'text_encoder_2', 'unet', 'transformer']
 
-    if constant_rank:
-        for sub in submodelname:
-            submodel = getattr(shared.sd_model, sub, None)
-            if submodel is not None:
-                for name, module in submodel.named_modules():
-                    if not hasattr(module, "svdhandler"):
-                        continue
-                    rank_overide = max(rank_overide, module.svdhandler.rank)
-    print(f"rank_overide: {rank_overide}")
     lora_state_dict = {}
     for sub in submodelname:
         submodel = getattr(shared.sd_model, sub, None)
@@ -154,39 +144,40 @@ def make_lora(basename, rank, auto_rank, rank_ratio, constant_rank):
             for name, module in submodel.named_modules():
                 if not hasattr(module, "svdhandler"):
                     continue
-                lora_state_dict.update(module.svdhandler.makeweights(rank_overide))
+                lora_state_dict.update(module.svdhandler.makeweights())
                 del module.svdhandler
 
-    save_file(lora_state_dict, path.join(shared.cmd_opts.lora_dir, basename+".safetensors"))
+    suffix = []
+    if maxrank and auto_rank and rank_ratio != 1:
+        suffix.append(f'maxrank{str(maxrank).replace(".","-")}')
+    else:
+        suffix.append(f'rank{str(maxrank).replace(".","-")}')
+    if auto_rank and rank_ratio != 1:
+        suffix.append(f'autorank{str(rank_ratio).replace(".","-")}')
+    pathstr = str(path.join(shared.cmd_opts.lora_dir, basename+f'_{"_".join(suffix)}.safetensors'))
+    save_file(lora_state_dict, pathstr)
+    shared.log.info(f'LoRA extracted to {pathstr} in {time.time()-t0} seconds')
 
 
 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(label="Loaded LoRA", interactive=False)
-            # create_refresh_button(loaded, lambda: None, gr.update(value=loaded_lora()), "testid")
+            loaded = gr.Textbox(value="Press refresh to query loaded LoRA", label="Loaded LoRA", interactive=False)
             create_refresh_button(loaded, lambda: None, lambda: {'value': loaded_lora()}, "testid")
         with gr.Row():
-            rank = gr.Number(value=0, label="Optional max rank")
+            rank = gr.Number(value=32, label="Max rank to extract", minimum=1)
         with gr.Row():
             auto_rank = gr.Checkbox(value=False, label="Automatically determine rank")
         with gr.Row(visible=False) as rank_options:
             rank_ratio = gr.Slider(minimum=0, maximum=1, value=1, label="Autorank ratio", visible=True)
-            constant_rank = gr.Checkbox(value=False, label="Constant rank", visible=True)
         with gr.Row():
             basename = gr.Textbox(label="Base name for LoRa")
         with gr.Row():
             extract = gr.Button(value="Extract Lora", variant='primary')
 
     auto_rank.change(fn=lambda x: gr_show(x), inputs=[auto_rank], outputs=[rank_options])
-    # extract.click(
-    #     fn=wrap_gradio_gpu_call(make_lora(basename, rank, auto_rank, rank_ratio, constant_rank),
-    #                             extra_outputs=None), _js='loraextract', inputs=[],
-    #     outputs=[])
-    extract.click(fn=make_lora, inputs=[basename, rank, auto_rank, rank_ratio, constant_rank], outputs=[])
-    # extract.click(fn= lambda: None)
+
+    extract.click(fn=make_lora, inputs=[basename, rank, auto_rank, rank_ratio], outputs=[])

From 2570d87ad3d97620049e82d6da45670ff265dca6 Mon Sep 17 00:00:00 2001
From: Vladimir Mandic <mandic00@live.com>
Date: Sun, 20 Oct 2024 08:30:06 -0400
Subject: [PATCH 3/7] typing, linting

Signed-off-by: Vladimir Mandic <mandic00@live.com>
---
 extensions-builtin/Lora/lora_extract.py | 64 ++++++++++++-------------
 1 file changed, 32 insertions(+), 32 deletions(-)

diff --git a/extensions-builtin/Lora/lora_extract.py b/extensions-builtin/Lora/lora_extract.py
index fc5e342a0..aa5c7d8d8 100644
--- a/extensions-builtin/Lora/lora_extract.py
+++ b/extensions-builtin/Lora/lora_extract.py
@@ -1,6 +1,6 @@
-import torch
+import os
 import time
-from os import path
+import torch
 from safetensors.torch import save_file
 import gradio as gr
 from modules import shared, devices
@@ -9,17 +9,17 @@ from modules.ui_common import create_refresh_button
 
 class SVDHandler:
     def __init__(self, maxrank=0, rank_ratio=1):
-        self.network_name = None
-        self.U = None
-        self.S = None
-        self.Vh = None
-        self.maxrank = maxrank
-        self.rank_ratio = rank_ratio
-        self.rank = 0
-        self.out_size = None
-        self.in_size = None
-        self.kernel_size = None
-        self.conv2d = False
+        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
@@ -32,8 +32,7 @@ class SVDHandler:
                 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)
+        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)
@@ -56,9 +55,9 @@ class SVDHandler:
         self.findrank()
         up = self.U[:, :self.rank] @ torch.diag(self.S[:self.rank])
         down = self.Vh[:self.rank, :]
-        if self.conv2d:
+        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])
+            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]),
@@ -71,7 +70,7 @@ def loaded_lora():
         return ""
     loaded = set()
     if hasattr(shared.sd_model, 'unet'):
-        for name, module in shared.sd_model.unet.named_modules():
+        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]
@@ -79,14 +78,14 @@ def loaded_lora():
     return ", ".join(list(loaded))
 
 
-def make_lora(basename, maxrank, auto_rank, rank_ratio):
+def make_lora(filename, maxrank, auto_rank, rank_ratio):
     if not shared.sd_loaded or not shared.native:
         return
     if loaded_lora() == "":
-        shared.log.warning("Lora extract: No LoRA detected. Aborting...")
+        shared.log.warning("LoRA extract: no LoRA detected")
         return
-    if not basename:
-        shared.log.warning("Lora extract: Base name required. Aborting...")
+    if not filename:
+        shared.log.warning("LoRA extract: target filename required")
         return
     t0 = time.time()
     maxrank = int(maxrank)
@@ -123,7 +122,8 @@ def make_lora(basename, maxrank, auto_rank, rank_ratio):
             with devices.inference_context():
                 module.svdhandler.decompose(module.weight, weights_backup)
 
-    # if hasattr(shared.sd_model, 'transformer'):  # TODO: Handle quant for Flux
+    # TODO: Handle quant for Flux
+    # if hasattr(shared.sd_model, 'transformer'):
     #     for name, module in shared.sd_model.transformer.named_modules():
     #         if "norm" in name and "linear" not in name:
     #             continue
@@ -141,7 +141,7 @@ def make_lora(basename, maxrank, auto_rank, rank_ratio):
     for sub in submodelname:
         submodel = getattr(shared.sd_model, sub, None)
         if submodel is not None:
-            for name, module in submodel.named_modules():
+            for _name, module in submodel.named_modules():
                 if not hasattr(module, "svdhandler"):
                     continue
                 lora_state_dict.update(module.svdhandler.makeweights())
@@ -154,9 +154,10 @@ def make_lora(basename, maxrank, auto_rank, rank_ratio):
         suffix.append(f'rank{str(maxrank).replace(".","-")}')
     if auto_rank and rank_ratio != 1:
         suffix.append(f'autorank{str(rank_ratio).replace(".","-")}')
-    pathstr = str(path.join(shared.cmd_opts.lora_dir, basename+f'_{"_".join(suffix)}.safetensors'))
+
+    pathstr = str(os.path.join(shared.cmd_opts.lora_dir, filename+f'_{"_".join(suffix)}.safetensors'))
     save_file(lora_state_dict, pathstr)
-    shared.log.info(f'LoRA extracted to {pathstr} in {time.time()-t0} seconds')
+    shared.log.info(f'LoRA extra: fn={pathstr} in {time.time()-t0} seconds')
 
 
 def create_ui():
@@ -168,16 +169,15 @@ def create_ui():
             loaded = gr.Textbox(value="Press refresh to query loaded LoRA", label="Loaded LoRA", interactive=False)
             create_refresh_button(loaded, lambda: None, lambda: {'value': loaded_lora()}, "testid")
         with gr.Row():
-            rank = gr.Number(value=32, label="Max rank to extract", minimum=1)
+            rank = gr.Slider(label="Maximum rank", value=32, minimum=1, maximum=256)
         with gr.Row():
             auto_rank = gr.Checkbox(value=False, label="Automatically determine rank")
         with gr.Row(visible=False) as rank_options:
-            rank_ratio = gr.Slider(minimum=0, maximum=1, value=1, label="Autorank ratio", visible=True)
+            rank_ratio = gr.Slider(label="Autorank ratio", value=1, minimum=0, maximum=1, step=0.05, visible=True)
         with gr.Row():
-            basename = gr.Textbox(label="Base name for LoRa")
+            filename = gr.Textbox(label="LoRA target filename")
         with gr.Row():
-            extract = gr.Button(value="Extract Lora", variant='primary')
+            extract = gr.Button(value="Extract LoRA", variant='primary')
 
     auto_rank.change(fn=lambda x: gr_show(x), inputs=[auto_rank], outputs=[rank_options])
-
-    extract.click(fn=make_lora, inputs=[basename, rank, auto_rank, rank_ratio], outputs=[])
+    extract.click(fn=make_lora, inputs=[filename, rank, auto_rank, rank_ratio], outputs=[])

From 64f363283fc830eee78e966a02b247cb9377e0d8 Mon Sep 17 00:00:00 2001
From: Vladimir Mandic <mandic00@live.com>
Date: Sun, 20 Oct 2024 09:26:07 -0400
Subject: [PATCH 4/7] messages,stats,save

Signed-off-by: Vladimir Mandic <mandic00@live.com>
---
 extensions-builtin/Lora/lora_extract.py | 101 +++++++++++++++++-------
 extensions-builtin/sdnext-modernui      |   2 +-
 modules/devices.py                      |  12 +--
 modules/ui_sections.py                  |   2 +-
 4 files changed, 80 insertions(+), 37 deletions(-)

diff --git a/extensions-builtin/Lora/lora_extract.py b/extensions-builtin/Lora/lora_extract.py
index aa5c7d8d8..de6404955 100644
--- a/extensions-builtin/Lora/lora_extract.py
+++ b/extensions-builtin/Lora/lora_extract.py
@@ -78,20 +78,36 @@ def loaded_lora():
     return ", ".join(list(loaded))
 
 
-def make_lora(filename, maxrank, auto_rank, rank_ratio):
+def make_meta():
+    return {
+        'todo': 'extra-lora'
+    }
+
+
+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() == "":
-        shared.log.warning("LoRA extract: no LoRA detected")
+        msg = "LoRA extract: no LoRA detected"
+        shared.log.warning(msg)
+        yield msg
         return
-    if not filename:
-        shared.log.warning("LoRA extract: target filename required")
+    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.state.begin('LoRA extract')
 
-    if hasattr(shared.sd_model, 'text_encoder') and shared.sd_model.text_encoder is not None:
+    shared.log.debug(f'LoRA extract: modules={modules} maxrank={maxrank} auto={auto_rank} ratio={rank_ratio} fn="{fn}"')
+    if 'te' in modules and getattr(shared.sd_model, 'text_encoder', None) is not None:
+        yield "LoRA extract: extracting TE-1"
         for name, module in shared.sd_model.text_encoder.named_modules():
             weights_backup = getattr(module, "network_weights_backup", None)
             if weights_backup is None or getattr(module, "network_current_names", None) is None:
@@ -101,8 +117,10 @@ def make_lora(filename, maxrank, auto_rank, rank_ratio):
             module.svdhandler.network_name = prefix + name.replace(".", "_")
             with devices.inference_context():
                 module.svdhandler.decompose(module.weight, weights_backup)
+    t1 = time.time()
 
-    if hasattr(shared.sd_model, 'text_encoder_2'):
+    if 'te' in modules and getattr(shared.sd_model, 'text_encoder_2', None) is not None:
+        yield "LoRA extract: extracting TE-2"
         for name, module in shared.sd_model.text_encoder_2.named_modules():
             weights_backup = getattr(module, "network_weights_backup", None)
             if weights_backup is None or getattr(module, "network_current_names", None) is None:
@@ -111,8 +129,10 @@ def make_lora(filename, maxrank, auto_rank, rank_ratio):
             module.svdhandler.network_name = "lora_te2_" + name.replace(".", "_")
             with devices.inference_context():
                 module.svdhandler.decompose(module.weight, weights_backup)
+    t2 = time.time()
 
-    if hasattr(shared.sd_model, 'unet'):
+    if 'unet' in modules and getattr(shared.sd_model, 'unet', None) is not None:
+        yield "LoRA extract: extracting UNet"
         for name, module in shared.sd_model.unet.named_modules():
             weights_backup = getattr(module, "network_weights_backup", None)
             if weights_backup is None or getattr(module, "network_current_names", None) is None:
@@ -121,9 +141,10 @@ def make_lora(filename, maxrank, auto_rank, rank_ratio):
             module.svdhandler.network_name = "lora_unet_" + name.replace(".", "_")
             with devices.inference_context():
                 module.svdhandler.decompose(module.weight, weights_backup)
+    t3 = time.time()
 
     # TODO: Handle quant for Flux
-    # if hasattr(shared.sd_model, 'transformer'):
+    # 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
@@ -135,29 +156,48 @@ def make_lora(filename, maxrank, auto_rank, rank_ratio):
     #         module.svdhandler.decompose(module.weight, weights_backup)
     #         module.svdhandler.findrank(rank, rank_ratio)
 
-    submodelname = ['text_encoder', 'text_encoder_2', 'unet', 'transformer']
-
     lora_state_dict = {}
-    for sub in submodelname:
+    for sub in ['text_encoder', 'text_encoder_2', 'unet', 'transformer']:
         submodel = getattr(shared.sd_model, sub, None)
         if submodel is not None:
+            yield f"LoRA extract: creating {sub}"
             for _name, module in submodel.named_modules():
                 if not hasattr(module, "svdhandler"):
                     continue
                 lora_state_dict.update(module.svdhandler.makeweights())
                 del module.svdhandler
+    shared.log.debug('LoRA extract: create done')
+    t4 = time.time()
 
-    suffix = []
-    if maxrank and auto_rank and rank_ratio != 1:
-        suffix.append(f'maxrank{str(maxrank).replace(".","-")}')
-    else:
-        suffix.append(f'rank{str(maxrank).replace(".","-")}')
-    if auto_rank and rank_ratio != 1:
-        suffix.append(f'autorank{str(rank_ratio).replace(".","-")}')
+    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:
+            shared.log.warning(f'LoRA extract: fn="{fn}" overwriting existing file')
+            os.remove(fn)
+        else:
+            msg = f'LoRA extract: fn="{fn}" file exists'
+            shared.log.warning(msg)
+            yield msg
+            return
 
-    pathstr = str(os.path.join(shared.cmd_opts.lora_dir, filename+f'_{"_".join(suffix)}.safetensors'))
-    save_file(lora_state_dict, pathstr)
-    shared.log.info(f'LoRA extra: fn={pathstr} in {time.time()-t0} seconds')
+    shared.state.end()
+    meta = make_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: 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)} time={t5-t0:.2f}'
+    shared.log.info(msg)
+    yield msg
 
 
 def create_ui():
@@ -168,16 +208,19 @@ def create_ui():
         with gr.Row():
             loaded = gr.Textbox(value="Press refresh to query loaded LoRA", label="Loaded LoRA", interactive=False)
             create_refresh_button(loaded, lambda: None, lambda: {'value': loaded_lora()}, "testid")
-        with gr.Row():
-            rank = gr.Slider(label="Maximum rank", value=32, minimum=1, maximum=256)
-        with gr.Row():
-            auto_rank = gr.Checkbox(value=False, label="Automatically determine rank")
-        with gr.Row(visible=False) as rank_options:
-            rank_ratio = gr.Slider(label="Autorank ratio", value=1, minimum=0, maximum=1, step=0.05, visible=True)
+        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_options])
-    extract.click(fn=make_lora, inputs=[filename, rank, auto_rank, rank_ratio], outputs=[])
+        auto_rank.change(fn=lambda x: gr_show(x), inputs=[auto_rank], outputs=[rank_ratio])
+        extract.click(fn=make_lora, inputs=[filename, rank, auto_rank, rank_ratio, modules, overwrite], outputs=[status])
diff --git a/extensions-builtin/sdnext-modernui b/extensions-builtin/sdnext-modernui
index 74e12fb5e..8afbad75d 160000
--- a/extensions-builtin/sdnext-modernui
+++ b/extensions-builtin/sdnext-modernui
@@ -1 +1 @@
-Subproject commit 74e12fb5e5e67c8a219ff1c4bf1cf6f986e0740e
+Subproject commit 8afbad75d6cd238270111ec77ff19b567855d8bd
diff --git a/modules/devices.py b/modules/devices.py
index 17e6c8f0f..a4b15f412 100644
--- a/modules/devices.py
+++ b/modules/devices.py
@@ -508,12 +508,12 @@ def test_for_nans(x, where):
     raise NansException(message)
 
 
-def normalize_device(device):
-    if torch.device(device).type in {"cpu", "mps", "meta"}:
-        return torch.device(device)
-    if torch.device(device).index is None:
-        return torch.device(str(device), index=0)
-    return torch.device(device)
+def normalize_device(dev):
+    if torch.device(dev).type in {"cpu", "mps", "meta"}:
+        return torch.device(dev)
+    if torch.device(dev).index is None:
+        return torch.device(str(dev), index=0)
+    return torch.device(dev)
 
 
 def same_device(d1, d2):
diff --git a/modules/ui_sections.py b/modules/ui_sections.py
index 739286bd1..a38037392 100644
--- a/modules/ui_sections.py
+++ b/modules/ui_sections.py
@@ -270,7 +270,7 @@ def create_sampler_options(tabname):
             sampler_options = gr.CheckboxGroup(label='Sampler options', elem_id=f"{tabname}_sampler_options", choices=options, value=values, type='value')
         with gr.Row(elem_classes=['flex-break']):
             shared.opts.data['schedulers_sigma'] = shared.opts.data.get('schedulers_sigma', 'default')
-            sampler_algo = gr.Radio(label='Sigma algorithm', elem_id=f"{tabname}_sigma_algo", choices=['default', 'karras', 'exponential', 'polyexponential'], value=shared.opts.data['schedulers_sigma'], type='value')
+            sampler_algo = gr.Radio(label='Sigma algorithm', elem_id=f"{tabname}_sigma_algo", choices=['default', 'karras', 'exponential', 'polyexponential'], value=shared.opts.data.schedulers_sigma, type='value')
         sampler_options.change(fn=set_sampler_original_options, inputs=[sampler_options, sampler_algo], outputs=[])
         sampler_algo.change(fn=set_sampler_original_options, inputs=[sampler_options, sampler_algo], outputs=[])
 

From a5185929e0372c1bd30910be2ada010f8fb5094b Mon Sep 17 00:00:00 2001
From: Vladimir Mandic <mandic00@live.com>
Date: Sun, 20 Oct 2024 10:03:56 -0400
Subject: [PATCH 5/7] metadata

Signed-off-by: Vladimir Mandic <mandic00@live.com>
---
 extensions-builtin/Lora/lora_extract.py | 38 +++++++++++++++++++++----
 1 file changed, 32 insertions(+), 6 deletions(-)

diff --git a/extensions-builtin/Lora/lora_extract.py b/extensions-builtin/Lora/lora_extract.py
index de6404955..7fa8cab6d 100644
--- a/extensions-builtin/Lora/lora_extract.py
+++ b/extensions-builtin/Lora/lora_extract.py
@@ -1,5 +1,7 @@
 import os
 import time
+import json
+import datetime
 import torch
 from safetensors.torch import save_file
 import gradio as gr
@@ -75,13 +77,36 @@ def loaded_lora():
             if current is not None:
                 current = [item[0] for item in current]
                 loaded.update(current)
-    return ", ".join(list(loaded))
+    return list(loaded)
 
 
-def make_meta():
-    return {
-        'todo': 'extra-lora'
+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.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):
@@ -184,7 +209,8 @@ def make_lora(fn, maxrank, auto_rank, rank_ratio, modules, overwrite):
             return
 
     shared.state.end()
-    meta = make_meta()
+    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:
@@ -207,7 +233,7 @@ def create_ui():
     with gr.Tab(label="Extract LoRA"):
         with gr.Row():
             loaded = gr.Textbox(value="Press refresh to query loaded LoRA", label="Loaded LoRA", interactive=False)
-            create_refresh_button(loaded, lambda: None, lambda: {'value': loaded_lora()}, "testid")
+            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'])

From 229cf67bf08ffe68967f158721bfbd413fae9ad6 Mon Sep 17 00:00:00 2001
From: Vladimir Mandic <mandic00@live.com>
Date: Sun, 20 Oct 2024 10:42:56 -0400
Subject: [PATCH 6/7] locking,progressbar

Signed-off-by: Vladimir Mandic <mandic00@live.com>
---
 extensions-builtin/Lora/lora_extract.py | 152 ++++++++++++++----------
 1 file changed, 86 insertions(+), 66 deletions(-)

diff --git a/extensions-builtin/Lora/lora_extract.py b/extensions-builtin/Lora/lora_extract.py
index 7fa8cab6d..761bf614f 100644
--- a/extensions-builtin/Lora/lora_extract.py
+++ b/extensions-builtin/Lora/lora_extract.py
@@ -7,6 +7,7 @@ from safetensors.torch import save_file
 import gradio as gr
 from modules import shared, devices
 from modules.ui_common import create_refresh_button
+from modules.call_queue import wrap_gradio_gpu_call
 
 
 class SVDHandler:
@@ -92,6 +93,7 @@ def make_meta(fn, maxrank, rank_ratio):
         "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}",
     }
@@ -128,71 +130,86 @@ def make_lora(fn, maxrank, auto_rank, rank_ratio, modules, overwrite):
     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')
 
-    shared.log.debug(f'LoRA extract: modules={modules} maxrank={maxrank} auto={auto_rank} ratio={rank_ratio} fn="{fn}"')
-    if 'te' in modules and getattr(shared.sd_model, 'text_encoder', None) is not None:
-        yield "LoRA extract: extracting TE-1"
-        for name, module in shared.sd_model.text_encoder.named_modules():
-            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)
-    t1 = time.time()
+    # bar_format='Progress {rate_fmt}{postfix} {bar} {percentage:3.0f}% {n_fmt}/{total_fmt} {elapsed} {remaining} ' + '\x1b[38;5;71m' + desc, ncols=80, colour='#327fba'
+    from rich.progress import Progress, TextColumn, BarColumn, TaskProgressColumn, TimeRemainingColumn, TimeElapsedColumn
+    with Progress(TextColumn('[cyan]LoRA extract'), BarColumn(), TaskProgressColumn(), TimeRemainingColumn(), TimeElapsedColumn(), TextColumn('[cyan]{task.description}'), console=shared.console) as progress:
 
-    if 'te' in modules and getattr(shared.sd_model, 'text_encoder_2', None) is not None:
-        yield "LoRA extract: extracting TE-2"
-        for name, module in shared.sd_model.text_encoder_2.named_modules():
-            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)
-    t2 = time.time()
-
-    if 'unet' in modules and getattr(shared.sd_model, 'unet', None) is not None:
-        yield "LoRA extract: extracting UNet"
-        for name, module in shared.sd_model.unet.named_modules():
-            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)
-    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:
-            yield f"LoRA extract: creating {sub}"
-            for _name, module in submodel.named_modules():
-                if not hasattr(module, "svdhandler"):
+        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
-                lora_state_dict.update(module.svdhandler.makeweights())
-                del module.svdhandler
-    shared.log.debug('LoRA extract: create done')
-    t4 = time.time()
+                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)
@@ -200,7 +217,6 @@ def make_lora(fn, maxrank, auto_rank, rank_ratio, modules, overwrite):
         fn += '.safetensors'
     if os.path.exists(fn):
         if overwrite:
-            shared.log.warning(f'LoRA extract: fn="{fn}" overwriting existing file')
             os.remove(fn)
         else:
             msg = f'LoRA extract: fn="{fn}" file exists'
@@ -219,9 +235,9 @@ def make_lora(fn, maxrank, auto_rank, rank_ratio, modules, overwrite):
         yield msg
         return
     t5 = time.time()
-    shared.log.debug(f'LoRA extract: te1={t1-t0:.2f} te2={t2-t1:.2f} unet={t3-t2:.2f} save={t5-t4:.2f}')
+    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)} time={t5-t0:.2f}'
+    msg = f'LoRA extract: fn="{fn}" keys={len(keys)}'
     shared.log.info(msg)
     yield msg
 
@@ -232,7 +248,7 @@ def create_ui():
 
     with gr.Tab(label="Extract LoRA"):
         with gr.Row():
-            loaded = gr.Textbox(value="Press refresh to query loaded LoRA", label="Loaded LoRA", interactive=False)
+            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():
@@ -249,4 +265,8 @@ def create_ui():
             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=make_lora, inputs=[filename, rank, auto_rank, rank_ratio, modules, overwrite], outputs=[status])
+        extract.click(
+            fn=wrap_gradio_gpu_call(make_lora, extra_outputs=[]),
+            inputs=[filename, rank, auto_rank, rank_ratio, modules, overwrite],
+            outputs=[status]
+        )

From a13f1ee63a14a85cc3b9e29449035374b385046c Mon Sep 17 00:00:00 2001
From: Vladimir Mandic <mandic00@live.com>
Date: Sun, 20 Oct 2024 10:46:23 -0400
Subject: [PATCH 7/7] cleanup

Signed-off-by: Vladimir Mandic <mandic00@live.com>
---
 extensions-builtin/sdnext-modernui | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/extensions-builtin/sdnext-modernui b/extensions-builtin/sdnext-modernui
index 8afbad75d..895addec9 160000
--- a/extensions-builtin/sdnext-modernui
+++ b/extensions-builtin/sdnext-modernui
@@ -1 +1 @@
-Subproject commit 8afbad75d6cd238270111ec77ff19b567855d8bd
+Subproject commit 895addec9ef65498ed44311d27db0adf699e512d