factor out stuff from load_lora_adapter().

Author: sayakpaul

src/diffusers/loaders/peft.py

@@ -29,13 +29,18 @@
     convert_unet_state_dict_to_peft,
     delete_adapter_layers,
     get_adapter_name,
-    get_peft_kwargs,
     is_peft_available,
     is_peft_version,
     logging,
     set_adapter_layers,
     set_weights_and_activate_adapters,
 )
+from ..utils.peft_utils import (
+    _create_lora_config,
+    _lora_loading_context,
+    _maybe_warn_for_unhandled_keys,
+    _maybe_warn_if_no_keys_found,
+)
 from .lora_base import _fetch_state_dict, _func_optionally_disable_offloading
 from .unet_loader_utils import _maybe_expand_lora_scales
 
@@ -64,26 +69,6 @@
 }
 
 
-def _maybe_raise_error_for_ambiguity(config):
-    rank_pattern = config["rank_pattern"].copy()
-    target_modules = config["target_modules"]
-
-    for key in list(rank_pattern.keys()):
-        # try to detect ambiguity
-        # `target_modules` can also be a str, in which case this loop would loop
-        # over the chars of the str. The technically correct way to match LoRA keys
-        # in PEFT is to use LoraModel._check_target_module_exists (lora_config, key).
-        # But this cuts it for now.
-        exact_matches = [mod for mod in target_modules if mod == key]
-        substring_matches = [mod for mod in target_modules if key in mod and mod != key]
-
-        if exact_matches and substring_matches:
-            if is_peft_version("<", "0.14.1"):
-                raise ValueError(
-                    "There are ambiguous keys present in this LoRA. To load it, please update your `peft` installation - `pip install -U peft`."
-                )
-
-
 class PeftAdapterMixin:
     """
     A class containing all functions for loading and using adapters weights that are supported in PEFT library. For
@@ -189,7 +174,7 @@ def load_lora_adapter(
                 https://huggingface.co/docs/peft/main/en/package_reference/hotswap
             metadata: TODO
         """
-        from peft import LoraConfig, inject_adapter_in_model, set_peft_model_state_dict
+        from peft import inject_adapter_in_model, set_peft_model_state_dict
         from peft.tuners.tuners_utils import BaseTunerLayer
 
         cache_dir = kwargs.pop("cache_dir", None)
@@ -214,7 +199,6 @@ def load_lora_adapter(
             )
 
         user_agent = {"file_type": "attn_procs_weights", "framework": "pytorch"}
