diff --git a/extensions-builtin/Lora/scripts/lora_script.py b/extensions-builtin/Lora/scripts/lora_script.py index ffbef47d9..dea2985b3 100644 --- a/extensions-builtin/Lora/scripts/lora_script.py +++ b/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(" 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): + t0 = time.time() + 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.load_networks(names, te_multipliers, unet_multipliers, dyn_dims) + t1 = time.time() + 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} dims={dyn_dims} load={t1-t0:.2f}') + + 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 + 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() diff --git a/modules/lora/lora.py b/modules/lora/lora.py new file mode 100644 index 000000000..33adfe05c --- /dev/null +++ b/modules/lora/lora.py @@ -0,0 +1,8 @@ +# import networks +# +# list_available_loras = networks.list_available_networks +# available_loras = networks.available_networks +# available_lora_aliases = networks.available_network_aliases +# available_lora_hash_lookup = networks.available_network_hash_lookup +# forbidden_lora_aliases = networks.forbidden_network_aliases +# loaded_loras = networks.loaded_networks diff --git a/modules/lora/lora_convert.py b/modules/lora/lora_convert.py new file mode 100644 index 000000000..6bf563125 --- /dev/null +++ b/modules/lora/lora_convert.py @@ -0,0 +1,477 @@ +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) + + 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) diff --git a/modules/lora/lora_extract.py b/modules/lora/lora_extract.py new file mode 100644 index 000000000..c2e0a275b --- /dev/null +++ b/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] + ) diff --git a/modules/lora/lyco_helpers.py b/modules/lora/lyco_helpers.py new file mode 100644 index 000000000..9a16d25ab --- /dev/null +++ b/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) + + +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) + + +# 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 length or new_m>factor: + break + m, n = new_m, new_n + if m > n: + n, m = m, n + return m, n diff --git a/modules/lora/network.py b/modules/lora/network.py new file mode 100644 index 000000000..0785ef9f4 --- /dev/null +++ b/modules/lora/network.py @@ -0,0 +1,187 @@ +import os +from collections import namedtuple +import enum + +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) -> 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 diff --git a/modules/lora/network_full.py b/modules/lora/network_full.py new file mode 100644 index 000000000..5eb0b2e4e --- /dev/null +++ b/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) diff --git a/modules/lora/network_glora.py b/modules/lora/network_glora.py new file mode 100644 index 000000000..ffcb25986 --- /dev/null +++ b/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) diff --git a/modules/lora/network_hada.py b/modules/lora/network_hada.py new file mode 100644 index 000000000..6fc142b3b --- /dev/null +++ b/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) diff --git a/modules/lora/network_ia3.py b/modules/lora/network_ia3.py new file mode 100644 index 000000000..479e42526 --- /dev/null +++ b/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) diff --git a/modules/lora/network_lokr.py b/modules/lora/network_lokr.py new file mode 100644 index 000000000..877d4005b --- /dev/null +++ b/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) diff --git a/modules/lora/network_lora.py b/modules/lora/network_lora.py new file mode 100644 index 000000000..6c1d7ea3f --- /dev/null +++ b/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() diff --git a/modules/lora/network_norm.py b/modules/lora/network_norm.py new file mode 100644 index 000000000..e8f1740e3 --- /dev/null +++ b/modules/lora/network_norm.py @@ -0,0 +1,23 @@ +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) diff --git a/modules/lora/network_oft.py b/modules/lora/network_oft.py new file mode 100644 index 000000000..808286066 --- /dev/null +++ b/modules/lora/network_oft.py @@ -0,0 +1,81 @@ +import torch +import modules.lora.network as network +from modules.lora.lyco_helpers import factorization +from einops import rearrange + + +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) diff --git a/modules/lora/network_overrides.py b/modules/lora/network_overrides.py new file mode 100644 index 000000000..5334f3c1b --- /dev/null +++ b/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 diff --git a/modules/lora/networks.py b/modules/lora/networks.py new file mode 100644 index 000000000..762705b67 --- /dev/null +++ b/modules/lora/networks.py @@ -0,0 +1,453 @@ +from typing import Union, List +import os +import re +import time +import concurrent +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 +import torch +import diffusers.models.lora +from modules import shared, devices, sd_models, sd_models_compile, errors, scripts, files_cache, model_quant + + +debug = os.environ.get('SD_LORA_DEBUG', None) is not None +extra_network_lora = None +available_networks = {} +available_network_aliases = {} +loaded_networks: List[network.Network] = [] +timer = { 'load': 0, 'apply': 0, 'restore': 0, 'deactivate': 0 } +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 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 + + +def load_diffusers(name, network_on_disk, lora_scale=shared.opts.extra_networks_default_multiplier) -> 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_network(name, network_on_disk) -> 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 + 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 + 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 load_networks(names, te_multipliers=None, unet_multipliers=None, dyn_dims=None): + 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() + timer['load'] = 0 + 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_network(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)}') + + devices.torch_gc() + + 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 shared.opts.diffusers_offload_mode == "balanced": + sd_models.apply_balanced_offload(shared.sd_model) + t1 = time.time() + timer['load'] += t1 - t0 + +def set_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): + 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 + device = self.weight.device + with devices.inference_context(): + if weights_backup is not None: + if updown is not None: + if 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 + weights_backup = weights_backup.clone().to(device) + weights_backup += updown.to(weights_backup) + if getattr(self, "quant_type", None) in ['nf4', 'fp4']: + bnb = model_quant.load_bnb('Load network: type=LoRA', silent=True) + if bnb is not None: + self.weight = bnb.nn.Params4bit(weights_backup, quant_state=self.quant_state, quant_type=self.quant_type, blocksize=self.blocksize) + else: + self.weight.copy_(weights_backup, non_blocking=True) + else: + self.weight.copy_(weights_backup, non_blocking=True) + if hasattr(self, "qweight") and hasattr(self, "freeze"): + self.freeze() + if bias_backup is not None: + if ex_bias is not None: + bias_backup = bias_backup.clone() + ex_bias.to(weights_backup) + self.bias.copy_(bias_backup) + else: + self.bias = None + self.to(device) + + +def maybe_backup_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.GroupNorm, torch.nn.LayerNorm, diffusers.models.lora.LoRACompatibleLinear, diffusers.models.lora.LoRACompatibleConv], wanted_names): # pylint: disable=W0613 + weights_backup = getattr(self, "network_weights_backup", None) + if weights_backup is None and wanted_names != (): # pylint: disable=C1803 + if getattr(self.weight, "quant_type", None) in ['nf4', 'fp4']: + 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(self.weight, quant_state=self.weight.quant_state, quant_type=self.weight.quant_type, blocksize=self.weight.blocksize,) + self.quant_state = self.weight.quant_state + self.quant_type = self.weight.quant_type + self.blocksize = self.weight.blocksize + else: + weights_backup = self.weight.clone() + else: + weights_backup = self.weight.clone() + if shared.opts.lora_offload_backup and weights_backup is not None: + 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: + bias_backup = self.bias.clone() + else: + bias_backup = None + if shared.opts.lora_offload_backup and bias_backup is not None: + bias_backup = bias_backup.to(devices.cpu) + self.network_bias_backup = bias_backup + + +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]): + """ + 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]): # noqa: C419 # pylint: disable=R1729 + maybe_backup_weights(self, wanted_names) + if current_names != wanted_names: + for net in loaded_networks: + # default workflow where module is known and has weights + module = net.modules.get(network_layer_name, None) + if module is not None and hasattr(self, 'weight'): + try: + with devices.inference_context(): + weight = self.weight # calculate quant weights once + updown, ex_bias = module.calc_updown(weight) + set_weights(self, updown, ex_bias) + 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 + if not loaded_networks: # restore from backup + t5 = time.time() + set_weights(self, None, None) + self.network_current_names = wanted_names + t1 = time.time() + timer['apply'] += t1 - t0 + +def network_load(): + sd_model = getattr(shared.sd_model, "pipe", shared.sd_model) # wrapped model compatiblility + for component_name in ['text_encoder','text_encoder_2', 'unet', 'transformer']: + component = getattr(sd_model, component_name, None) + if component is not None: + for _, module in component.named_modules(): + network_apply_weights(module) + + +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() + 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 infotext_pasted(infotext, params): # pylint: disable=W0613 + if "AddNet Module 1" in [x[1] for x in scripts.scripts_txt2img.infotext_fields]: + return # if the other extension is active, it will handle those fields, no need to do anything + added = [] + for k in params: + if not k.startswith("AddNet Model "): + continue + num = k[13:] + if params.get("AddNet Module " + num) != "LoRA": + continue + name = params.get("AddNet Model " + num) + if name is None: + continue + m = re_network_name.match(name) + if m: + name = m.group(1) + multiplier = params.get("AddNet Weight A " + num, "1.0") + added.append(f"") + if added: + params["Prompt"] += "\n" + "".join(added) + + +list_available_networks() diff --git a/modules/lora/ui_extra_networks_lora.py b/modules/lora/ui_extra_networks_lora.py new file mode 100644 index 000000000..73cce47a3 --- /dev/null +++ b/modules/lora/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" "), + "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, shared.cmd_opts.lyco_dir] diff --git a/modules/processing_diffusers.py b/modules/processing_diffusers.py index 2164134b1..83d3b1b69 100644 --- a/modules/processing_diffusers.py +++ b/modules/processing_diffusers.py @@ -8,6 +8,8 @@ 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.networks import network_load +from modules.lora.networks import timer as network_timer debug = shared.log.trace if os.environ.get('SD_DIFFUSERS_DEBUG', None) is not None else lambda *args, **kwargs: None @@ -424,6 +426,9 @@ def process_diffusers(p: processing.StableDiffusionProcessing): p.prompts = p.all_prompts[p.iteration * p.batch_size:(p.iteration+1) * p.batch_size] 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] + network_timer['apply'] = 0 + network_timer['restore'] = 0 + network_load() 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 diff --git a/modules/shared.py b/modules/shared.py index a89cbbc95..a5af83f5e 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -908,6 +908,7 @@ options_templates.update(options_section(('extra_networks', "Networks"), { "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"), + "lora_offload_backup": OptionInfo(True, "Offload LoRA Backup Weights"), })) options_templates.update(options_section((None, "Internal options"), { diff --git a/scripts/lora_script.py b/scripts/lora_script.py new file mode 100644 index 000000000..a153a2caa --- /dev/null +++ b/scripts/lora_script.py @@ -0,0 +1,62 @@ +import re +import modules.lora.networks as networks +from modules.lora.lora_extract import create_ui +from modules.lora.network import NetworkOnDisk +from modules.lora.ui_extra_networks_lora import ExtraNetworksPageLora +from modules.lora.extra_networks_lora import ExtraNetworkLora +from modules import script_callbacks, extra_networks, ui_extra_networks, ui_models, shared # pylint: disable=unused-import + + +re_lora = re.compile("