diff --git a/extensions-builtin/Lora/lora_convert.py b/extensions-builtin/Lora/lora_convert.py index 8432d8208..c64cb9671 100644 --- a/extensions-builtin/Lora/lora_convert.py +++ b/extensions-builtin/Lora/lora_convert.py @@ -1,6 +1,7 @@ import os import re import bisect +import torch from typing import Dict from modules import shared @@ -174,6 +175,8 @@ class KeyConvert: if search_key.startswith(map_key): key = key.replace(map_key, self.UNET_CONVERSION_MAP[map_key]).replace("oft", "lora") # pylint: disable=unsubscriptable-object 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 # SegMoE begin expert_key = key + "_experts_0" expert_module = shared.sd_model.network_layer_mapping.get(expert_key, None) @@ -253,3 +256,200 @@ def convert_diffusers_name_to_compvis(key, is_sd2): else: return f"1_model_transformer_resblocks_{m[0]}_{m[1].replace('self_attn', 'attn')}" return key + + +# 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_kohya_flux_lora_to_diffusers(state_dict): + 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 + 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 + + # 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_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) diff --git a/extensions-builtin/Lora/lora_patches.py b/extensions-builtin/Lora/lora_patches.py index f31e2cc44..a0f3566ff 100644 --- a/extensions-builtin/Lora/lora_patches.py +++ b/extensions-builtin/Lora/lora_patches.py @@ -20,6 +20,11 @@ class LoraPatches: def apply(self): if self.active or shared.opts.lora_force_diffusers: return + try: + import bitsandbytes + self.Linear4bit_forward = patches.patch(__name__, bitsandbytes.nn.Linear4bit, 'forward', networks.network_Linear4bit_forward) + except: + pass if "Model" in shared.opts.optimum_quanto_weights or "Text Encoder" in shared.opts.optimum_quanto_weights: from optimum import quanto # pylint: disable=no-name-in-module self.QLinear_forward = patches.patch(__name__, quanto.nn.QLinear, 'forward', networks.network_QLinear_forward) # pylint: disable=attribute-defined-outside-init @@ -42,6 +47,11 @@ class LoraPatches: def undo(self): if not self.active or shared.opts.lora_force_diffusers: return + try: + import bitsandbytes + self.Linear4bit_forward = patches.undo(__name__, bitsandbytes.nn.Linear4bit, 'forward') + except: + pass if "Model" in shared.opts.optimum_quanto_weights or "Text Encoder" in shared.opts.optimum_quanto_weights: from optimum import quanto # pylint: disable=no-name-in-module self.QLinear_forward = patches.undo(__name__, quanto.nn.QLinear, 'forward') # pylint: disable=E1128, attribute-defined-outside-init diff --git a/extensions-builtin/Lora/network_lora.py b/extensions-builtin/Lora/network_lora.py index 7dfded536..5e6eaef6c 100644 --- a/extensions-builtin/Lora/network_lora.py +++ b/extensions-builtin/Lora/network_lora.py @@ -27,7 +27,7 @@ class NetworkModuleLora(network.NetworkModule): 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"} + 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) @@ -55,12 +55,13 @@ class NetworkModuleLora(network.NetworkModule): return module def calc_updown(self, target): # pylint: disable=W0237 - up = self.up_model.weight.to(target.device, dtype=target.dtype) - down = self.down_model.weight.to(target.device, dtype=target.dtype) + 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) + 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: diff --git a/extensions-builtin/Lora/network_overrides.py b/extensions-builtin/Lora/network_overrides.py index 9123d0039..24afb0c28 100644 --- a/extensions-builtin/Lora/network_overrides.py +++ b/extensions-builtin/Lora/network_overrides.py @@ -30,7 +30,6 @@ force_models = [ # forced always 'kandinsky', 'hunyuandit', 'auraflow', - 'f1', ] force_classes = [ # forced always diff --git a/extensions-builtin/Lora/networks.py b/extensions-builtin/Lora/networks.py index 6e20b5a59..eb5b888a5 100644 --- a/extensions-builtin/Lora/networks.py +++ b/extensions-builtin/Lora/networks.py @@ -50,23 +50,29 @@ convert_diffusers_name_to_compvis = lora_convert.convert_diffusers_name_to_compv def assign_network_names_to_compvis_modules(sd_model): network_layer_mapping = {} if shared.native: - if not hasattr(shared.sd_model, 'text_encoder') or not hasattr(shared.sd_model, 'unet'): - sd_model.network_layer_mapping = {} - return - for name, module in shared.sd_model.text_encoder.named_modules(): - prefix = "lora_te1_" if shared.sd_model_type == "sdxl" else "lora_te_" - network_name = prefix + name.replace(".", "_") - network_layer_mapping[network_name] = module - module.network_layer_name = network_name - if shared.sd_model_type == "sdxl": + if hasattr(shared.sd_model, 'text_encoder'): + for name, module in shared.sd_model.text_encoder.named_modules(): + prefix = "lora_te1_" if hasattr(shared.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(shared.sd_model, 'text_encoder_2'): for name, module in shared.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 - for name, module in shared.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(shared.sd_model, 'unet'): + for name, module in shared.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(shared.sd_model, 'transformer'): + for name, module in shared.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: + continue + module.network_layer_name = network_name else: if not hasattr(shared.sd_model, 'cond_stage_model'): sd_model.network_layer_mapping = {} @@ -128,6 +134,8 @@ def load_network(name, network_on_disk) -> network.Network: 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 assign_network_names_to_compvis_modules(shared.sd_model) # this should not be needed but is here as an emergency fix for an unknown error people are experiencing in 1.2.0 keys_failed_to_match = {} matched_networks = {} @@ -288,6 +296,12 @@ def network_restore_weights_from_backup(self: Union[torch.nn.Conv2d, torch.nn.Li elif hasattr(self, "qweight") and hasattr(self, "freeze"): self.weight = torch.nn.Parameter(weights_backup.to(self.weight.device, copy=True)) self.freeze() + elif getattr(self, "quant_type", None) is not None: + import bitsandbytes + device = self.weight.device + self.weight = bitsandbytes.nn.Params4bit(weights_backup, quant_state=self.quant_state, + quant_type=self.quant_type, blocksize=self.blocksize) + self.weight.to(device) else: self.weight.copy_(weights_backup) if bias_backup is not None: @@ -321,16 +335,27 @@ def network_apply_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn if current_names != (): raise RuntimeError("no backup weights found and current weights are not unchanged") if isinstance(self, torch.nn.MultiheadAttention): - weights_backup = (self.in_proj_weight.to(devices.cpu, copy=True), self.out_proj.weight.to(devices.cpu, copy=True)) + weights_backup = (self.in_proj_weight.clone().to(devices.cpu), self.out_proj.weight.clone().to(devices.cpu)) + elif getattr(self.weight, "quant_type", None) == "nf4" or getattr(self.weight, "quant_type", None) == "nf4": + import bitsandbytes + with devices.inference_context(): + weights_backup = bitsandbytes.functional.dequantize_4bit(self.weight, + quant_state=self.weight.quant_state, + quant_type=self.weight.quant_type, + blocksize=self.weight.blocksize, + ).to(devices.cpu) + self.quant_state = self.weight.quant_state + self.quant_type = self.weight.quant_type + self.blocksize = self.weight.blocksize else: - weights_backup = self.weight.to(devices.cpu, copy=True) + weights_backup = self.weight.clone().to(devices.cpu) self.network_weights_backup = weights_backup bias_backup = getattr(self, "network_bias_backup", None) if bias_backup is None: if isinstance(self, torch.nn.MultiheadAttention) and self.out_proj.bias is not None: - bias_backup = self.out_proj.bias.to(devices.cpu, copy=True) + bias_backup = self.out_proj.bias.clone().to(devices.cpu) elif getattr(self, 'bias', None) is not None: - bias_backup = self.bias.to(devices.cpu, copy=True) + bias_backup = self.bias.clone().to(devices.cpu) else: bias_backup = None self.network_bias_backup = bias_backup @@ -348,7 +373,19 @@ def network_apply_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn if len(weight.shape) == 4 and weight.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 - self.weight = torch.nn.Parameter(weight + updown) + if getattr(self.weight, "quant_type", None) == "nf4" or self.weight.numel() != updown.numel(): + import bitsandbytes + device = self.weight.device + weight = bitsandbytes.functional.dequantize_4bit(self.weight, + quant_state=self.weight.quant_state, + quant_type=self.weight.quant_type, + blocksize=self.weight.blocksize) + self.weight = bitsandbytes.nn.Params4bit(weight + updown, quant_state=self.quant_state, + quant_type=shared.opts.lora_quant.lower(), + blocksize=self.blocksize) + self.weight.to(device) + else: + self.weight = torch.nn.Parameter(weight + updown) if hasattr(self, "qweight") and hasattr(self, "freeze"): self.freeze() if ex_bias is not None and hasattr(self, 'bias'): @@ -438,6 +475,14 @@ def network_Linear_load_state_dict(self, *args, **kwargs): return originals.Linear_load_state_dict(self, *args, **kwargs) +def network_Linear4bit_forward(self, input): # pylint: disable=W0622 + network_apply_weights(self) + return originals.Linear4bit_forward(self, input) +# +# def network_Linear4bit_load_state_dict(self, *args, **kwargs): +# network_reset_cached_weight(self) +# return originals.Linear4bit_load_state_dict(self, *args, **kwargs) + def network_Conv2d_forward(self, input): # pylint: disable=W0622 network_apply_weights(self) return originals.Conv2d_forward(self, input) diff --git a/modules/shared.py b/modules/shared.py index 0852f3dac..87a75206f 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -861,6 +861,7 @@ options_templates.update(options_section(('extra_networks', "Networks"), { "extra_networks_styles": OptionInfo(True, "Show built-in styles"), "lora_preferred_name": OptionInfo("filename", "LoRA preferred name", gr.Radio, {"choices": ["filename", "alias"]}), "lora_add_hashes_to_infotext": OptionInfo(False, "LoRA add hash info"), + "lora_quant": OptionInfo("FP4","LoRA precision for merged layers in quantized models", gr.Radio, {"choices": ["FP4", "NF4"]}), "lora_force_diffusers": OptionInfo(False if not cmd_opts.use_openvino else True, "LoRA force loading of all models using Diffusers"), "lora_maybe_diffusers": OptionInfo(False, "LoRA force loading of specific models using Diffusers"), "lora_fuse_diffusers": OptionInfo(False if not cmd_opts.use_openvino else True, "LoRA use merge when using alternative method"),