diff --git a/extensions-builtin/Lora/lora_extract.py b/extensions-builtin/Lora/lora_extract.py new file mode 100644 index 000000000..761bf614f --- /dev/null +++ b/extensions-builtin/Lora/lora_extract.py @@ -0,0 +1,272 @@ +import os +import time +import json +import datetime +import torch +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, 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') + + # 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', 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] + ) 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): diff --git a/extensions-builtin/sdnext-modernui b/extensions-builtin/sdnext-modernui index 74e12fb5e..895addec9 160000 --- a/extensions-builtin/sdnext-modernui +++ b/extensions-builtin/sdnext-modernui @@ -1 +1 @@ -Subproject commit 74e12fb5e5e67c8a219ff1c4bf1cf6f986e0740e +Subproject commit 895addec9ef65498ed44311d27db0adf699e512d 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=[])