From 0f6a9a2435cc0f38a5e3a12c75602dd57788979c Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Tue, 2 Jan 2024 18:53:35 +0530 Subject: [PATCH 01/11] handle rest of the stuff related to deprecated lora stuff. --- .../controlnetxs/controlnetxs.py | 6 ++---- src/diffusers/models/attention_processor.py | 17 +---------------- .../models/autoencoders/autoencoder_kl.py | 4 ++-- .../autoencoder_kl_temporal_decoder.py | 4 ++-- .../autoencoders/consistency_decoder_vae.py | 4 ++-- src/diffusers/models/controlnet.py | 10 ++++------ src/diffusers/models/prior_transformer.py | 10 ++++------ src/diffusers/models/unet_2d_condition.py | 4 ++-- src/diffusers/models/unet_3d_condition.py | 4 ++-- src/diffusers/models/unet_motion_model.py | 4 ++-- src/diffusers/models/uvit_2d.py | 4 ++-- .../pipelines/audioldm2/modeling_audioldm2.py | 4 ++-- .../versatile_diffusion/modeling_text_unet.py | 4 ++-- .../wuerstchen/modeling_wuerstchen_prior.py | 4 ++-- 14 files changed, 31 insertions(+), 52 deletions(-) diff --git a/examples/research_projects/controlnetxs/controlnetxs.py b/examples/research_projects/controlnetxs/controlnetxs.py index c6419b44daeb..20c8d0fdf0f1 100644 --- a/examples/research_projects/controlnetxs/controlnetxs.py +++ b/examples/research_projects/controlnetxs/controlnetxs.py @@ -494,9 +494,7 @@ def attn_processors(self) -> Dict[str, AttentionProcessor]: """ return self.control_model.attn_processors - def set_attn_processor( - self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]], _remove_lora=False - ): + def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): r""" Sets the attention processor to use to compute attention. @@ -509,7 +507,7 @@ def set_attn_processor( processor. This is strongly recommended when setting trainable attention processors. """ - self.control_model.set_attn_processor(processor, _remove_lora) + self.control_model.set_attn_processor(processor) def set_default_attn_processor(self): """ diff --git a/src/diffusers/models/attention_processor.py b/src/diffusers/models/attention_processor.py index 23a3e2bb3791..ac9563e186bb 100644 --- a/src/diffusers/models/attention_processor.py +++ b/src/diffusers/models/attention_processor.py @@ -373,29 +373,14 @@ def set_attention_slice(self, slice_size: int) -> None: self.set_processor(processor) - def set_processor(self, processor: "AttnProcessor", _remove_lora: bool = False) -> None: + def set_processor(self, processor: "AttnProcessor") -> None: r""" Set the attention processor to use. Args: processor (`AttnProcessor`): The attention processor to use. - _remove_lora (`bool`, *optional*, defaults to `False`): - Set to `True` to remove LoRA layers from the model. """ - if not USE_PEFT_BACKEND and hasattr(self, "processor") and _remove_lora and self.to_q.lora_layer is not None: - deprecate( - "set_processor to offload LoRA", - "0.26.0", - "In detail, removing LoRA layers via calling `set_default_attn_processor` is deprecated. Please make sure to call `pipe.unload_lora_weights()` instead.", - ) - # TODO(Patrick, Sayak) - this can be deprecated once PEFT LoRA integration is complete - # We need to remove all LoRA layers - # Don't forget to remove ALL `_remove_lora` from the codebase - for module in self.modules(): - if hasattr(module, "set_lora_layer"): - module.set_lora_layer(None) - # if current processor is in `self._modules` and if passed `processor` is not, we need to # pop `processor` from `self._modules` if ( diff --git a/src/diffusers/models/autoencoders/autoencoder_kl.py b/src/diffusers/models/autoencoders/autoencoder_kl.py index ae2d90c548f8..e601be0331eb 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl.py @@ -208,9 +208,9 @@ def set_attn_processor( def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): if hasattr(module, "set_processor"): if not isinstance(processor, dict): - module.set_processor(processor, _remove_lora=_remove_lora) + module.set_processor(processor) else: - module.set_processor(processor.pop(f"{name}.processor"), _remove_lora=_remove_lora) + module.set_processor(processor.pop(f"{name}.processor")) for sub_name, child in module.named_children(): fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_temporal_decoder.py b/src/diffusers/models/autoencoders/autoencoder_kl_temporal_decoder.py index 0b7f8d1f5336..ab44ae5e0bcf 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_temporal_decoder.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_temporal_decoder.py @@ -293,9 +293,9 @@ def set_attn_processor( def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): if hasattr(module, "set_processor"): if not isinstance(processor, dict): - module.set_processor(processor, _remove_lora=_remove_lora) + module.set_processor(processor) else: - module.set_processor(processor.pop(f"{name}.processor"), _remove_lora=_remove_lora) + module.set_processor(processor.pop(f"{name}.processor")) for sub_name, child in module.named_children(): fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) diff --git a/src/diffusers/models/autoencoders/consistency_decoder_vae.py b/src/diffusers/models/autoencoders/consistency_decoder_vae.py index d92423eafc31..bb960504c324 100644 --- a/src/diffusers/models/autoencoders/consistency_decoder_vae.py +++ b/src/diffusers/models/autoencoders/consistency_decoder_vae.py @@ -238,9 +238,9 @@ def set_attn_processor( def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): if hasattr(module, "set_processor"): if not isinstance(processor, dict): - module.set_processor(processor, _remove_lora=_remove_lora) + module.set_processor(processor) else: - module.set_processor(processor.pop(f"{name}.processor"), _remove_lora=_remove_lora) + module.set_processor(processor.pop(f"{name}.processor")) for sub_name, child in module.named_children(): fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) diff --git a/src/diffusers/models/controlnet.py b/src/diffusers/models/controlnet.py index 3139bb2a5c6c..1102f4f9d36d 100644 --- a/src/diffusers/models/controlnet.py +++ b/src/diffusers/models/controlnet.py @@ -534,9 +534,7 @@ def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: return processors # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_attn_processor - def set_attn_processor( - self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]], _remove_lora=False - ): + def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): r""" Sets the attention processor to use to compute attention. @@ -560,9 +558,9 @@ def set_attn_processor( def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): if hasattr(module, "set_processor"): if not isinstance(processor, dict): - module.set_processor(processor, _remove_lora=_remove_lora) + module.set_processor(processor) else: - module.set_processor(processor.pop(f"{name}.processor"), _remove_lora=_remove_lora) + module.set_processor(processor.pop(f"{name}.processor")) for sub_name, child in module.named_children(): fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) @@ -584,7 +582,7 @@ def set_default_attn_processor(self): f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}" ) - self.set_attn_processor(processor, _remove_lora=True) + self.set_attn_processor(processor) # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_attention_slice def set_attention_slice(self, slice_size: Union[str, int, List[int]]) -> None: diff --git a/src/diffusers/models/prior_transformer.py b/src/diffusers/models/prior_transformer.py index 6c5e406ad378..8ada0a7c08a5 100644 --- a/src/diffusers/models/prior_transformer.py +++ b/src/diffusers/models/prior_transformer.py @@ -192,9 +192,7 @@ def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: return processors # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_attn_processor - def set_attn_processor( - self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]], _remove_lora=False - ): + def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): r""" Sets the attention processor to use to compute attention. @@ -218,9 +216,9 @@ def set_attn_processor( def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): if hasattr(module, "set_processor"): if not isinstance(processor, dict): - module.set_processor(processor, _remove_lora=_remove_lora) + module.set_processor(processor) else: - module.set_processor(processor.pop(f"{name}.processor"), _remove_lora=_remove_lora) + module.set_processor(processor.pop(f"{name}.processor")) for sub_name, child in module.named_children(): fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) @@ -242,7 +240,7 @@ def set_default_attn_processor(self): f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}" ) - self.set_attn_processor(processor, _remove_lora=True) + self.set_attn_processor(processor) def forward( self, diff --git a/src/diffusers/models/unet_2d_condition.py b/src/diffusers/models/unet_2d_condition.py index 623e4d88d564..397b11b74073 100644 --- a/src/diffusers/models/unet_2d_condition.py +++ b/src/diffusers/models/unet_2d_condition.py @@ -669,9 +669,9 @@ def set_attn_processor( def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): if hasattr(module, "set_processor"): if not isinstance(processor, dict): - module.set_processor(processor, _remove_lora=_remove_lora) + module.set_processor(processor) else: - module.set_processor(processor.pop(f"{name}.processor"), _remove_lora=_remove_lora) + module.set_processor(processor.pop(f"{name}.processor")) for sub_name, child in module.named_children(): fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) diff --git a/src/diffusers/models/unet_3d_condition.py b/src/diffusers/models/unet_3d_condition.py index 3c76b5aa8452..1b49a746858a 100644 --- a/src/diffusers/models/unet_3d_condition.py +++ b/src/diffusers/models/unet_3d_condition.py @@ -401,9 +401,9 @@ def set_attn_processor( def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): if hasattr(module, "set_processor"): if not isinstance(processor, dict): - module.set_processor(processor, _remove_lora=_remove_lora) + module.set_processor(processor) else: - module.set_processor(processor.pop(f"{name}.processor"), _remove_lora=_remove_lora) + module.set_processor(processor.pop(f"{name}.processor")) for sub_name, child in module.named_children(): fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) diff --git a/src/diffusers/models/unet_motion_model.py b/src/diffusers/models/unet_motion_model.py index 0bbc573e7df1..4375515ae043 100644 --- a/src/diffusers/models/unet_motion_model.py +++ b/src/diffusers/models/unet_motion_model.py @@ -575,9 +575,9 @@ def set_attn_processor( def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): if hasattr(module, "set_processor"): if not isinstance(processor, dict): - module.set_processor(processor, _remove_lora=_remove_lora) + module.set_processor(processor) else: - module.set_processor(processor.pop(f"{name}.processor"), _remove_lora=_remove_lora) + module.set_processor(processor.pop(f"{name}.processor")) for sub_name, child in module.named_children(): fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) diff --git a/src/diffusers/models/uvit_2d.py b/src/diffusers/models/uvit_2d.py index 14dd8aee8e89..9df01aebf229 100644 --- a/src/diffusers/models/uvit_2d.py +++ b/src/diffusers/models/uvit_2d.py @@ -263,9 +263,9 @@ def set_attn_processor( def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): if hasattr(module, "set_processor"): if not isinstance(processor, dict): - module.set_processor(processor, _remove_lora=_remove_lora) + module.set_processor(processor) else: - module.set_processor(processor.pop(f"{name}.processor"), _remove_lora=_remove_lora) + module.set_processor(processor.pop(f"{name}.processor")) for sub_name, child in module.named_children(): fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) diff --git a/src/diffusers/pipelines/audioldm2/modeling_audioldm2.py b/src/diffusers/pipelines/audioldm2/modeling_audioldm2.py index e855c2f0d6f1..4c688640fbba 100644 --- a/src/diffusers/pipelines/audioldm2/modeling_audioldm2.py +++ b/src/diffusers/pipelines/audioldm2/modeling_audioldm2.py @@ -564,9 +564,9 @@ def set_attn_processor( def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): if hasattr(module, "set_processor"): if not isinstance(processor, dict): - module.set_processor(processor, _remove_lora=_remove_lora) + module.set_processor(processor) else: - module.set_processor(processor.pop(f"{name}.processor"), _remove_lora=_remove_lora) + module.set_processor(processor.pop(f"{name}.processor")) for sub_name, child in module.named_children(): fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) diff --git a/src/diffusers/pipelines/deprecated/versatile_diffusion/modeling_text_unet.py b/src/diffusers/pipelines/deprecated/versatile_diffusion/modeling_text_unet.py index 7c9936a0bd4e..fc92db7236eb 100644 --- a/src/diffusers/pipelines/deprecated/versatile_diffusion/modeling_text_unet.py +++ b/src/diffusers/pipelines/deprecated/versatile_diffusion/modeling_text_unet.py @@ -874,9 +874,9 @@ def set_attn_processor( def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): if hasattr(module, "set_processor"): if not isinstance(processor, dict): - module.set_processor(processor, _remove_lora=_remove_lora) + module.set_processor(processor) else: - module.set_processor(processor.pop(f"{name}.processor"), _remove_lora=_remove_lora) + module.set_processor(processor.pop(f"{name}.processor")) for sub_name, child in module.named_children(): fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) diff --git a/src/diffusers/pipelines/wuerstchen/modeling_wuerstchen_prior.py b/src/diffusers/pipelines/wuerstchen/modeling_wuerstchen_prior.py index a7d9e32fb6c9..42bbf2e41c52 100644 --- a/src/diffusers/pipelines/wuerstchen/modeling_wuerstchen_prior.py +++ b/src/diffusers/pipelines/wuerstchen/modeling_wuerstchen_prior.py @@ -117,9 +117,9 @@ def set_attn_processor( def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): if hasattr(module, "set_processor"): if not isinstance(processor, dict): - module.set_processor(processor, _remove_lora=_remove_lora) + module.set_processor(processor) else: - module.set_processor(processor.pop(f"{name}.processor"), _remove_lora=_remove_lora) + module.set_processor(processor.pop(f"{name}.processor")) for sub_name, child in module.named_children(): fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) From fc550651a291fba0754cad861abc7b9b1a72f4b3 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Tue, 2 Jan 2024 19:00:52 +0530 Subject: [PATCH 02/11] fix: copies --- src/diffusers/models/autoencoders/autoencoder_kl.py | 6 ++---- .../models/autoencoders/autoencoder_kl_temporal_decoder.py | 4 +--- .../models/autoencoders/consistency_decoder_vae.py | 6 ++---- src/diffusers/models/unet_2d_condition.py | 6 ++---- src/diffusers/models/unet_3d_condition.py | 6 ++---- src/diffusers/models/unet_motion_model.py | 6 ++---- src/diffusers/models/uvit_2d.py | 6 ++---- src/diffusers/pipelines/audioldm2/modeling_audioldm2.py | 6 ++---- .../deprecated/versatile_diffusion/modeling_text_unet.py | 6 ++---- .../pipelines/wuerstchen/modeling_wuerstchen_prior.py | 6 ++---- 10 files changed, 19 insertions(+), 39 deletions(-) diff --git a/src/diffusers/models/autoencoders/autoencoder_kl.py b/src/diffusers/models/autoencoders/autoencoder_kl.py index e601be0331eb..10a3ae58de9f 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl.py @@ -182,9 +182,7 @@ def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: return processors # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_attn_processor - def set_attn_processor( - self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]], _remove_lora=False - ): + def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): r""" Sets the attention processor to use to compute attention. @@ -232,7 +230,7 @@ def set_default_attn_processor(self): f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}" ) - self.set_attn_processor(processor, _remove_lora=True) + self.set_attn_processor(processor) @apply_forward_hook def encode( diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_temporal_decoder.py b/src/diffusers/models/autoencoders/autoencoder_kl_temporal_decoder.py index ab44ae5e0bcf..f16d1c15a685 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_temporal_decoder.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_temporal_decoder.py @@ -267,9 +267,7 @@ def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: return processors # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_attn_processor - def set_attn_processor( - self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]], _remove_lora=False - ): + def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): r""" Sets the attention processor to use to compute attention. diff --git a/src/diffusers/models/autoencoders/consistency_decoder_vae.py b/src/diffusers/models/autoencoders/consistency_decoder_vae.py index bb960504c324..ca670fec4b28 100644 --- a/src/diffusers/models/autoencoders/consistency_decoder_vae.py +++ b/src/diffusers/models/autoencoders/consistency_decoder_vae.py @@ -212,9 +212,7 @@ def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: return processors # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_attn_processor - def set_attn_processor( - self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]], _remove_lora=False - ): + def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): r""" Sets the attention processor to use to compute attention. @@ -262,7 +260,7 @@ def set_default_attn_processor(self): f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}" ) - self.set_attn_processor(processor, _remove_lora=True) + self.set_attn_processor(processor) @apply_forward_hook def encode( diff --git a/src/diffusers/models/unet_2d_condition.py b/src/diffusers/models/unet_2d_condition.py index 397b11b74073..4554016bdd53 100644 --- a/src/diffusers/models/unet_2d_condition.py +++ b/src/diffusers/models/unet_2d_condition.py @@ -643,9 +643,7 @@ def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: return processors - def set_attn_processor( - self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]], _remove_lora=False - ): + def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): r""" Sets the attention processor to use to compute attention. @@ -692,7 +690,7 @@ def set_default_attn_processor(self): f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}" ) - self.set_attn_processor(processor, _remove_lora=True) + self.set_attn_processor(processor) def set_attention_slice(self, slice_size): r""" diff --git a/src/diffusers/models/unet_3d_condition.py b/src/diffusers/models/unet_3d_condition.py index 1b49a746858a..2fd629b2f8d9 100644 --- a/src/diffusers/models/unet_3d_condition.py +++ b/src/diffusers/models/unet_3d_condition.py @@ -375,9 +375,7 @@ def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: List[i fn_recursive_set_attention_slice(module, reversed_slice_size) # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_attn_processor - def set_attn_processor( - self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]], _remove_lora=False - ): + def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): r""" Sets the attention processor to use to compute attention. @@ -465,7 +463,7 @@ def set_default_attn_processor(self): f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}" ) - self.set_attn_processor(processor, _remove_lora=True) + self.set_attn_processor(processor) def _set_gradient_checkpointing(self, module, value: bool = False) -> None: if isinstance(module, (CrossAttnDownBlock3D, DownBlock3D, CrossAttnUpBlock3D, UpBlock3D)): diff --git a/src/diffusers/models/unet_motion_model.py b/src/diffusers/models/unet_motion_model.py index 4375515ae043..b5f0302b4a43 100644 --- a/src/diffusers/models/unet_motion_model.py +++ b/src/diffusers/models/unet_motion_model.py @@ -549,9 +549,7 @@ def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: return processors # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_attn_processor - def set_attn_processor( - self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]], _remove_lora=False - ): + def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): r""" Sets the attention processor to use to compute attention. @@ -641,7 +639,7 @@ def set_default_attn_processor(self) -> None: f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}" ) - self.set_attn_processor(processor, _remove_lora=True) + self.set_attn_processor(processor) def _set_gradient_checkpointing(self, module, value: bool = False) -> None: if isinstance(module, (CrossAttnDownBlockMotion, DownBlockMotion, CrossAttnUpBlockMotion, UpBlockMotion)): diff --git a/src/diffusers/models/uvit_2d.py b/src/diffusers/models/uvit_2d.py index 9df01aebf229..a49c77a51b02 100644 --- a/src/diffusers/models/uvit_2d.py +++ b/src/diffusers/models/uvit_2d.py @@ -237,9 +237,7 @@ def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: return processors # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_attn_processor - def set_attn_processor( - self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]], _remove_lora=False - ): + def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): r""" Sets the attention processor to use to compute attention. @@ -287,7 +285,7 @@ def set_default_attn_processor(self): f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}" ) - self.set_attn_processor(processor, _remove_lora=True) + self.set_attn_processor(processor) class UVit2DConvEmbed(nn.Module): diff --git a/src/diffusers/pipelines/audioldm2/modeling_audioldm2.py b/src/diffusers/pipelines/audioldm2/modeling_audioldm2.py index 4c688640fbba..d39b2c99ddd0 100644 --- a/src/diffusers/pipelines/audioldm2/modeling_audioldm2.py +++ b/src/diffusers/pipelines/audioldm2/modeling_audioldm2.py @@ -538,9 +538,7 @@ def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: return processors # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_attn_processor - def set_attn_processor( - self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]], _remove_lora=False - ): + def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): r""" Sets the attention processor to use to compute attention. @@ -588,7 +586,7 @@ def set_default_attn_processor(self): f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}" ) - self.set_attn_processor(processor, _remove_lora=True) + self.set_attn_processor(processor) # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_attention_slice def set_attention_slice(self, slice_size): diff --git a/src/diffusers/pipelines/deprecated/versatile_diffusion/modeling_text_unet.py b/src/diffusers/pipelines/deprecated/versatile_diffusion/modeling_text_unet.py index fc92db7236eb..6f95112c3d50 100644 --- a/src/diffusers/pipelines/deprecated/versatile_diffusion/modeling_text_unet.py +++ b/src/diffusers/pipelines/deprecated/versatile_diffusion/modeling_text_unet.py @@ -848,9 +848,7 @@ def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: return processors - def set_attn_processor( - self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]], _remove_lora=False - ): + def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): r""" Sets the attention processor to use to compute attention. @@ -897,7 +895,7 @@ def set_default_attn_processor(self): f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}" ) - self.set_attn_processor(processor, _remove_lora=True) + self.set_attn_processor(processor) def set_attention_slice(self, slice_size): r""" diff --git a/src/diffusers/pipelines/wuerstchen/modeling_wuerstchen_prior.py b/src/diffusers/pipelines/wuerstchen/modeling_wuerstchen_prior.py index 42bbf2e41c52..d4502639cebc 100644 --- a/src/diffusers/pipelines/wuerstchen/modeling_wuerstchen_prior.py +++ b/src/diffusers/pipelines/wuerstchen/modeling_wuerstchen_prior.py @@ -91,9 +91,7 @@ def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: return processors # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_attn_processor - def set_attn_processor( - self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]], _remove_lora=False - ): + def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): r""" Sets the attention processor to use to compute attention. @@ -141,7 +139,7 @@ def set_default_attn_processor(self): f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}" ) - self.set_attn_processor(processor, _remove_lora=True) + self.set_attn_processor(processor) def _set_gradient_checkpointing(self, module, value=False): self.gradient_checkpointing = value From 504800de4d6ee62e45db35f7b1ddcacf23eb6767 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Wed, 3 Jan 2024 10:20:11 +0530 Subject: [PATCH 03/11] don't modify the uNet in-place. --- tests/lora/test_lora_layers_old_backend.py | 43 ++++++++++++++-------- 1 file changed, 27 insertions(+), 16 deletions(-) diff --git a/tests/lora/test_lora_layers_old_backend.py b/tests/lora/test_lora_layers_old_backend.py index 7d6d30169455..0d9bbc1e302f 100644 --- a/tests/lora/test_lora_layers_old_backend.py +++ b/tests/lora/test_lora_layers_old_backend.py @@ -61,7 +61,8 @@ ) -def text_encoder_attn_modules(text_encoder): +def text_encoder_attn_modules(text_encoder: nn.Module): + """Fetches the attention modules from `text_encoder`.""" attn_modules = [] if isinstance(text_encoder, (CLIPTextModel, CLIPTextModelWithProjection)): @@ -75,7 +76,8 @@ def text_encoder_attn_modules(text_encoder): return attn_modules -def text_encoder_lora_state_dict(text_encoder): +def text_encoder_lora_state_dict(text_encoder: nn.Module): + """Returns the LoRA state dict of the `text_encoder`. Assumes that `_modify_text_encoder()` was already called on it.""" state_dict = {} for name, module in text_encoder_attn_modules(text_encoder): @@ -95,11 +97,15 @@ def text_encoder_lora_state_dict(text_encoder): def create_unet_lora_layers(unet: nn.Module, rank=4, mock_weights=True): + """Creates and returns the LoRA state dict for the UNet.""" + # So that we accidentally don't end up using the in-place modified UNet. + copied_unet = copy.deepcopy(unet) + unet_lora_parameters = [] - for attn_processor_name, attn_processor in unet.attn_processors.items(): + for attn_processor_name, attn_processor in copied_unet.attn_processors.items(): # Parse the attention module. - attn_module = unet + attn_module = copied_unet for n in attn_processor_name.split(".")[:-1]: attn_module = getattr(attn_module, n) @@ -145,30 +151,33 @@ def create_unet_lora_layers(unet: nn.Module, rank=4, mock_weights=True): unet_lora_parameters.extend(attn_module.to_v.lora_layer.parameters()) unet_lora_parameters.extend(attn_module.to_out[0].lora_layer.parameters()) - return unet_lora_parameters, unet_lora_state_dict(unet) + return unet_lora_parameters, unet_lora_state_dict(copied_unet) def create_3d_unet_lora_layers(unet: nn.Module, rank=4, mock_weights=True): - for attn_processor_name in unet.attn_processors.keys(): + """Creates and returns the LoRA state dict for the 3D UNet.""" + copied_unet = copy.deepcopy(unet) + + for attn_processor_name in copied_unet.attn_processors.keys(): has_cross_attention = attn_processor_name.endswith("attn2.processor") and not ( attn_processor_name.startswith("transformer_in") or "temp_attentions" in attn_processor_name.split(".") ) - cross_attention_dim = unet.config.cross_attention_dim if has_cross_attention else None + cross_attention_dim = copied_unet.config.cross_attention_dim if has_cross_attention else None if attn_processor_name.startswith("mid_block"): - hidden_size = unet.config.block_out_channels[-1] + hidden_size = copied_unet.config.block_out_channels[-1] elif attn_processor_name.startswith("up_blocks"): block_id = int(attn_processor_name[len("up_blocks.")]) - hidden_size = list(reversed(unet.config.block_out_channels))[block_id] + hidden_size = list(reversed(copied_unet.config.block_out_channels))[block_id] elif attn_processor_name.startswith("down_blocks"): block_id = int(attn_processor_name[len("down_blocks.")]) - hidden_size = unet.config.block_out_channels[block_id] + hidden_size = copied_unet.config.block_out_channels[block_id] elif attn_processor_name.startswith("transformer_in"): # Note that the `8 * ...` comes from: https://github.com/huggingface/diffusers/blob/7139f0e874f10b2463caa8cbd585762a309d12d6/src/diffusers/models/unet_3d_condition.py#L148 - hidden_size = 8 * unet.config.attention_head_dim + hidden_size = 8 * copied_unet.config.attention_head_dim # Parse the attention module. - attn_module = unet + attn_module = copied_unet for n in attn_processor_name.split(".")[:-1]: attn_module = getattr(attn_module, n) @@ -216,10 +225,11 @@ def create_3d_unet_lora_layers(unet: nn.Module, rank=4, mock_weights=True): attn_module.to_v.lora_layer.up.weight += 1 attn_module.to_out[0].lora_layer.up.weight += 1 - return unet_lora_state_dict(unet) + return unet_lora_state_dict(copied_unet) def set_lora_weights(lora_attn_parameters, randn_weight=False, var=1.0): + """Randomizes the LoRA params if specified.""" if not isinstance(lora_attn_parameters, dict): with torch.no_grad(): for parameter in lora_attn_parameters: @@ -1441,6 +1451,7 @@ def test_save_load_fused_lora_modules(self): class UNet2DConditionLoRAModelTests(unittest.TestCase): model_class = UNet2DConditionModel main_input_name = "sample" + lora_rank = 4 @property def dummy_input(self): @@ -1489,7 +1500,7 @@ def test_lora_processors(self): with torch.no_grad(): sample1 = model(**inputs_dict).sample - _, lora_params = create_unet_lora_layers(model) + _, lora_params = create_unet_lora_layers(model, rank=self.lora_rank) # make sure we can set a list of attention processors model.load_attn_procs(lora_params) @@ -1522,7 +1533,7 @@ def test_lora_on_off(self, expected_max_diff=1e-3): with torch.no_grad(): old_sample = model(**inputs_dict).sample - _, lora_params = create_unet_lora_layers(model) + _, lora_params = create_unet_lora_layers(model, rank=self.lora_rank) model.load_attn_procs(lora_params) with torch.no_grad(): @@ -1552,7 +1563,7 @@ def test_lora_xformers_on_off(self, expected_max_diff=6e-4): torch.manual_seed(0) model = self.model_class(**init_dict) model.to(torch_device) - _, lora_params = create_unet_lora_layers(model) + _, lora_params = create_unet_lora_layers(model, rank=self.lora_rank) model.load_attn_procs(lora_params) # default From c5edc4b592ef6f860c173965a1b5ec5a5fe912f2 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Wed, 3 Jan 2024 10:30:17 +0530 Subject: [PATCH 04/11] fix: temporal autoencoder. --- src/diffusers/loaders/lora.py | 2 +- .../models/autoencoders/autoencoder_kl_temporal_decoder.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/diffusers/loaders/lora.py b/src/diffusers/loaders/lora.py index bbd01a995061..424e95f0843e 100644 --- a/src/diffusers/loaders/lora.py +++ b/src/diffusers/loaders/lora.py @@ -980,7 +980,7 @@ def unload_lora_weights(self): if not USE_PEFT_BACKEND: if version.parse(__version__) > version.parse("0.23"): - logger.warn( + logger.warning( "You are using `unload_lora_weights` to disable and unload lora weights. If you want to iteratively enable and disable adapter weights," "you can use `pipe.enable_lora()` or `pipe.disable_lora()`. After installing the latest version of PEFT." ) diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_temporal_decoder.py b/src/diffusers/models/autoencoders/autoencoder_kl_temporal_decoder.py index f16d1c15a685..dbafb4571d4a 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_temporal_decoder.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_temporal_decoder.py @@ -312,7 +312,7 @@ def set_default_attn_processor(self): f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}" ) - self.set_attn_processor(processor, _remove_lora=True) + self.set_attn_processor(processor) @apply_forward_hook def encode( From cbb2d112f5b6f9a83d9b38df8aac7d4dceda8266 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Wed, 3 Jan 2024 10:39:31 +0530 Subject: [PATCH 05/11] manually remove lora layers. --- tests/lora/test_lora_layers_old_backend.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/lora/test_lora_layers_old_backend.py b/tests/lora/test_lora_layers_old_backend.py index 0d9bbc1e302f..8edb03455ac8 100644 --- a/tests/lora/test_lora_layers_old_backend.py +++ b/tests/lora/test_lora_layers_old_backend.py @@ -1539,7 +1539,10 @@ def test_lora_on_off(self, expected_max_diff=1e-3): with torch.no_grad(): sample = model(**inputs_dict, cross_attention_kwargs={"scale": 0.0}).sample - model.set_default_attn_processor() + # Unload LoRA. + for module in model.modules(): + if hasattr(module, "set_lora_layer"): + module.set_lora_layer(None) with torch.no_grad(): new_sample = model(**inputs_dict).sample From eed8e156492b69e762e83b1690b876116ede91c8 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Wed, 3 Jan 2024 18:11:53 +0530 Subject: [PATCH 06/11] don't copy unet. --- tests/lora/test_lora_layers_old_backend.py | 37 +++++++++++++--------- 1 file changed, 22 insertions(+), 15 deletions(-) diff --git a/tests/lora/test_lora_layers_old_backend.py b/tests/lora/test_lora_layers_old_backend.py index 8edb03455ac8..34715d64421c 100644 --- a/tests/lora/test_lora_layers_old_backend.py +++ b/tests/lora/test_lora_layers_old_backend.py @@ -99,13 +99,11 @@ def text_encoder_lora_state_dict(text_encoder: nn.Module): def create_unet_lora_layers(unet: nn.Module, rank=4, mock_weights=True): """Creates and returns the LoRA state dict for the UNet.""" # So that we accidentally don't end up using the in-place modified UNet. - copied_unet = copy.deepcopy(unet) - unet_lora_parameters = [] - for attn_processor_name, attn_processor in copied_unet.attn_processors.items(): + for attn_processor_name, attn_processor in unet.attn_processors.items(): # Parse the attention module. - attn_module = copied_unet + attn_module = unet for n in attn_processor_name.split(".")[:-1]: attn_module = getattr(attn_module, n) @@ -151,33 +149,37 @@ def create_unet_lora_layers(unet: nn.Module, rank=4, mock_weights=True): unet_lora_parameters.extend(attn_module.to_v.lora_layer.parameters()) unet_lora_parameters.extend(attn_module.to_out[0].lora_layer.parameters()) - return unet_lora_parameters, unet_lora_state_dict(copied_unet) + unet_lora_state_dict = unet_lora_state_dict(unet) + # Unload LoRA. + for module in unet.modules(): + if hasattr(module, "set_lora_layer"): + module.set_lora_layer(None) + + return unet_lora_parameters, unet_lora_state_dict def create_3d_unet_lora_layers(unet: nn.Module, rank=4, mock_weights=True): """Creates and returns the LoRA state dict for the 3D UNet.""" - copied_unet = copy.deepcopy(unet) - - for attn_processor_name in copied_unet.attn_processors.keys(): + for attn_processor_name in unet.attn_processors.keys(): has_cross_attention = attn_processor_name.endswith("attn2.processor") and not ( attn_processor_name.startswith("transformer_in") or "temp_attentions" in attn_processor_name.split(".") ) - cross_attention_dim = copied_unet.config.cross_attention_dim if has_cross_attention else None + cross_attention_dim = unet.config.cross_attention_dim if has_cross_attention else None if attn_processor_name.startswith("mid_block"): - hidden_size = copied_unet.config.block_out_channels[-1] + hidden_size = unet.config.block_out_channels[-1] elif attn_processor_name.startswith("up_blocks"): block_id = int(attn_processor_name[len("up_blocks.")]) - hidden_size = list(reversed(copied_unet.config.block_out_channels))[block_id] + hidden_size = list(reversed(unet.config.block_out_channels))[block_id] elif attn_processor_name.startswith("down_blocks"): block_id = int(attn_processor_name[len("down_blocks.")]) - hidden_size = copied_unet.config.block_out_channels[block_id] + hidden_size = unet.config.block_out_channels[block_id] elif attn_processor_name.startswith("transformer_in"): # Note that the `8 * ...` comes from: https://github.com/huggingface/diffusers/blob/7139f0e874f10b2463caa8cbd585762a309d12d6/src/diffusers/models/unet_3d_condition.py#L148 - hidden_size = 8 * copied_unet.config.attention_head_dim + hidden_size = 8 * unet.config.attention_head_dim # Parse the attention module. - attn_module = copied_unet + attn_module = unet for n in attn_processor_name.split(".")[:-1]: attn_module = getattr(attn_module, n) @@ -225,7 +227,12 @@ def create_3d_unet_lora_layers(unet: nn.Module, rank=4, mock_weights=True): attn_module.to_v.lora_layer.up.weight += 1 attn_module.to_out[0].lora_layer.up.weight += 1 - return unet_lora_state_dict(copied_unet) + # Unload LoRA. + for module in unet.modules(): + if hasattr(module, "set_lora_layer"): + module.set_lora_layer(None) + + return unet_lora_state_dict(unet) def set_lora_weights(lora_attn_parameters, randn_weight=False, var=1.0): From 760130b402eb5baa110ced5568858e09728d61ba Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Wed, 3 Jan 2024 18:35:28 +0530 Subject: [PATCH 07/11] alright --- tests/lora/test_lora_layers_old_backend.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/lora/test_lora_layers_old_backend.py b/tests/lora/test_lora_layers_old_backend.py index 34715d64421c..8d0030c8966a 100644 --- a/tests/lora/test_lora_layers_old_backend.py +++ b/tests/lora/test_lora_layers_old_backend.py @@ -149,13 +149,13 @@ def create_unet_lora_layers(unet: nn.Module, rank=4, mock_weights=True): unet_lora_parameters.extend(attn_module.to_v.lora_layer.parameters()) unet_lora_parameters.extend(attn_module.to_out[0].lora_layer.parameters()) - unet_lora_state_dict = unet_lora_state_dict(unet) + unet_lora_sd = unet_lora_state_dict(unet) # Unload LoRA. for module in unet.modules(): if hasattr(module, "set_lora_layer"): module.set_lora_layer(None) - return unet_lora_parameters, unet_lora_state_dict + return unet_lora_parameters, unet_lora_sd def create_3d_unet_lora_layers(unet: nn.Module, rank=4, mock_weights=True): From cd8b195fb20cc9c236deff72ed1d273d523694dd Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Wed, 3 Jan 2024 18:46:22 +0530 Subject: [PATCH 08/11] remove lora attn processors from unet3d --- tests/lora/test_lora_layers_old_backend.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/tests/lora/test_lora_layers_old_backend.py b/tests/lora/test_lora_layers_old_backend.py index 8d0030c8966a..dfeccd526d00 100644 --- a/tests/lora/test_lora_layers_old_backend.py +++ b/tests/lora/test_lora_layers_old_backend.py @@ -1654,9 +1654,6 @@ def test_lora_processors(self): model.load_attn_procs(unet_lora_params) model.to(torch_device) - # test that attn processors can be set to itself - model.set_attn_processor(model.attn_processors) - with torch.no_grad(): sample2 = model(**inputs_dict, cross_attention_kwargs={"scale": 0.0}).sample sample3 = model(**inputs_dict, cross_attention_kwargs={"scale": 0.5}).sample From 828d65524e15bc4716150bd183d5fabc974d0f8f Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Wed, 3 Jan 2024 18:55:46 +0530 Subject: [PATCH 09/11] fix: unet3d. --- tests/lora/test_lora_layers_old_backend.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/tests/lora/test_lora_layers_old_backend.py b/tests/lora/test_lora_layers_old_backend.py index dfeccd526d00..644f4e44a189 100644 --- a/tests/lora/test_lora_layers_old_backend.py +++ b/tests/lora/test_lora_layers_old_backend.py @@ -227,12 +227,14 @@ def create_3d_unet_lora_layers(unet: nn.Module, rank=4, mock_weights=True): attn_module.to_v.lora_layer.up.weight += 1 attn_module.to_out[0].lora_layer.up.weight += 1 - # Unload LoRA. + unet_lora_sd = unet_lora_state_dict(unet) + + # Unload LoRA. for module in unet.modules(): if hasattr(module, "set_lora_layer"): module.set_lora_layer(None) - return unet_lora_state_dict(unet) + return unet_lora_sd def set_lora_weights(lora_attn_parameters, randn_weight=False, var=1.0): @@ -1654,6 +1656,9 @@ def test_lora_processors(self): model.load_attn_procs(unet_lora_params) model.to(torch_device) + # test that attn processors can be set to itself + model.set_attn_processor(model.attn_processors) + with torch.no_grad(): sample2 = model(**inputs_dict, cross_attention_kwargs={"scale": 0.0}).sample sample3 = model(**inputs_dict, cross_attention_kwargs={"scale": 0.5}).sample From cb71deaf5991bb37ea008589621c2b1d2576c95d Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Wed, 3 Jan 2024 18:58:11 +0530 Subject: [PATCH 10/11] styl --- tests/lora/test_lora_layers_old_backend.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/lora/test_lora_layers_old_backend.py b/tests/lora/test_lora_layers_old_backend.py index 644f4e44a189..09bb87c85163 100644 --- a/tests/lora/test_lora_layers_old_backend.py +++ b/tests/lora/test_lora_layers_old_backend.py @@ -227,9 +227,9 @@ def create_3d_unet_lora_layers(unet: nn.Module, rank=4, mock_weights=True): attn_module.to_v.lora_layer.up.weight += 1 attn_module.to_out[0].lora_layer.up.weight += 1 - unet_lora_sd = unet_lora_state_dict(unet) + unet_lora_sd = unet_lora_state_dict(unet) - # Unload LoRA. + # Unload LoRA. for module in unet.modules(): if hasattr(module, "set_lora_layer"): module.set_lora_layer(None) From 5d8691129adca606468816f857098b662893f797 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Wed, 3 Jan 2024 20:22:45 +0530 Subject: [PATCH 11/11] Empty-Commit