From 4f67ac2967240a35985a0cc895105b0ed72249b1 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Mon, 1 Jan 2024 14:54:04 +0530 Subject: [PATCH 1/8] introduce integrations module. --- src/diffusers/integrations/__init__.py | 30 ++++ src/diffusers/integrations/peft.py | 182 +++++++++++++++++++++++++ src/diffusers/models/modeling_utils.py | 3 +- 3 files changed, 214 insertions(+), 1 deletion(-) create mode 100644 src/diffusers/integrations/__init__.py create mode 100644 src/diffusers/integrations/peft.py diff --git a/src/diffusers/integrations/__init__.py b/src/diffusers/integrations/__init__.py new file mode 100644 index 000000000000..ec2408f348c5 --- /dev/null +++ b/src/diffusers/integrations/__init__.py @@ -0,0 +1,30 @@ +# Copyright 2023 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import TYPE_CHECKING + +from ..utils import DIFFUSERS_SLOW_IMPORT, _LazyModule, is_peft_available, is_torch_available + + +_import_structure = {} + +if is_torch_available() and is_peft_available(): + _import_structure = {"peft": ["PeftAdapterMixin"]} + +if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: + if is_torch_available() and is_peft_available(): + from .peft import PeftAdapterMixin +else: + import sys + + sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__) diff --git a/src/diffusers/integrations/peft.py b/src/diffusers/integrations/peft.py new file mode 100644 index 000000000000..098b4ee50073 --- /dev/null +++ b/src/diffusers/integrations/peft.py @@ -0,0 +1,182 @@ +# coding=utf-8 +# Copyright 2023 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import List, Union + +from ..utils import MIN_PEFT_VERSION, check_peft_version + + +class PeftAdapterMixin: + """ + A class containing all functions for loading and using adapters weights that are supported in PEFT library. For + more details about adapters and injecting them on a transformer-based model, check out the documentation of PEFT + library: https://huggingface.co/docs/peft/index. + + + With this mixin, if the correct PEFT version is installed, it is possible to: + + - Attach new adapters in the model. + - Attach multiple adapters and iteratively activate / deactivate them. + - Activate / deactivate all adapters from the model. + - Get a list of the active adapters. + """ + + _hf_peft_config_loaded = False + + def add_adapter(self, adapter_config, adapter_name: str = "default") -> None: + r""" + Adds a new adapter to the current model for training. If no adapter name is passed, a default name is assigned + to the adapter to follow the convention of the PEFT library. + + If you are not familiar with adapters and PEFT methods, we invite you to read more about them in the PEFT + [documentation](https://huggingface.co/docs/peft). + + Args: + adapter_config (`[~peft.PeftConfig]`): + The configuration of the adapter to add; supported adapters are non-prefix tuning and adaption prompt + methods. + adapter_name (`str`, *optional*, defaults to `"default"`): + The name of the adapter to add. If no name is passed, a default name is assigned to the adapter. + """ + check_peft_version(min_version=MIN_PEFT_VERSION) + + from peft import PeftConfig, inject_adapter_in_model + + if not self._hf_peft_config_loaded: + self._hf_peft_config_loaded = True + elif adapter_name in self.peft_config: + raise ValueError(f"Adapter with name {adapter_name} already exists. Please use a different name.") + + if not isinstance(adapter_config, PeftConfig): + raise ValueError( + f"adapter_config should be an instance of PeftConfig. Got {type(adapter_config)} instead." + ) + + # Unlike transformers, here we don't need to retrieve the name_or_path of the unet as the loading logic is + # handled by the `load_lora_layers` or `LoraLoaderMixin`. Therefore we set it to `None` here. + adapter_config.base_model_name_or_path = None + inject_adapter_in_model(adapter_config, self, adapter_name) + self.set_adapter(adapter_name) + + def set_adapter(self, adapter_name: Union[str, List[str]]) -> None: + """ + Sets a specific adapter by forcing the model to only use that adapter and disables the other adapters. + + If you are not familiar with adapters and PEFT methods, we invite you to read more about them on the PEFT + official documentation: https://huggingface.co/docs/peft + + Args: + adapter_name (Union[str, List[str]])): + The list of adapters to set or the adapter name in case of single adapter. + """ + check_peft_version(min_version=MIN_PEFT_VERSION) + + if not self._hf_peft_config_loaded: + raise ValueError("No adapter loaded. Please load an adapter first.") + + if isinstance(adapter_name, str): + adapter_name = [adapter_name] + + missing = set(adapter_name) - set(self.peft_config) + if len(missing) > 0: + raise ValueError( + f"Following adapter(s) could not be found: {', '.join(missing)}. Make sure you are passing the correct adapter name(s)." + f" current loaded adapters are: {list(self.peft_config.keys())}" + ) + + from peft.tuners.tuners_utils import BaseTunerLayer + + _adapters_has_been_set = False + + for _, module in self.named_modules(): + if isinstance(module, BaseTunerLayer): + if hasattr(module, "set_adapter"): + module.set_adapter(adapter_name) + # Previous versions of PEFT does not support multi-adapter inference + elif not hasattr(module, "set_adapter") and len(adapter_name) != 1: + raise ValueError( + "You are trying to set multiple adapters and you have a PEFT version that does not support multi-adapter inference. Please upgrade to the latest version of PEFT." + " `pip install -U peft` or `pip install -U git+https://github.com/huggingface/peft.git`" + ) + else: + module.active_adapter = adapter_name + _adapters_has_been_set = True + + if not _adapters_has_been_set: + raise ValueError( + "Did not succeeded in setting the adapter. Please make sure you are using a model that supports adapters." + ) + + def disable_adapters(self) -> None: + r""" + Disable all adapters attached to the model and fallback to inference with the base model only. + + If you are not familiar with adapters and PEFT methods, we invite you to read more about them on the PEFT + official documentation: https://huggingface.co/docs/peft + """ + check_peft_version(min_version=MIN_PEFT_VERSION) + + if not self._hf_peft_config_loaded: + raise ValueError("No adapter loaded. Please load an adapter first.") + + from peft.tuners.tuners_utils import BaseTunerLayer + + for _, module in self.named_modules(): + if isinstance(module, BaseTunerLayer): + if hasattr(module, "enable_adapters"): + module.enable_adapters(enabled=False) + else: + # support for older PEFT versions + module.disable_adapters = True + + def enable_adapters(self) -> None: + """ + Enable adapters that are attached to the model. The model will use `self.active_adapters()` to retrieve the + list of adapters to enable. + + If you are not familiar with adapters and PEFT methods, we invite you to read more about them on the PEFT + official documentation: https://huggingface.co/docs/peft + """ + check_peft_version(min_version=MIN_PEFT_VERSION) + + if not self._hf_peft_config_loaded: + raise ValueError("No adapter loaded. Please load an adapter first.") + + from peft.tuners.tuners_utils import BaseTunerLayer + + for _, module in self.named_modules(): + if isinstance(module, BaseTunerLayer): + if hasattr(module, "enable_adapters"): + module.enable_adapters(enabled=True) + else: + # support for older PEFT versions + module.disable_adapters = False + + def active_adapters(self) -> List[str]: + """ + Gets the current list of active adapters of the model. + + If you are not familiar with adapters and PEFT methods, we invite you to read more about them on the PEFT + official documentation: https://huggingface.co/docs/peft + """ + check_peft_version(min_version=MIN_PEFT_VERSION) + + if not self._hf_peft_config_loaded: + raise ValueError("No adapter loaded. Please load an adapter first.") + + from peft.tuners.tuners_utils import BaseTunerLayer + + for _, module in self.named_modules(): + if isinstance(module, BaseTunerLayer): + return module.active_adapter diff --git a/src/diffusers/models/modeling_utils.py b/src/diffusers/models/modeling_utils.py index 546c5b20f937..4fcc15333936 100644 --- a/src/diffusers/models/modeling_utils.py +++ b/src/diffusers/models/modeling_utils.py @@ -29,6 +29,7 @@ from torch import Tensor, nn from .. import __version__ +from ..integrations import PeftAdapterMixin from ..utils import ( CONFIG_NAME, FLAX_WEIGHTS_NAME, @@ -183,7 +184,7 @@ def load(module: torch.nn.Module, prefix: str = ""): return error_msgs -class ModelMixin(torch.nn.Module, PushToHubMixin): +class ModelMixin(torch.nn.Module, PushToHubMixin, PeftAdapterMixin): r""" Base class for all models. From 8bbdb665cbe8fb4e4f20e910534820191e34bea2 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Mon, 1 Jan 2024 14:58:51 +0530 Subject: [PATCH 2/8] remove duplicate methods. --- src/diffusers/models/modeling_utils.py | 150 ------------------------- 1 file changed, 150 deletions(-) diff --git a/src/diffusers/models/modeling_utils.py b/src/diffusers/models/modeling_utils.py index 4fcc15333936..ab4cab32cd96 100644 --- a/src/diffusers/models/modeling_utils.py +++ b/src/diffusers/models/modeling_utils.py @@ -33,12 +33,10 @@ from ..utils import ( CONFIG_NAME, FLAX_WEIGHTS_NAME, - MIN_PEFT_VERSION, SAFETENSORS_WEIGHTS_NAME, WEIGHTS_NAME, _add_variant, _get_model_file, - check_peft_version, deprecate, is_accelerate_available, is_torch_version, @@ -198,7 +196,6 @@ class ModelMixin(torch.nn.Module, PushToHubMixin, PeftAdapterMixin): _automatically_saved_args = ["_diffusers_version", "_class_name", "_name_or_path"] _supports_gradient_checkpointing = False _keys_to_ignore_on_load_unexpected = None - _hf_peft_config_loaded = False def __init__(self): super().__init__() @@ -304,153 +301,6 @@ def disable_xformers_memory_efficient_attention(self) -> None: """ self.set_use_memory_efficient_attention_xformers(False) - def add_adapter(self, adapter_config, adapter_name: str = "default") -> None: - r""" - Adds a new adapter to the current model for training. If no adapter name is passed, a default name is assigned - to the adapter to follow the convention of the PEFT library. - - If you are not familiar with adapters and PEFT methods, we invite you to read more about them in the PEFT - [documentation](https://huggingface.co/docs/peft). - - Args: - adapter_config (`[~peft.PeftConfig]`): - The configuration of the adapter to add; supported adapters are non-prefix tuning and adaption prompt - methods. - adapter_name (`str`, *optional*, defaults to `"default"`): - The name of the adapter to add. If no name is passed, a default name is assigned to the adapter. - """ - check_peft_version(min_version=MIN_PEFT_VERSION) - - from peft import PeftConfig, inject_adapter_in_model - - if not self._hf_peft_config_loaded: - self._hf_peft_config_loaded = True - elif adapter_name in self.peft_config: - raise ValueError(f"Adapter with name {adapter_name} already exists. Please use a different name.") - - if not isinstance(adapter_config, PeftConfig): - raise ValueError( - f"adapter_config should be an instance of PeftConfig. Got {type(adapter_config)} instead." - ) - - # Unlike transformers, here we don't need to retrieve the name_or_path of the unet as the loading logic is - # handled by the `load_lora_layers` or `LoraLoaderMixin`. Therefore we set it to `None` here. - adapter_config.base_model_name_or_path = None - inject_adapter_in_model(adapter_config, self, adapter_name) - self.set_adapter(adapter_name) - - def set_adapter(self, adapter_name: Union[str, List[str]]) -> None: - """ - Sets a specific adapter by forcing the model to only use that adapter and disables the other adapters. - - If you are not familiar with adapters and PEFT methods, we invite you to read more about them on the PEFT - official documentation: https://huggingface.co/docs/peft - - Args: - adapter_name (Union[str, List[str]])): - The list of adapters to set or the adapter name in case of single adapter. - """ - check_peft_version(min_version=MIN_PEFT_VERSION) - - if not self._hf_peft_config_loaded: - raise ValueError("No adapter loaded. Please load an adapter first.") - - if isinstance(adapter_name, str): - adapter_name = [adapter_name] - - missing = set(adapter_name) - set(self.peft_config) - if len(missing) > 0: - raise ValueError( - f"Following adapter(s) could not be found: {', '.join(missing)}. Make sure you are passing the correct adapter name(s)." - f" current loaded adapters are: {list(self.peft_config.keys())}" - ) - - from peft.tuners.tuners_utils import BaseTunerLayer - - _adapters_has_been_set = False - - for _, module in self.named_modules(): - if isinstance(module, BaseTunerLayer): - if hasattr(module, "set_adapter"): - module.set_adapter(adapter_name) - # Previous versions of PEFT does not support multi-adapter inference - elif not hasattr(module, "set_adapter") and len(adapter_name) != 1: - raise ValueError( - "You are trying to set multiple adapters and you have a PEFT version that does not support multi-adapter inference. Please upgrade to the latest version of PEFT." - " `pip install -U peft` or `pip install -U git+https://github.com/huggingface/peft.git`" - ) - else: - module.active_adapter = adapter_name - _adapters_has_been_set = True - - if not _adapters_has_been_set: - raise ValueError( - "Did not succeeded in setting the adapter. Please make sure you are using a model that supports adapters." - ) - - def disable_adapters(self) -> None: - r""" - Disable all adapters attached to the model and fallback to inference with the base model only. - - If you are not familiar with adapters and PEFT methods, we invite you to read more about them on the PEFT - official documentation: https://huggingface.co/docs/peft - """ - check_peft_version(min_version=MIN_PEFT_VERSION) - - if not self._hf_peft_config_loaded: - raise ValueError("No adapter loaded. Please load an adapter first.") - - from peft.tuners.tuners_utils import BaseTunerLayer - - for _, module in self.named_modules(): - if isinstance(module, BaseTunerLayer): - if hasattr(module, "enable_adapters"): - module.enable_adapters(enabled=False) - else: - # support for older PEFT versions - module.disable_adapters = True - - def enable_adapters(self) -> None: - """ - Enable adapters that are attached to the model. The model will use `self.active_adapters()` to retrieve the - list of adapters to enable. - - If you are not familiar with adapters and PEFT methods, we invite you to read more about them on the PEFT - official documentation: https://huggingface.co/docs/peft - """ - check_peft_version(min_version=MIN_PEFT_VERSION) - - if not self._hf_peft_config_loaded: - raise ValueError("No adapter loaded. Please load an adapter first.") - - from peft.tuners.tuners_utils import BaseTunerLayer - - for _, module in self.named_modules(): - if isinstance(module, BaseTunerLayer): - if hasattr(module, "enable_adapters"): - module.enable_adapters(enabled=True) - else: - # support for older PEFT versions - module.disable_adapters = False - - def active_adapters(self) -> List[str]: - """ - Gets the current list of active adapters of the model. - - If you are not familiar with adapters and PEFT methods, we invite you to read more about them on the PEFT - official documentation: https://huggingface.co/docs/peft - """ - check_peft_version(min_version=MIN_PEFT_VERSION) - - if not self._hf_peft_config_loaded: - raise ValueError("No adapter loaded. Please load an adapter first.") - - from peft.tuners.tuners_utils import BaseTunerLayer - - for _, module in self.named_modules(): - if isinstance(module, BaseTunerLayer): - return module.active_adapter - def save_pretrained( self, save_directory: Union[str, os.PathLike], From 9e461c1561df72ea0b13d34ac3dabd4cf0b7a99f Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Mon, 1 Jan 2024 15:07:08 +0530 Subject: [PATCH 3/8] better imports. --- src/diffusers/integrations/__init__.py | 8 +++----- src/diffusers/integrations/peft.py | 8 +++++++- 2 files changed, 10 insertions(+), 6 deletions(-) diff --git a/src/diffusers/integrations/__init__.py b/src/diffusers/integrations/__init__.py index ec2408f348c5..98d9306e1141 100644 --- a/src/diffusers/integrations/__init__.py +++ b/src/diffusers/integrations/__init__.py @@ -18,12 +18,10 @@ _import_structure = {} -if is_torch_available() and is_peft_available(): - _import_structure = {"peft": ["PeftAdapterMixin"]} +_import_structure = {"peft": ["PeftAdapterMixin"]} -if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: - if is_torch_available() and is_peft_available(): - from .peft import PeftAdapterMixin +if TYPE_CHECKING: + from .peft import PeftAdapterMixin else: import sys diff --git a/src/diffusers/integrations/peft.py b/src/diffusers/integrations/peft.py index 098b4ee50073..a7f1e1c938f0 100644 --- a/src/diffusers/integrations/peft.py +++ b/src/diffusers/integrations/peft.py @@ -14,7 +14,7 @@ # limitations under the License. from typing import List, Union -from ..utils import MIN_PEFT_VERSION, check_peft_version +from ..utils import MIN_PEFT_VERSION, check_peft_version, is_peft_available class PeftAdapterMixin: @@ -51,6 +51,9 @@ def add_adapter(self, adapter_config, adapter_name: str = "default") -> None: """ check_peft_version(min_version=MIN_PEFT_VERSION) + if not is_peft_available(): + raise ImportError("PEFT is not available. Please install PEFT to use this function: `pip install peft`.") + from peft import PeftConfig, inject_adapter_in_model if not self._hf_peft_config_loaded: @@ -172,6 +175,9 @@ def active_adapters(self) -> List[str]: """ check_peft_version(min_version=MIN_PEFT_VERSION) + if not is_peft_available(): + raise ImportError("PEFT is not available. Please install PEFT to use this function: `pip install peft`.") + if not self._hf_peft_config_loaded: raise ValueError("No adapter loaded. Please load an adapter first.") From 763f28258b5d1b52130faaaea2758226d193431b Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Wed, 3 Jan 2024 11:03:19 +0530 Subject: [PATCH 4/8] move to loaders.py --- src/diffusers/integrations/__init__.py | 28 ------------------- src/diffusers/loaders/__init__.py | 6 +++- .../{integrations => loaders}/peft.py | 0 src/diffusers/models/modeling_utils.py | 2 +- 4 files changed, 6 insertions(+), 30 deletions(-) delete mode 100644 src/diffusers/integrations/__init__.py rename src/diffusers/{integrations => loaders}/peft.py (100%) diff --git a/src/diffusers/integrations/__init__.py b/src/diffusers/integrations/__init__.py deleted file mode 100644 index 98d9306e1141..000000000000 --- a/src/diffusers/integrations/__init__.py +++ /dev/null @@ -1,28 +0,0 @@ -# Copyright 2023 The HuggingFace Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -from typing import TYPE_CHECKING - -from ..utils import DIFFUSERS_SLOW_IMPORT, _LazyModule, is_peft_available, is_torch_available - - -_import_structure = {} - -_import_structure = {"peft": ["PeftAdapterMixin"]} - -if TYPE_CHECKING: - from .peft import PeftAdapterMixin -else: - import sys - - sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__) diff --git a/src/diffusers/loaders/__init__.py b/src/diffusers/loaders/__init__.py index 45c8c97c76eb..d7855206a287 100644 --- a/src/diffusers/loaders/__init__.py +++ b/src/diffusers/loaders/__init__.py @@ -1,7 +1,7 @@ from typing import TYPE_CHECKING from ..utils import DIFFUSERS_SLOW_IMPORT, _LazyModule, deprecate -from ..utils.import_utils import is_torch_available, is_transformers_available +from ..utils.import_utils import is_peft_available, is_torch_available, is_transformers_available def text_encoder_lora_state_dict(text_encoder): @@ -64,6 +64,8 @@ def text_encoder_attn_modules(text_encoder): _import_structure["textual_inversion"] = ["TextualInversionLoaderMixin"] _import_structure["ip_adapter"] = ["IPAdapterMixin"] +_import_structure["peft"] = ["PeftAdapterMixin"] + if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: if is_torch_available(): @@ -76,6 +78,8 @@ def text_encoder_attn_modules(text_encoder): from .lora import LoraLoaderMixin, StableDiffusionXLLoraLoaderMixin from .single_file import FromSingleFileMixin from .textual_inversion import TextualInversionLoaderMixin + + from .peft import PeftAdapterMixin else: import sys diff --git a/src/diffusers/integrations/peft.py b/src/diffusers/loaders/peft.py similarity index 100% rename from src/diffusers/integrations/peft.py rename to src/diffusers/loaders/peft.py diff --git a/src/diffusers/models/modeling_utils.py b/src/diffusers/models/modeling_utils.py index ab4cab32cd96..855f5fcab051 100644 --- a/src/diffusers/models/modeling_utils.py +++ b/src/diffusers/models/modeling_utils.py @@ -29,7 +29,7 @@ from torch import Tensor, nn from .. import __version__ -from ..integrations import PeftAdapterMixin +from ..loaders import PeftAdapterMixin from ..utils import ( CONFIG_NAME, FLAX_WEIGHTS_NAME, From 1240a58fb4f2713c248f542ded9a9e1e87cf4165 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Wed, 3 Jan 2024 18:17:00 +0530 Subject: [PATCH 5/8] remove peftadaptermixin from modelmixin. --- src/diffusers/models/modeling_utils.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/diffusers/models/modeling_utils.py b/src/diffusers/models/modeling_utils.py index 855f5fcab051..445c3ca71caf 100644 --- a/src/diffusers/models/modeling_utils.py +++ b/src/diffusers/models/modeling_utils.py @@ -29,7 +29,6 @@ from torch import Tensor, nn from .. import __version__ -from ..loaders import PeftAdapterMixin from ..utils import ( CONFIG_NAME, FLAX_WEIGHTS_NAME, @@ -182,7 +181,7 @@ def load(module: torch.nn.Module, prefix: str = ""): return error_msgs -class ModelMixin(torch.nn.Module, PushToHubMixin, PeftAdapterMixin): +class ModelMixin(torch.nn.Module, PushToHubMixin): r""" Base class for all models. From 488463b2768a9f5210a0bb19e2b79ae62088f98d Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Wed, 3 Jan 2024 18:24:01 +0530 Subject: [PATCH 6/8] add: peftadaptermixin selectively. --- src/diffusers/models/prior_transformer.py | 4 ++-- src/diffusers/models/unet_2d_condition.py | 4 ++-- src/diffusers/models/uvit_2d.py | 3 ++- .../pipelines/wuerstchen/modeling_wuerstchen_prior.py | 4 ++-- 4 files changed, 8 insertions(+), 7 deletions(-) diff --git a/src/diffusers/models/prior_transformer.py b/src/diffusers/models/prior_transformer.py index 6c5e406ad378..e2c02a99332a 100644 --- a/src/diffusers/models/prior_transformer.py +++ b/src/diffusers/models/prior_transformer.py @@ -6,7 +6,7 @@ from torch import nn from ..configuration_utils import ConfigMixin, register_to_config -from ..loaders import UNet2DConditionLoadersMixin +from ..loaders import PeftAdapterMixin, UNet2DConditionLoadersMixin from ..utils import BaseOutput from .attention import BasicTransformerBlock from .attention_processor import ( @@ -33,7 +33,7 @@ class PriorTransformerOutput(BaseOutput): predicted_image_embedding: torch.FloatTensor -class PriorTransformer(ModelMixin, ConfigMixin, UNet2DConditionLoadersMixin): +class PriorTransformer(ModelMixin, ConfigMixin, UNet2DConditionLoadersMixin, PeftAdapterMixin): """ A Prior Transformer model. diff --git a/src/diffusers/models/unet_2d_condition.py b/src/diffusers/models/unet_2d_condition.py index 623e4d88d564..7ae043ebfe24 100644 --- a/src/diffusers/models/unet_2d_condition.py +++ b/src/diffusers/models/unet_2d_condition.py @@ -19,7 +19,7 @@ import torch.utils.checkpoint from ..configuration_utils import ConfigMixin, register_to_config -from ..loaders import UNet2DConditionLoadersMixin +from ..loaders import PeftAdapterMixin, UNet2DConditionLoadersMixin from ..utils import USE_PEFT_BACKEND, BaseOutput, deprecate, logging, scale_lora_layers, unscale_lora_layers from .activations import get_activation from .attention_processor import ( @@ -68,7 +68,7 @@ class UNet2DConditionOutput(BaseOutput): sample: torch.FloatTensor = None -class UNet2DConditionModel(ModelMixin, ConfigMixin, UNet2DConditionLoadersMixin): +class UNet2DConditionModel(ModelMixin, ConfigMixin, UNet2DConditionLoadersMixin, PeftAdapterMixin): r""" A conditional 2D UNet model that takes a noisy sample, conditional state, and a timestep and returns a sample shaped output. diff --git a/src/diffusers/models/uvit_2d.py b/src/diffusers/models/uvit_2d.py index 14dd8aee8e89..014fbb70b84f 100644 --- a/src/diffusers/models/uvit_2d.py +++ b/src/diffusers/models/uvit_2d.py @@ -21,6 +21,7 @@ from torch.utils.checkpoint import checkpoint from ..configuration_utils import ConfigMixin, register_to_config +from ..loaders import PeftAdapterMixin from .attention import BasicTransformerBlock, SkipFFTransformerBlock from .attention_processor import ( ADDED_KV_ATTENTION_PROCESSORS, @@ -35,7 +36,7 @@ from .resnet import Downsample2D, Upsample2D -class UVit2DModel(ModelMixin, ConfigMixin): +class UVit2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin): _supports_gradient_checkpointing = True @register_to_config diff --git a/src/diffusers/pipelines/wuerstchen/modeling_wuerstchen_prior.py b/src/diffusers/pipelines/wuerstchen/modeling_wuerstchen_prior.py index a7d9e32fb6c9..8276ece87034 100644 --- a/src/diffusers/pipelines/wuerstchen/modeling_wuerstchen_prior.py +++ b/src/diffusers/pipelines/wuerstchen/modeling_wuerstchen_prior.py @@ -20,7 +20,7 @@ import torch.nn as nn from ...configuration_utils import ConfigMixin, register_to_config -from ...loaders import UNet2DConditionLoadersMixin +from ...loaders import PeftAdapterMixin, UNet2DConditionLoadersMixin from ...models.attention_processor import ( ADDED_KV_ATTENTION_PROCESSORS, CROSS_ATTENTION_PROCESSORS, @@ -34,7 +34,7 @@ from .modeling_wuerstchen_common import AttnBlock, ResBlock, TimestepBlock, WuerstchenLayerNorm -class WuerstchenPrior(ModelMixin, ConfigMixin, UNet2DConditionLoadersMixin): +class WuerstchenPrior(ModelMixin, ConfigMixin, UNet2DConditionLoadersMixin, PeftAdapterMixin): unet_name = "prior" _supports_gradient_checkpointing = True From b6e7465b0c3d8733ead8244183d1372d9271935d Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Wed, 3 Jan 2024 18:31:44 +0530 Subject: [PATCH 7/8] add: entry to _toctree --- docs/source/en/_toctree.yml | 2 ++ docs/source/en/api/loaders/peft.md | 25 +++++++++++++++++++++++++ 2 files changed, 27 insertions(+) create mode 100644 docs/source/en/api/loaders/peft.md diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index 0c05f0ef7ffa..e70fcf118b9b 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -210,6 +210,8 @@ title: Textual Inversion - local: api/loaders/unet title: UNet + - local: api/loaders/peft + title: PEFT title: Loaders - sections: - local: api/models/overview diff --git a/docs/source/en/api/loaders/peft.md b/docs/source/en/api/loaders/peft.md new file mode 100644 index 000000000000..1aff389d3142 --- /dev/null +++ b/docs/source/en/api/loaders/peft.md @@ -0,0 +1,25 @@ + + +# PEFT + +Diffusers supports working with adapters (such as [LoRA](../../using-diffusers/loading_adapters)) via the [`peft` library](https://huggingface.co/docs/peft/index). We provide a `PeftAdapterMixin` class to handle this for modeling classes in Diffusers (such as [`UNet2DConditionModel`]). + + + +Refer to [this doc](../../tutorials/using_peft_for_inference.md) to get an overview of how to work with `peft` in Diffusers for inference. + + + +## PeftAdapterMixin + +[[autodoc]] loaders.peft.PeftAdapterMixin \ No newline at end of file From 0107460d697880b6b0ee8d3e0b689ab522258688 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Wed, 3 Jan 2024 20:23:31 +0530 Subject: [PATCH 8/8] Empty-Commit