From 2570d87ad3d97620049e82d6da45670ff265dca6 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Sun, 20 Oct 2024 08:30:06 -0400 Subject: [PATCH] typing, linting Signed-off-by: Vladimir Mandic --- 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=[])