Merge pull request #3461 from AI-Casanova/flux-lora

Native Flux nf4 LoRA

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)

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

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:

extensions-builtin/Lora/network_overrides.py

@@ -30,7 +30,6 @@ force_models = [ # forced always
     'kandinsky',
     'hunyuandit',
     'auraflow',
-    'f1',
 ]
 
 force_classes = [ # forced always

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)

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"),