From a28a6e5bdaa1f638efbe782bacd16295d11c38a1 Mon Sep 17 00:00:00 2001 From: Fabio Rigano Date: Fri, 15 Dec 2023 18:30:57 +0100 Subject: [PATCH 1/4] Add unload_ip_adapter method --- src/diffusers/loaders/ip_adapter.py | 35 ++++++++++++++++++++++++++++- 1 file changed, 34 insertions(+), 1 deletion(-) diff --git a/src/diffusers/loaders/ip_adapter.py b/src/diffusers/loaders/ip_adapter.py index 158bde436374..eb4311218562 100644 --- a/src/diffusers/loaders/ip_adapter.py +++ b/src/diffusers/loaders/ip_adapter.py @@ -132,7 +132,7 @@ def load_ip_adapter( if keys != ["image_proj", "ip_adapter"]: raise ValueError("Required keys are (`image_proj` and `ip_adapter`) missing from the state dict.") - # load CLIP image encoer here if it has not been registered to the pipeline yet + # load CLIP image encoder here if it has not been registered to the pipeline yet if hasattr(self, "image_encoder") and getattr(self, "image_encoder", None) is None: if not isinstance(pretrained_model_name_or_path_or_dict, dict): logger.info(f"loading image_encoder from {pretrained_model_name_or_path_or_dict}") @@ -141,12 +141,14 @@ def load_ip_adapter( subfolder=os.path.join(subfolder, "image_encoder"), ).to(self.device, dtype=self.dtype) self.image_encoder = image_encoder + self.register_to_config(image_encoder=CLIPVisionModelWithProjection) else: raise ValueError("`image_encoder` cannot be None when using IP Adapters.") # create feature extractor if it has not been registered to the pipeline yet if hasattr(self, "feature_extractor") and getattr(self, "feature_extractor", None) is None: self.feature_extractor = CLIPImageProcessor() + self.register_to_config(feature_extractor=["transformers", "CLIPImageProcessor"]) # load ip-adapter into unet self.unet._load_ip_adapter_weights(state_dict) @@ -155,3 +157,34 @@ def set_ip_adapter_scale(self, scale): for attn_processor in self.unet.attn_processors.values(): if isinstance(attn_processor, (IPAdapterAttnProcessor, IPAdapterAttnProcessor2_0)): attn_processor.scale = scale + + def unload_ip_adapter(self): + """ + Unloads the IP Adapter weights + + Examples: + + ```python + >>> # Assuming `pipeline` is already loaded with the IP Adapter weights. + >>> pipeline.unload_ip_adapter() + >>> ... + ``` + """ + # remove CLIP image encoder + if hasattr(self, "image_encoder"): + self.image_encoder = None + self.register_to_config(image_encoder=[None, None]) + + # remove feature extractor + if hasattr(self, "feature_extractor"): + self.feature_extractor = None + self.register_to_config(feature_extractor=[None, None]) + + # remove hidden encoder + self.unet.encoder_hid_proj = None + self.config.encoder_hid_dim_type = None + + # restore original Unet layers + device = self.unet.device + self.unet = self.unet.from_pretrained(self.config._name_or_path, subfolder="unet", torch_dtype=self.dtype) + self.unet.to(device) From ce9b12aecfb9c600959bad237f801d84e8a94f1b Mon Sep 17 00:00:00 2001 From: Fabio Rigano Date: Sat, 16 Dec 2023 10:12:56 +0100 Subject: [PATCH 2/4] Update attn_processors with original layers --- src/diffusers/loaders/ip_adapter.py | 10 ++++------ src/diffusers/loaders/unet.py | 13 +++++++++++++ 2 files changed, 17 insertions(+), 6 deletions(-) diff --git a/src/diffusers/loaders/ip_adapter.py b/src/diffusers/loaders/ip_adapter.py index eb4311218562..b4518325985d 100644 --- a/src/diffusers/loaders/ip_adapter.py +++ b/src/diffusers/loaders/ip_adapter.py @@ -171,12 +171,12 @@ def unload_ip_adapter(self): ``` """ # remove CLIP image encoder - if hasattr(self, "image_encoder"): + if hasattr(self, "image_encoder") and getattr(self, "image_encoder", None) is not None: self.image_encoder = None self.register_to_config(image_encoder=[None, None]) # remove feature extractor - if hasattr(self, "feature_extractor"): + if hasattr(self, "feature_extractor") and getattr(self, "feature_extractor", None) is not None: self.feature_extractor = None self.register_to_config(feature_extractor=[None, None]) @@ -184,7 +184,5 @@ def unload_ip_adapter(self): self.unet.encoder_hid_proj = None self.config.encoder_hid_dim_type = None - # restore original Unet layers - device = self.unet.device - self.unet = self.unet.from_pretrained(self.config._name_or_path, subfolder="unet", torch_dtype=self.dtype) - self.unet.to(device) + # restore original Unet attention processors layers + self.unet._restore_attn_processors() diff --git a/src/diffusers/loaders/unet.py b/src/diffusers/loaders/unet.py index 7309c3fc709c..bf23c44c56e0 100644 --- a/src/diffusers/loaders/unet.py +++ b/src/diffusers/loaders/unet.py @@ -686,6 +686,13 @@ def _load_ip_adapter_weights(self, state_dict): # because `Resampler` also has `attn_processors`. self.encoder_hid_proj = None + # store original attn_processors if not already stored + if ( + getattr(self, "original_attn_processors", None) is None + or getattr(self, "original_attn_processors", None) == {} + ): + self.original_attn_processors = self.attn_processors + # set ip-adapter cross-attention processors & load state_dict attn_procs = {} key_id = 1 @@ -824,3 +831,9 @@ def _load_ip_adapter_weights(self, state_dict): self.config.encoder_hid_dim_type = "ip_image_proj" delete_adapter_layers + + def _restore_attn_processors(self): + if hasattr(self, "original_attn_processors") and getattr(self, "original_attn_processors", None) is not None: + self.set_attn_processor(self.original_attn_processors) + else: + raise AttributeError("No original attn_processors to restore") From a5fcc6a78b443fa7ee23c3719150bc4f1582c236 Mon Sep 17 00:00:00 2001 From: fabiorigano Date: Tue, 19 Dec 2023 10:16:07 +0100 Subject: [PATCH 3/4] Add test --- src/diffusers/loaders/ip_adapter.py | 2 +- .../test_ip_adapter_stable_diffusion.py | 20 +++++++++++++++++++ 2 files changed, 21 insertions(+), 1 deletion(-) diff --git a/src/diffusers/loaders/ip_adapter.py b/src/diffusers/loaders/ip_adapter.py index b4518325985d..4b2ce60e4471 100644 --- a/src/diffusers/loaders/ip_adapter.py +++ b/src/diffusers/loaders/ip_adapter.py @@ -141,7 +141,7 @@ def load_ip_adapter( subfolder=os.path.join(subfolder, "image_encoder"), ).to(self.device, dtype=self.dtype) self.image_encoder = image_encoder - self.register_to_config(image_encoder=CLIPVisionModelWithProjection) + self.register_to_config(image_encoder=["transformers", "CLIPVisionModelWithProjection"]) else: raise ValueError("`image_encoder` cannot be None when using IP Adapters.") diff --git a/tests/pipelines/ip_adapters/test_ip_adapter_stable_diffusion.py b/tests/pipelines/ip_adapters/test_ip_adapter_stable_diffusion.py index dfc39d61bb08..289d2b7d6573 100644 --- a/tests/pipelines/ip_adapters/test_ip_adapter_stable_diffusion.py +++ b/tests/pipelines/ip_adapters/test_ip_adapter_stable_diffusion.py @@ -31,6 +31,7 @@ StableDiffusionXLInpaintPipeline, StableDiffusionXLPipeline, ) +from diffusers.models.attention_processor import AttnProcessor, AttnProcessor2_0 from diffusers.utils import load_image from diffusers.utils.testing_utils import ( enable_full_determinism, @@ -228,6 +229,25 @@ def test_text_to_image_full_face(self): assert np.allclose(image_slice, expected_slice, atol=1e-4, rtol=1e-4) + def test_unload(self): + image_encoder = self.get_image_encoder(repo_id="h94/IP-Adapter", subfolder="models/image_encoder") + pipeline = StableDiffusionPipeline.from_pretrained( + "runwayml/stable-diffusion-v1-5", image_encoder=image_encoder, safety_checker=None, torch_dtype=self.dtype + ) + pipeline.to(torch_device) + pipeline.load_ip_adapter("h94/IP-Adapter", subfolder="models", weight_name="ip-adapter_sd15.bin") + pipeline.set_ip_adapter_scale(0.7) + + pipeline.unload_ip_adapter() + + assert getattr(pipeline, "image_encoder") is None + assert getattr(pipeline, "feature_extractor") is None + processors = [ + isinstance(attn_proc, (AttnProcessor, AttnProcessor2_0)) + for name, attn_proc in pipeline.unet.attn_processors.items() + ] + assert processors == [True] * len(processors) + @slow @require_torch_gpu From 412194fa7dd2c1733c2dc8d160ac79d62dc15574 Mon Sep 17 00:00:00 2001 From: fabiorigano Date: Tue, 19 Dec 2023 10:56:35 +0100 Subject: [PATCH 4/4] Use set_default_attn_processor --- src/diffusers/loaders/ip_adapter.py | 2 +- src/diffusers/loaders/unet.py | 13 ------------- 2 files changed, 1 insertion(+), 14 deletions(-) diff --git a/src/diffusers/loaders/ip_adapter.py b/src/diffusers/loaders/ip_adapter.py index 4b2ce60e4471..ae15dc0ea876 100644 --- a/src/diffusers/loaders/ip_adapter.py +++ b/src/diffusers/loaders/ip_adapter.py @@ -185,4 +185,4 @@ def unload_ip_adapter(self): self.config.encoder_hid_dim_type = None # restore original Unet attention processors layers - self.unet._restore_attn_processors() + self.unet.set_default_attn_processor() diff --git a/src/diffusers/loaders/unet.py b/src/diffusers/loaders/unet.py index 669e00fdea9c..7dec43571b1c 100644 --- a/src/diffusers/loaders/unet.py +++ b/src/diffusers/loaders/unet.py @@ -760,13 +760,6 @@ def _load_ip_adapter_weights(self, state_dict): # because `Resampler` also has `attn_processors`. self.encoder_hid_proj = None - # store original attn_processors if not already stored - if ( - getattr(self, "original_attn_processors", None) is None - or getattr(self, "original_attn_processors", None) == {} - ): - self.original_attn_processors = self.attn_processors - # set ip-adapter cross-attention processors & load state_dict attn_procs = {} key_id = 1 @@ -810,9 +803,3 @@ def _load_ip_adapter_weights(self, state_dict): self.encoder_hid_proj = image_projection.to(device=self.device, dtype=self.dtype) self.config.encoder_hid_dim_type = "ip_image_proj" - - def _restore_attn_processors(self): - if hasattr(self, "original_attn_processors") and getattr(self, "original_attn_processors", None) is not None: - self.set_attn_processor(self.original_attn_processors) - else: - raise AttributeError("No original attn_processors to restore")