[core / PEFT / LoRA] Integrate PEFT into Unet

src/diffusers/models/attention.py

@@ -17,6 +17,7 @@
 import torch.nn.functional as F
 from torch import nn
 
+from ..utils import USE_PEFT_BACKEND
 from ..utils.torch_utils import maybe_allow_in_graph
 from .activations import get_activation
 from .attention_processor import Attention
@@ -300,6 +301,7 @@ def __init__(
  super().__init__()
  inner_dim = int(dim * mult)
  dim_out = dim_out if dim_out is not None else dim
+ linear_cls = LoRACompatibleLinear if not USE_PEFT_BACKEND else nn.Linear
 
  if activation_fn == "gelu":
  act_fn = GELU(dim, inner_dim)
@@ -316,14 +318,15 @@ def __init__(
  # project dropout
  self.net.append(nn.Dropout(dropout))
  # project out
- self.net.append(LoRACompatibleLinear(inner_dim, dim_out))
+ self.net.append(linear_cls(inner_dim, dim_out))
  # FF as used in Vision Transformer, MLP-Mixer, etc. have a final dropout
  if final_dropout:
  self.net.append(nn.Dropout(dropout))
 
  def forward(self, hidden_states: torch.Tensor, scale: float = 1.0) -> torch.Tensor:
+ compatible_cls = (GEGLU,) if USE_PEFT_BACKEND else (GEGLU, LoRACompatibleLinear)
  for module in self.net:
- if isinstance(module, (LoRACompatibleLinear, GEGLU)):
+ if isinstance(module, compatible_cls):
  hidden_states = module(hidden_states, scale)
  else:
  hidden_states = module(hidden_states)
@@ -368,7 +371,9 @@ class GEGLU(nn.Module):
 
  def __init__(self, dim_in: int, dim_out: int):
  super().__init__()
- self.proj = LoRACompatibleLinear(dim_in, dim_out * 2)
+ linear_cls = LoRACompatibleLinear if not USE_PEFT_BACKEND else nn.Linear
+
+ self.proj = linear_cls(dim_in, dim_out * 2)
 
  def gelu(self, gate: torch.Tensor) -> torch.Tensor:
  if gate.device.type != "mps":
@@ -377,7 +382,8 @@ def gelu(self, gate: torch.Tensor) -> torch.Tensor:
  return F.gelu(gate.to(dtype=torch.float32)).to(dtype=gate.dtype)
 
  def forward(self, hidden_states, scale: float = 1.0):
- hidden_states, gate = self.proj(hidden_states, scale).chunk(2, dim=-1)
+ args = () if USE_PEFT_BACKEND else (scale,)
+ hidden_states, gate = self.proj(hidden_states, *args).chunk(2, dim=-1)
  return hidden_states * self.gelu(gate)
 
 

src/diffusers/models/attention_processor.py

@@ -18,7 +18,7 @@
 import torch.nn.functional as F
 from torch import nn
 
-from ..utils import deprecate, logging
+from ..utils import USE_PEFT_BACKEND, deprecate, logging
 from ..utils.import_utils import is_xformers_available
 from ..utils.torch_utils import maybe_allow_in_graph
 from .lora import LoRACompatibleLinear, LoRALinearLayer
@@ -137,22 +137,27 @@ def __init__(
  f"unknown cross_attention_norm: {cross_attention_norm}. Should be None, 'layer_norm' or 'group_norm'"
  )
 
- self.to_q = LoRACompatibleLinear(query_dim, self.inner_dim, bias=bias)
+ if USE_PEFT_BACKEND:
+ linear_cls = nn.Linear
+ else:
+ linear_cls = LoRACompatibleLinear
+
+ self.to_q = linear_cls(query_dim, self.inner_dim, bias=bias)
 
  if not self.only_cross_attention:
  # only relevant for the `AddedKVProcessor` classes
- self.to_k = LoRACompatibleLinear(self.cross_attention_dim, self.inner_dim, bias=bias)
- self.to_v = LoRACompatibleLinear(self.cross_attention_dim, self.inner_dim, bias=bias)
+ self.to_k = linear_cls(self.cross_attention_dim, self.inner_dim, bias=bias)
+ self.to_v = linear_cls(self.cross_attention_dim, self.inner_dim, bias=bias)
  else:
  self.to_k = None
  self.to_v = None
 
  if self.added_kv_proj_dim is not None:
- self.add_k_proj = LoRACompatibleLinear(added_kv_proj_dim, self.inner_dim)
- self.add_v_proj = LoRACompatibleLinear(added_kv_proj_dim, self.inner_dim)
+ self.add_k_proj = linear_cls(added_kv_proj_dim, self.inner_dim)
+ self.add_v_proj = linear_cls(added_kv_proj_dim, self.inner_dim)
 
  self.to_out = nn.ModuleList([])
- self.to_out.append(LoRACompatibleLinear(self.inner_dim, query_dim, bias=out_bias))
+ self.to_out.append(linear_cls(self.inner_dim, query_dim, bias=out_bias))
  self.to_out.append(nn.Dropout(dropout))
 
  # set attention processor
@@ -545,6 +550,8 @@ def __call__(
  ):
  residual = hidden_states
 
+ args = () if USE_PEFT_BACKEND else (scale,)
+
  if attn.spatial_norm is not None:
  hidden_states = attn.spatial_norm(hidden_states, temb)
 
@@ -562,15 +569,15 @@ def __call__(
  if attn.group_norm is not None:
  hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
 
- query = attn.to_q(hidden_states, scale=scale)
+ query = attn.to_q(hidden_states, *args)
 
  if encoder_hidden_states is None:
  encoder_hidden_states = hidden_states
  elif attn.norm_cross:
  encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
 
- key = attn.to_k(encoder_hidden_states, scale=scale)
- value = attn.to_v(encoder_hidden_states, scale=scale)
+ key = attn.to_k(encoder_hidden_states, *args)
+ value = attn.to_v(encoder_hidden_states, *args)
 
  query = attn.head_to_batch_dim(query)
  key = attn.head_to_batch_dim(key)
@@ -581,7 +588,7 @@ def __call__(
  hidden_states = attn.batch_to_head_dim(hidden_states)
 
  # linear proj
- hidden_states = attn.to_out[0](hidden_states, scale=scale)
+ hidden_states = attn.to_out[0](hidden_states, *args)
  # dropout
  hidden_states = attn.to_out[1](hidden_states)
 
@@ -1007,15 +1014,20 @@ def __call__(
  if attn.group_norm is not None:
  hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
 
- query = attn.to_q(hidden_states, scale=scale)
+ args = () if USE_PEFT_BACKEND else (scale,)
+ query = attn.to_q(hidden_states, *args)
 
  if encoder_hidden_states is None:
  encoder_hidden_states = hidden_states
  elif attn.norm_cross:
  encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
 
- key = attn.to_k(encoder_hidden_states, scale=scale)
- value = attn.to_v(encoder_hidden_states, scale=scale)
+ key = (
+ attn.to_k(encoder_hidden_states, scale=scale) if not USE_PEFT_BACKEND else attn.to_k(encoder_hidden_states)
+ )
+ value = (
+ attn.to_v(encoder_hidden_states, scale=scale) if not USE_PEFT_BACKEND else attn.to_v(encoder_hidden_states)
+ )
 
  inner_dim = key.shape[-1]
  head_dim = inner_dim // attn.heads
@@ -1035,7 +1047,9 @@ def __call__(
  hidden_states = hidden_states.to(query.dtype)
 
  # linear proj
- hidden_states = attn.to_out[0](hidden_states, scale=scale)
+ hidden_states = (
+ attn.to_out[0](hidden_states, scale=scale) if not USE_PEFT_BACKEND else attn.to_out[0](hidden_states)
+ )
  # dropout
  hidden_states = attn.to_out[1](hidden_states)
 

src/diffusers/models/embeddings.py

@@ -18,6 +18,7 @@
 import torch
 from torch import nn
 
+from ..utils import USE_PEFT_BACKEND
 from .activations import get_activation
 from .lora import LoRACompatibleLinear
 
@@ -166,8 +167,9 @@ def __init__(
  cond_proj_dim=None,
  ):
  super().__init__()
+ linear_cls = nn.Linear if USE_PEFT_BACKEND else LoRACompatibleLinear
 
- self.linear_1 = LoRACompatibleLinear(in_channels, time_embed_dim)
+ self.linear_1 = linear_cls(in_channels, time_embed_dim)
 
  if cond_proj_dim is not None:
  self.cond_proj = nn.Linear(cond_proj_dim, in_channels, bias=False)
@@ -180,7 +182,7 @@ def __init__(
  time_embed_dim_out = out_dim
  else:
  time_embed_dim_out = time_embed_dim
- self.linear_2 = LoRACompatibleLinear(time_embed_dim, time_embed_dim_out)
+ self.linear_2 = linear_cls(time_embed_dim, time_embed_dim_out)
 
  if post_act_fn is None:
  self.post_act = None

src/diffusers/models/modeling_utils.py

@@ -32,10 +32,12 @@
  DIFFUSERS_CACHE,
  FLAX_WEIGHTS_NAME,
  HF_HUB_OFFLINE,
+ MIN_PEFT_VERSION,
  SAFETENSORS_WEIGHTS_NAME,
  WEIGHTS_NAME,
  _add_variant,
  _get_model_file,
+ check_peft_version,
  deprecate,
  is_accelerate_available,
  is_torch_version,
@@ -187,6 +189,7 @@ class ModelMixin(torch.nn.Module, PushToHubMixin):
  _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__()
@@ -292,6 +295,153 @@ def disable_xformers_memory_efficient_attention(self):
  """
  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],