-
         state_dict, metadata = _fetch_state_dict(
             pretrained_model_name_or_path_or_dict=pretrained_model_name_or_path_or_dict,
             weight_name=weight_name,
@@ -273,38 +257,8 @@ def load_lora_adapter(
                     k.removeprefix(f"{prefix}."): v for k, v in network_alphas.items() if k in alpha_keys
                 }
 
-            if metadata is not None:
-                lora_config_kwargs = metadata
-            else:
-                lora_config_kwargs = get_peft_kwargs(
-                    rank, network_alpha_dict=network_alphas, peft_state_dict=state_dict
-                )
-            _maybe_raise_error_for_ambiguity(lora_config_kwargs)
-
-            if "use_dora" in lora_config_kwargs:
-                if lora_config_kwargs["use_dora"]:
-                    if is_peft_version("<", "0.9.0"):
-                        raise ValueError(
-                            "You need `peft` 0.9.0 at least to use DoRA-enabled LoRAs. Please upgrade your installation of `peft`."
-                        )
-                else:
-                    if is_peft_version("<", "0.9.0"):
-                        lora_config_kwargs.pop("use_dora")
-
-            if "lora_bias" in lora_config_kwargs:
-                if lora_config_kwargs["lora_bias"]:
-                    if is_peft_version("<=", "0.13.2"):
-                        raise ValueError(
-                            "You need `peft` 0.14.0 at least to use `lora_bias` in LoRAs. Please upgrade your installation of `peft`."
-                        )
-                else:
-                    if is_peft_version("<=", "0.13.2"):
-                        lora_config_kwargs.pop("lora_bias")
-
-            try:
-                lora_config = LoraConfig(**lora_config_kwargs)
-            except TypeError as e:
-                raise TypeError("`LoraConfig` class could not be instantiated.") from e
+            # create LoraConfig
+            lora_config = _create_lora_config(state_dict, network_alphas, metadata, rank, self.lora_layer_modules)
 
             # adapter_name
             if adapter_name is None:
@@ -315,132 +269,98 @@ def load_lora_adapter(
             # Now we remove any existing hooks to `_pipeline`.
 
             # In case the pipeline has been already offloaded to CPU - temporarily remove the hooks
-            # otherwise loading LoRA weights will lead to an error
-            is_model_cpu_offload, is_sequential_cpu_offload = self._optionally_disable_offloading(_pipeline)
-
-            peft_kwargs = {}
-            if is_peft_version(">=", "0.13.1"):
-                peft_kwargs["low_cpu_mem_usage"] = low_cpu_mem_usage
-
-            if hotswap or (self._prepare_lora_hotswap_kwargs is not None):
-                if is_peft_version(">", "0.14.0"):
-                    from peft.utils.hotswap import (
-                        check_hotswap_configs_compatible,
-                        hotswap_adapter_from_state_dict,
-                        prepare_model_for_compiled_hotswap,
-                    )
-                else:
-                    msg = (
-                        "Hotswapping requires PEFT > v0.14. Please upgrade PEFT to a higher version or install it "
-                        "from source."
-                    )
-                    raise ImportError(msg)
-
-            if hotswap:
-
-                def map_state_dict_for_hotswap(sd):
-                    # For hotswapping, we need the adapter name to be present in the state dict keys
-                    new_sd = {}
-                    for k, v in sd.items():
-                        if k.endswith("lora_A.weight") or key.endswith("lora_B.weight"):
-                            k = k[: -len(".weight")] + f".{adapter_name}.weight"
-                        elif k.endswith("lora_B.bias"):  # lora_bias=True option
-                            k = k[: -len(".bias")] + f".{adapter_name}.bias"
-                        new_sd[k] = v
-                    return new_sd
-
-            # To handle scenarios where we cannot successfully set state dict. If it's unsuccessful,
-            # we should also delete the `peft_config` associated to the `adapter_name`.
-            try:
-                if hotswap:
-                    state_dict = map_state_dict_for_hotswap(state_dict)
-                    check_hotswap_configs_compatible(self.peft_config[adapter_name], lora_config)
-                    try:
-                        hotswap_adapter_from_state_dict(
-                            model=self,
-                            state_dict=state_dict,
-                            adapter_name=adapter_name,
-                            config=lora_config,
+            # otherwise loading LoRA weights will lead to an error. So, we use a context manager here
+            # that takes care of enabling and disabling offloading in the pipeline automatically.
+            with _lora_loading_context(_pipeline):
+                peft_kwargs = {}
+                if is_peft_version(">=", "0.13.1"):
+                    peft_kwargs["low_cpu_mem_usage"] = low_cpu_mem_usage
+
+                if hotswap or (self._prepare_lora_hotswap_kwargs is not None):
+                    if is_peft_version(">", "0.14.0"):
+                        from peft.utils.hotswap import (
+                            check_hotswap_configs_compatible,
+                            hotswap_adapter_from_state_dict,
+                            prepare_model_for_compiled_hotswap,
                         )
-                    except Exception as e:
-                        logger.error(f"Hotswapping {adapter_name} was unsuccessful with the following error: \n{e}")
-                        raise
-                    # the hotswap function raises if there are incompatible keys, so if we reach this point we can set
-                    # it to None
-                    incompatible_keys = None
-                else:
-                    inject_adapter_in_model(lora_config, self, adapter_name=adapter_name, **peft_kwargs)
-                    incompatible_keys = set_peft_model_state_dict(self, state_dict, adapter_name, **peft_kwargs)
-
-                    if self._prepare_lora_hotswap_kwargs is not None:
-                        # For hotswapping of compiled models or adapters with different ranks.
-                        # If the user called enable_lora_hotswap, we need to ensure it is called:
-                        # - after the first adapter was loaded
-                        # - before the model is compiled and the 2nd adapter is being hotswapped in
-                        # Therefore, it needs to be called here
-                        prepare_model_for_compiled_hotswap(
-                            self, config=lora_config, **self._prepare_lora_hotswap_kwargs
-                        )
-                        # We only want to call prepare_model_for_compiled_hotswap once
-                        self._prepare_lora_hotswap_kwargs = None
-
-                # Set peft config loaded flag to True if module has been successfully injected and incompatible keys retrieved
-                if not self._hf_peft_config_loaded:
-                    self._hf_peft_config_loaded = True
-            except Exception as e:
-                # In case `inject_adapter_in_model()` was unsuccessful even before injecting the `peft_config`.
-                if hasattr(self, "peft_config"):
-                    for module in self.modules():
-                        if isinstance(module, BaseTunerLayer):
-                            active_adapters = module.active_adapters
-                            for active_adapter in active_adapters:
-                                if adapter_name in active_adapter:
-                                    module.delete_adapter(adapter_name)
-
-                    self.peft_config.pop(adapter_name)
-                logger.error(f"Loading {adapter_name} was unsuccessful with the following error: \n{e}")
-                raise
-
-            warn_msg = ""
-            if incompatible_keys is not None:
-                # Check only for unexpected keys.
-                unexpected_keys = getattr(incompatible_keys, "unexpected_keys", None)
-                if unexpected_keys:
-                    lora_unexpected_keys = [k for k in unexpected_keys if "lora_" in k and adapter_name in k]
-                    if lora_unexpected_keys:
-                        warn_msg = (
-                            f"Loading adapter weights from state_dict led to unexpected keys found in the model:"
-                            f" {', '.join(lora_unexpected_keys)}. "
-                        )
-
-                # Filter missing keys specific to the current adapter.
-                missing_keys = getattr(incompatible_keys, "missing_keys", None)
-                if missing_keys:
-                    lora_missing_keys = [k for k in missing_keys if "lora_" in k and adapter_name in k]
-                    if lora_missing_keys:
-                        warn_msg += (
-                            f"Loading adapter weights from state_dict led to missing keys in the model:"
-                            f" {', '.join(lora_missing_keys)}."
+                    else:
+                        msg = (
+                            "Hotswapping requires PEFT > v0.14. Please upgrade PEFT to a higher version or install it "
+                            "from source."
                         )
+                        raise ImportError(msg)
 
-            if warn_msg:
-                logger.warning(warn_msg)
+                if hotswap:
 
-            # Offload back.
-            if is_model_cpu_offload:
-                _pipeline.enable_model_cpu_offload()
-            elif is_sequential_cpu_offload:
-                _pipeline.enable_sequential_cpu_offload()
+                    def map_state_dict_for_hotswap(sd):
+                        # For hotswapping, we need the adapter name to be present in the state dict keys
+                        new_sd = {}
+                        for k, v in sd.items():
+                            if k.endswith("lora_A.weight") or key.endswith("lora_B.weight"):
+                                k = k[: -len(".weight")] + f".{adapter_name}.weight"
+                            elif k.endswith("lora_B.bias"):  # lora_bias=True option
+                                k = k[: -len(".bias")] + f".{adapter_name}.bias"
+                            new_sd[k] = v
+                        return new_sd
+
+                # To handle scenarios where we cannot successfully set state dict. If it's unsuccessful,
+                # we should also delete the `peft_config` associated to the `adapter_name`.
+                try:
+                    if hotswap:
+                        state_dict = map_state_dict_for_hotswap(state_dict)
+                        check_hotswap_configs_compatible(self.peft_config[adapter_name], lora_config)
+                        try:
+                            hotswap_adapter_from_state_dict(
+                                model=self,
+                                state_dict=state_dict,
+                                adapter_name=adapter_name,
+                                config=lora_config,
+                            )
+                        except Exception as e:
+                            logger.error(
+                                f"Hotswapping {adapter_name} was unsuccessful with the following error: \n{e}"
+                            )
+                            raise
+                        # the hotswap function raises if there are incompatible keys, so if we reach this point we can set
+                        # it to None
+                        incompatible_keys = None
+                    else:
+                        inject_adapter_in_model(lora_config, self, adapter_name=adapter_name, **peft_kwargs)
+                        incompatible_keys = set_peft_model_state_dict(self, state_dict, adapter_name, **peft_kwargs)
+
+                        if self._prepare_lora_hotswap_kwargs is not None:
+                            # For hotswapping of compiled models or adapters with different ranks.
+                            # If the user called enable_lora_hotswap, we need to ensure it is called:
+                            # - after the first adapter was loaded
+                            # - before the model is compiled and the 2nd adapter is being hotswapped in
+                            # Therefore, it needs to be called here
+                            prepare_model_for_compiled_hotswap(
+                                self, config=lora_config, **self._prepare_lora_hotswap_kwargs
+                            )
+                            # We only want to call prepare_model_for_compiled_hotswap once
+                            self._prepare_lora_hotswap_kwargs = None
+
+                    # Set peft config loaded flag to True if module has been successfully injected and incompatible keys retrieved
+                    if not self._hf_peft_config_loaded:
+                        self._hf_peft_config_loaded = True
+                except Exception as e:
+                    # In case `inject_adapter_in_model()` was unsuccessful even before injecting the `peft_config`.
+                    if hasattr(self, "peft_config"):
+                        for module in self.modules():
+                            if isinstance(module, BaseTunerLayer):
+                                active_adapters = module.active_adapters
+                                for active_adapter in active_adapters:
+                                    if adapter_name in active_adapter:
+                                        module.delete_adapter(adapter_name)
+
+                        self.peft_config.pop(adapter_name)
+                    logger.error(f"Loading {adapter_name} was unsuccessful with the following error: \n{e}")
+                    raise
             # Unsafe code />
 
-        if prefix is not None and not state_dict:
-            logger.warning(
-                f"No LoRA keys associated to {self.__class__.__name__} found with the {prefix=}. "
-                "This is safe to ignore if LoRA state dict didn't originally have any "
-                f"{self.__class__.__name__} related params. You can also try specifying `prefix=None` "
-                "to resolve the warning. Otherwise, open an issue if you think it's unexpected: "
-                "https://github.com/huggingface/diffusers/issues/new"
-            )
+            _maybe_warn_for_unhandled_keys(incompatible_keys, adapter_name)
+
+        _maybe_warn_if_no_keys_found(state_dict, prefix, model_class_name=self.__class__.__name__)
 
     def save_lora_adapter(
         self,