typing, linting

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

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=[])