From bd5ac8e0489458f8c55d6f521842e4b011e2b848 Mon Sep 17 00:00:00 2001 From: AI-Casanova <54461896+AI-Casanova@users.noreply.github.com> Date: Sun, 29 Sep 2024 15:56:53 -0500 Subject: [PATCH] Add option, fix diffusers keys --- extensions-builtin/Lora/lora_convert.py | 2 ++ extensions-builtin/Lora/networks.py | 21 +++++++++++---------- modules/shared.py | 1 + 3 files changed, 14 insertions(+), 10 deletions(-) diff --git a/extensions-builtin/Lora/lora_convert.py b/extensions-builtin/Lora/lora_convert.py index 15a5e51e1..c64cb9671 100644 --- a/extensions-builtin/Lora/lora_convert.py +++ b/extensions-builtin/Lora/lora_convert.py @@ -175,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) diff --git a/extensions-builtin/Lora/networks.py b/extensions-builtin/Lora/networks.py index c6341d4c7..eb5b888a5 100644 --- a/extensions-builtin/Lora/networks.py +++ b/extensions-builtin/Lora/networks.py @@ -134,8 +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': - sd = lora_convert._convert_kohya_flux_lora_to_diffusers(sd) or sd # if kohya flux lora, convert state_dict + 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 = {} @@ -296,11 +296,11 @@ 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.weight, "quant_type", None) is not None: + 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.weight.quant_state, - quant_type=self.weight.quant_type, blocksize=self.weight.blocksize) + 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) @@ -337,7 +337,6 @@ def network_apply_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn if isinstance(self, torch.nn.MultiheadAttention): 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": - # weights_backup = self.weight.__deepcopy__("") import bitsandbytes with devices.inference_context(): weights_backup = bitsandbytes.functional.dequantize_4bit(self.weight, @@ -345,6 +344,9 @@ def network_apply_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn 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.clone().to(devices.cpu) self.network_weights_backup = weights_backup @@ -378,10 +380,9 @@ def network_apply_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn 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.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) diff --git a/modules/shared.py b/modules/shared.py index 6b0bc0d5f..6f4eb06fb 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -862,6 +862,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"),