From 2fbdeb10bf21ed84e1080db11a516c94acbec929 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Tue, 6 Oct 2026 10:57:00 +0000 Subject: [PATCH] remove lora utils from tests/lora --- tests/lora/test_lora_layers_mochi.py | 120 - tests/lora/test_lora_layers_sana.py | 116 - tests/lora/test_lora_layers_sd3.py | 186 -- tests/lora/utils.py | 1973 ----------------- tests/pipelines/mochi/test_mochi.py | 13 + tests/pipelines/sana/test_sana.py | 10 + .../test_pipeline_stable_diffusion_3.py | 18 + ...est_pipeline_stable_diffusion_3_img2img.py | 58 + utils/extract_tests_from_mixin.py | 4 +- 9 files changed, 101 insertions(+), 2397 deletions(-) delete mode 100644 tests/lora/test_lora_layers_mochi.py delete mode 100644 tests/lora/test_lora_layers_sana.py delete mode 100644 tests/lora/test_lora_layers_sd3.py delete mode 100644 tests/lora/utils.py diff --git a/tests/lora/test_lora_layers_mochi.py b/tests/lora/test_lora_layers_mochi.py deleted file mode 100644 index 1b806c191fbf..000000000000 --- a/tests/lora/test_lora_layers_mochi.py +++ /dev/null @@ -1,120 +0,0 @@ -# Copyright 2026 HuggingFace Inc. -# -# 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. - -import sys -import unittest - -import torch -from transformers import AutoTokenizer, T5EncoderModel - -from diffusers import AutoencoderKLMochi, FlowMatchEulerDiscreteScheduler, MochiPipeline, MochiTransformer3DModel - -from ..testing_utils import ( - floats_tensor, - require_peft_backend, - skip_mps, -) - - -sys.path.append(".") - -from .utils import PeftLoraLoaderMixinTests # noqa: E402 - - -@require_peft_backend -@skip_mps -class MochiLoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): - pipeline_class = MochiPipeline - scheduler_cls = FlowMatchEulerDiscreteScheduler - scheduler_kwargs = {} - - transformer_kwargs = { - "patch_size": 2, - "num_attention_heads": 2, - "attention_head_dim": 8, - "num_layers": 2, - "pooled_projection_dim": 16, - "in_channels": 12, - "out_channels": None, - "qk_norm": "rms_norm", - "text_embed_dim": 32, - "time_embed_dim": 4, - "activation_fn": "swiglu", - "max_sequence_length": 16, - } - transformer_cls = MochiTransformer3DModel - vae_kwargs = { - "latent_channels": 12, - "out_channels": 3, - "encoder_block_out_channels": (32, 32, 32, 32), - "decoder_block_out_channels": (32, 32, 32, 32), - "layers_per_block": (1, 1, 1, 1, 1), - } - vae_cls = AutoencoderKLMochi - tokenizer_cls, tokenizer_id = AutoTokenizer, "hf-internal-testing/tiny-random-t5" - text_encoder_cls, text_encoder_id = T5EncoderModel, "hf-internal-testing/tiny-random-t5" - - text_encoder_target_modules = ["q", "k", "v", "o"] - - supports_text_encoder_loras = False - - @property - def output_shape(self): - return (1, 7, 16, 16, 3) - - def get_dummy_inputs(self, with_generator=True): - batch_size = 1 - sequence_length = 16 - num_channels = 4 - num_frames = 7 - num_latent_frames = 3 - sizes = (2, 2) - - generator = torch.manual_seed(0) - noise = floats_tensor((batch_size, num_latent_frames, num_channels) + sizes) - input_ids = torch.randint(1, sequence_length, size=(batch_size, sequence_length), generator=generator) - - pipeline_inputs = { - "prompt": "dance monkey", - "num_frames": num_frames, - "num_inference_steps": 4, - "guidance_scale": 6.0, - # Cannot reduce because convolution kernel becomes bigger than sample - "height": 16, - "width": 16, - "max_sequence_length": sequence_length, - "output_type": "np", - } - if with_generator: - pipeline_inputs.update({"generator": generator}) - - return noise, input_ids, pipeline_inputs - - def test_simple_inference_with_text_lora_denoiser_fused_multi(self): - super().test_simple_inference_with_text_lora_denoiser_fused_multi(expected_atol=9e-3) - - def test_simple_inference_with_text_denoiser_lora_unfused(self): - super().test_simple_inference_with_text_denoiser_lora_unfused(expected_atol=9e-3) - - @unittest.skip("Not supported in Mochi.") - def test_simple_inference_with_text_denoiser_block_scale(self): - pass - - @unittest.skip("Not supported in Mochi.") - def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(self): - pass - - @unittest.skip("Not supported in CogVideoX.") - def test_simple_inference_with_text_denoiser_multi_adapter_block_lora(self): - pass diff --git a/tests/lora/test_lora_layers_sana.py b/tests/lora/test_lora_layers_sana.py deleted file mode 100644 index 46f1b1a51e2e..000000000000 --- a/tests/lora/test_lora_layers_sana.py +++ /dev/null @@ -1,116 +0,0 @@ -# coding=utf-8 -# Copyright 2026 HuggingFace Inc. -# -# 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. -import sys -import unittest - -import torch -from transformers import Gemma2Model, GemmaTokenizer - -from diffusers import AutoencoderDC, FlowMatchEulerDiscreteScheduler, SanaPipeline, SanaTransformer2DModel - -from ..testing_utils import floats_tensor, require_peft_backend - - -sys.path.append(".") - -from .utils import PeftLoraLoaderMixinTests # noqa: E402 - - -@require_peft_backend -class SanaLoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): - pipeline_class = SanaPipeline - scheduler_cls = FlowMatchEulerDiscreteScheduler - scheduler_kwargs = {"shift": 7.0} - transformer_kwargs = { - "patch_size": 1, - "in_channels": 4, - "out_channels": 4, - "num_layers": 1, - "num_attention_heads": 2, - "attention_head_dim": 4, - "num_cross_attention_heads": 2, - "cross_attention_head_dim": 4, - "cross_attention_dim": 8, - "caption_channels": 8, - "sample_size": 32, - } - transformer_cls = SanaTransformer2DModel - vae_kwargs = { - "in_channels": 3, - "latent_channels": 4, - "attention_head_dim": 2, - "encoder_block_types": ( - "ResBlock", - "EfficientViTBlock", - ), - "decoder_block_types": ( - "ResBlock", - "EfficientViTBlock", - ), - "encoder_block_out_channels": (8, 8), - "decoder_block_out_channels": (8, 8), - "encoder_qkv_multiscales": ((), (5,)), - "decoder_qkv_multiscales": ((), (5,)), - "encoder_layers_per_block": (1, 1), - "decoder_layers_per_block": [1, 1], - "downsample_block_type": "conv", - "upsample_block_type": "interpolate", - "decoder_norm_types": "rms_norm", - "decoder_act_fns": "silu", - "scaling_factor": 0.41407, - } - vae_cls = AutoencoderDC - tokenizer_cls, tokenizer_id = GemmaTokenizer, "hf-internal-testing/dummy-gemma" - text_encoder_cls, text_encoder_id = Gemma2Model, "hf-internal-testing/dummy-gemma-for-diffusers" - - supports_text_encoder_loras = False - - @property - def output_shape(self): - return (1, 32, 32, 3) - - def get_dummy_inputs(self, with_generator=True): - batch_size = 1 - sequence_length = 16 - num_channels = 4 - sizes = (32, 32) - - generator = torch.manual_seed(0) - noise = floats_tensor((batch_size, num_channels) + sizes) - input_ids = torch.randint(1, sequence_length, size=(batch_size, sequence_length), generator=generator) - - pipeline_inputs = { - "prompt": "", - "negative_prompt": "", - "num_inference_steps": 4, - "guidance_scale": 4.5, - "height": 32, - "width": 32, - "max_sequence_length": sequence_length, - "output_type": "np", - "complex_human_instruction": None, - } - if with_generator: - pipeline_inputs.update({"generator": generator}) - - return noise, input_ids, pipeline_inputs - - @unittest.skip("Not supported in SANA.") - def test_simple_inference_with_text_denoiser_block_scale(self): - pass - - @unittest.skip("Not supported in SANA.") - def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(self): - pass diff --git a/tests/lora/test_lora_layers_sd3.py b/tests/lora/test_lora_layers_sd3.py deleted file mode 100644 index 0f078f9bad0e..000000000000 --- a/tests/lora/test_lora_layers_sd3.py +++ /dev/null @@ -1,186 +0,0 @@ -# coding=utf-8 -# Copyright 2026 HuggingFace Inc. -# -# 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. -import gc -import sys -import unittest - -import numpy as np -import torch -from transformers import AutoTokenizer, CLIPTextModelWithProjection, CLIPTokenizer, T5EncoderModel - -from diffusers import ( - FlowMatchEulerDiscreteScheduler, - SD3Transformer2DModel, - StableDiffusion3Img2ImgPipeline, - StableDiffusion3Pipeline, -) -from diffusers.utils import load_image -from diffusers.utils.import_utils import is_accelerate_available - -from ..testing_utils import ( - backend_empty_cache, - is_flaky, - nightly, - numpy_cosine_similarity_distance, - require_big_accelerator, - require_peft_backend, - require_torch_accelerator, - torch_device, -) - - -sys.path.append(".") - -from .utils import PeftLoraLoaderMixinTests # noqa: E402 - - -if is_accelerate_available(): - from accelerate.utils import release_memory - - -@require_peft_backend -class SD3LoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): - pipeline_class = StableDiffusion3Pipeline - scheduler_cls = FlowMatchEulerDiscreteScheduler - scheduler_kwargs = {} - transformer_kwargs = { - "sample_size": 32, - "patch_size": 1, - "in_channels": 4, - "num_layers": 1, - "attention_head_dim": 8, - "num_attention_heads": 4, - "caption_projection_dim": 32, - "joint_attention_dim": 32, - "pooled_projection_dim": 64, - "out_channels": 4, - } - transformer_cls = SD3Transformer2DModel - vae_kwargs = { - "sample_size": 32, - "in_channels": 3, - "out_channels": 3, - "block_out_channels": (4,), - "layers_per_block": 1, - "latent_channels": 4, - "norm_num_groups": 1, - "use_quant_conv": False, - "use_post_quant_conv": False, - "shift_factor": 0.0609, - "scaling_factor": 1.5035, - } - has_three_text_encoders = True - tokenizer_cls, tokenizer_id = CLIPTokenizer, "hf-internal-testing/tiny-random-clip" - tokenizer_2_cls, tokenizer_2_id = CLIPTokenizer, "hf-internal-testing/tiny-random-clip" - tokenizer_3_cls, tokenizer_3_id = AutoTokenizer, "hf-internal-testing/tiny-random-t5" - text_encoder_cls, text_encoder_id = CLIPTextModelWithProjection, "hf-internal-testing/tiny-sd3-text_encoder" - text_encoder_2_cls, text_encoder_2_id = CLIPTextModelWithProjection, "hf-internal-testing/tiny-sd3-text_encoder-2" - text_encoder_3_cls, text_encoder_3_id = T5EncoderModel, "hf-internal-testing/tiny-random-t5" - - @property - def output_shape(self): - return (1, 32, 32, 3) - - @require_torch_accelerator - def test_sd3_lora(self): - """ - Test loading the loras that are saved with the diffusers and peft formats. - Related PR: https://github.com/huggingface/diffusers/pull/8584 - """ - components = self.get_dummy_components() - pipe = self.pipeline_class(**components[0]) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - - lora_model_id = "hf-internal-testing/tiny-sd3-loras" - - lora_filename = "lora_diffusers_format.safetensors" - pipe.load_lora_weights(lora_model_id, weight_name=lora_filename) - pipe.unload_lora_weights() - - lora_filename = "lora_peft_format.safetensors" - pipe.load_lora_weights(lora_model_id, weight_name=lora_filename) - - @unittest.skip("Not supported in SD3.") - def test_simple_inference_with_text_denoiser_block_scale(self): - pass - - @unittest.skip("Not supported in SD3.") - def test_simple_inference_with_text_denoiser_multi_adapter_block_lora(self): - pass - - @unittest.skip("Not supported in SD3.") - def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(self): - pass - - @is_flaky - def test_multiple_wrong_adapter_name_raises_error(self): - super().test_multiple_wrong_adapter_name_raises_error() - - -@nightly -@require_torch_accelerator -@require_peft_backend -@require_big_accelerator -class SD3LoraIntegrationTests(unittest.TestCase): - pipeline_class = StableDiffusion3Img2ImgPipeline - repo_id = "stabilityai/stable-diffusion-3-medium-diffusers" - - def setUp(self): - super().setUp() - gc.collect() - backend_empty_cache(torch_device) - - def tearDown(self): - super().tearDown() - gc.collect() - backend_empty_cache(torch_device) - - def get_inputs(self, device, seed=0): - init_image = load_image( - "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/bird_canny.png" - ) - if str(device).startswith("mps"): - generator = torch.manual_seed(seed) - else: - generator = torch.Generator(device="cpu").manual_seed(seed) - - return { - "prompt": "corgi", - "num_inference_steps": 2, - "guidance_scale": 5.0, - "output_type": "np", - "generator": generator, - "image": init_image, - } - - def test_sd3_img2img_lora(self): - pipe = self.pipeline_class.from_pretrained(self.repo_id, torch_dtype=torch.float16) - pipe.load_lora_weights("zwloong/sd3-lora-training-rank16-v2") - pipe.fuse_lora() - pipe.unload_lora_weights() - pipe = pipe.to(torch_device) - - inputs = self.get_inputs(torch_device) - - image = pipe(**inputs).images[0] - image_slice = image[0, -3:, -3:] - expected_slice = np.array([0.5649, 0.5405, 0.5488, 0.5688, 0.5449, 0.5513, 0.5337, 0.5107, 0.5059]) - - max_diff = numpy_cosine_similarity_distance(expected_slice.flatten(), image_slice.flatten()) - - assert max_diff < 1e-4, f"Outputs are not close enough, got {max_diff}" - pipe.unload_lora_weights() - release_memory(pipe) diff --git a/tests/lora/utils.py b/tests/lora/utils.py deleted file mode 100644 index dacfd61cc9f0..000000000000 --- a/tests/lora/utils.py +++ /dev/null @@ -1,1973 +0,0 @@ -# coding=utf-8 -# Copyright 2026 HuggingFace Inc. -# -# 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. -import inspect -import os -import tempfile -import unittest -from itertools import product - -import numpy as np -import pytest -import torch -from parameterized import parameterized - -from diffusers import ( - AutoencoderKL, - UNet2DConditionModel, -) -from diffusers.hooks.group_offloading import _GROUP_OFFLOADING, apply_group_offloading -from diffusers.utils import logging -from diffusers.utils.import_utils import is_peft_available - -from ..testing_utils import ( - CaptureLogger, - check_if_dicts_are_equal, - floats_tensor, - require_peft_backend, - require_peft_version_greater, - require_torch_accelerator, - require_transformers_version_greater, - torch_device, -) - - -if is_peft_available(): - from peft import LoraConfig - from peft.tuners.tuners_utils import BaseTunerLayer - from peft.utils import get_peft_model_state_dict - - -def state_dicts_almost_equal(sd1, sd2): - sd1 = dict(sorted(sd1.items())) - sd2 = dict(sorted(sd2.items())) - - models_are_equal = True - for ten1, ten2 in zip(sd1.values(), sd2.values()): - if (ten1 - ten2).abs().max() > 1e-3: - models_are_equal = False - - return models_are_equal - - -def check_if_lora_correctly_set(model) -> bool: - """ - Checks if the LoRA layers are correctly set with peft - """ - for module in model.modules(): - if isinstance(module, BaseTunerLayer): - return True - return False - - -def check_module_lora_metadata(parsed_metadata: dict, lora_metadatas: dict, module_key: str): - extracted = { - k.removeprefix(f"{module_key}."): v for k, v in parsed_metadata.items() if k.startswith(f"{module_key}.") - } - check_if_dicts_are_equal(extracted, lora_metadatas[f"{module_key}_lora_adapter_metadata"]) - - -def initialize_dummy_state_dict(state_dict): - if not all(v.device.type == "meta" for _, v in state_dict.items()): - raise ValueError("`state_dict` has non-meta values.") - return {k: torch.randn(v.shape, device=torch_device, dtype=v.dtype) for k, v in state_dict.items()} - - -POSSIBLE_ATTENTION_KWARGS_NAMES = ["cross_attention_kwargs", "joint_attention_kwargs", "attention_kwargs"] - - -def determine_attention_kwargs_name(pipeline_class): - call_signature_keys = inspect.signature(pipeline_class.__call__).parameters.keys() - - # TODO(diffusers): Discuss a common naming convention across library for 1.0.0 release - for possible_attention_kwargs in POSSIBLE_ATTENTION_KWARGS_NAMES: - if possible_attention_kwargs in call_signature_keys: - attention_kwargs_name = possible_attention_kwargs - break - assert attention_kwargs_name is not None - return attention_kwargs_name - - -@require_peft_backend -class PeftLoraLoaderMixinTests: - pipeline_class = None - - scheduler_cls = None - scheduler_kwargs = None - - has_two_text_encoders = False - has_three_text_encoders = False - text_encoder_cls, text_encoder_id, text_encoder_subfolder = None, None, "" - text_encoder_2_cls, text_encoder_2_id, text_encoder_2_subfolder = None, None, "" - text_encoder_3_cls, text_encoder_3_id, text_encoder_3_subfolder = None, None, "" - tokenizer_cls, tokenizer_id, tokenizer_subfolder = None, None, "" - tokenizer_2_cls, tokenizer_2_id, tokenizer_2_subfolder = None, None, "" - tokenizer_3_cls, tokenizer_3_id, tokenizer_3_subfolder = None, None, "" - supports_text_encoder_loras = True - - unet_kwargs = None - transformer_cls = None - transformer_kwargs = None - vae_cls = AutoencoderKL - vae_kwargs = None - - text_encoder_target_modules = ["q_proj", "k_proj", "v_proj", "out_proj"] - denoiser_target_modules = ["to_q", "to_k", "to_v", "to_out.0"] - - cached_non_lora_output = None - - def get_base_pipe_output(self): - if self.cached_non_lora_output is None: - self.cached_non_lora_output = self._compute_baseline_output() - return self.cached_non_lora_output - - def get_dummy_components(self, scheduler_cls=None, use_dora=False, lora_alpha=None): - if self.unet_kwargs and self.transformer_kwargs: - raise ValueError("Both `unet_kwargs` and `transformer_kwargs` cannot be specified.") - if self.has_two_text_encoders and self.has_three_text_encoders: - raise ValueError("Both `has_two_text_encoders` and `has_three_text_encoders` cannot be True.") - - scheduler_cls = scheduler_cls if scheduler_cls is not None else self.scheduler_cls - rank = 4 - lora_alpha = rank if lora_alpha is None else lora_alpha - - torch.manual_seed(0) - if self.unet_kwargs is not None: - unet = UNet2DConditionModel(**self.unet_kwargs) - else: - transformer = self.transformer_cls(**self.transformer_kwargs) - - scheduler = scheduler_cls(**self.scheduler_kwargs) - - torch.manual_seed(0) - vae = self.vae_cls(**self.vae_kwargs) - - text_encoder = self.text_encoder_cls.from_pretrained( - self.text_encoder_id, subfolder=self.text_encoder_subfolder - ) - tokenizer = self.tokenizer_cls.from_pretrained(self.tokenizer_id, subfolder=self.tokenizer_subfolder) - - if self.text_encoder_2_cls is not None: - text_encoder_2 = self.text_encoder_2_cls.from_pretrained( - self.text_encoder_2_id, subfolder=self.text_encoder_2_subfolder - ) - tokenizer_2 = self.tokenizer_2_cls.from_pretrained( - self.tokenizer_2_id, subfolder=self.tokenizer_2_subfolder - ) - - if self.text_encoder_3_cls is not None: - text_encoder_3 = self.text_encoder_3_cls.from_pretrained( - self.text_encoder_3_id, subfolder=self.text_encoder_3_subfolder - ) - tokenizer_3 = self.tokenizer_3_cls.from_pretrained( - self.tokenizer_3_id, subfolder=self.tokenizer_3_subfolder - ) - - text_lora_config = LoraConfig( - r=rank, - lora_alpha=lora_alpha, - target_modules=self.text_encoder_target_modules, - init_lora_weights=False, - use_dora=use_dora, - ) - - denoiser_lora_config = LoraConfig( - r=rank, - lora_alpha=lora_alpha, - target_modules=self.denoiser_target_modules, - init_lora_weights=False, - use_dora=use_dora, - ) - - pipeline_components = { - "scheduler": scheduler, - "vae": vae, - "text_encoder": text_encoder, - "tokenizer": tokenizer, - } - # Denoiser - if self.unet_kwargs is not None: - pipeline_components.update({"unet": unet}) - elif self.transformer_kwargs is not None: - pipeline_components.update({"transformer": transformer}) - - # Remaining text encoders. - if self.text_encoder_2_cls is not None: - pipeline_components.update({"tokenizer_2": tokenizer_2, "text_encoder_2": text_encoder_2}) - if self.text_encoder_3_cls is not None: - pipeline_components.update({"tokenizer_3": tokenizer_3, "text_encoder_3": text_encoder_3}) - - # Remaining stuff - init_params = inspect.signature(self.pipeline_class.__init__).parameters - if "safety_checker" in init_params: - pipeline_components.update({"safety_checker": None}) - if "feature_extractor" in init_params: - pipeline_components.update({"feature_extractor": None}) - if "image_encoder" in init_params: - pipeline_components.update({"image_encoder": None}) - - return pipeline_components, text_lora_config, denoiser_lora_config - - @property - def output_shape(self): - raise NotImplementedError - - def get_dummy_inputs(self, with_generator=True): - batch_size = 1 - sequence_length = 10 - num_channels = 4 - sizes = (32, 32) - - generator = torch.manual_seed(0) - noise = floats_tensor((batch_size, num_channels) + sizes) - input_ids = torch.randint(1, sequence_length, size=(batch_size, sequence_length), generator=generator) - - pipeline_inputs = { - "prompt": "A painting of a squirrel eating a burger", - "num_inference_steps": 5, - "guidance_scale": 6.0, - "output_type": "np", - } - if with_generator: - pipeline_inputs.update({"generator": generator}) - - return noise, input_ids, pipeline_inputs - - def _compute_baseline_output(self): - components, _, _ = self.get_dummy_components(self.scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - - # Always ensure the inputs are without the `generator`. Make sure to pass the `generator` - # explicitly. - _, _, inputs = self.get_dummy_inputs(with_generator=False) - return pipe(**inputs, generator=torch.manual_seed(0))[0] - - def _get_lora_state_dicts(self, modules_to_save): - state_dicts = {} - for module_name, module in modules_to_save.items(): - if module is not None: - state_dicts[f"{module_name}_lora_layers"] = get_peft_model_state_dict(module) - return state_dicts - - def _get_lora_adapter_metadata(self, modules_to_save): - metadatas = {} - for module_name, module in modules_to_save.items(): - if module is not None: - metadatas[f"{module_name}_lora_adapter_metadata"] = module.peft_config["default"].to_dict() - return metadatas - - def _get_modules_to_save(self, pipe, has_denoiser=False): - modules_to_save = {} - lora_loadable_modules = self.pipeline_class._lora_loadable_modules - - if ( - "text_encoder" in lora_loadable_modules - and hasattr(pipe, "text_encoder") - and getattr(pipe.text_encoder, "peft_config", None) is not None - ): - modules_to_save["text_encoder"] = pipe.text_encoder - - if ( - "text_encoder_2" in lora_loadable_modules - and hasattr(pipe, "text_encoder_2") - and getattr(pipe.text_encoder_2, "peft_config", None) is not None - ): - modules_to_save["text_encoder_2"] = pipe.text_encoder_2 - - if has_denoiser: - if "unet" in lora_loadable_modules and hasattr(pipe, "unet"): - modules_to_save["unet"] = pipe.unet - - if "transformer" in lora_loadable_modules and hasattr(pipe, "transformer"): - modules_to_save["transformer"] = pipe.transformer - - return modules_to_save - - def add_adapters_to_pipeline(self, pipe, text_lora_config=None, denoiser_lora_config=None, adapter_name="default"): - if text_lora_config is not None: - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - pipe.text_encoder.add_adapter(text_lora_config, adapter_name=adapter_name) - self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder" - ) - - if denoiser_lora_config is not None: - denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet - denoiser.add_adapter(denoiser_lora_config, adapter_name=adapter_name) - self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.") - else: - denoiser = None - - if text_lora_config is not None and self.has_two_text_encoders or self.has_three_text_encoders: - if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: - pipe.text_encoder_2.add_adapter(text_lora_config, adapter_name=adapter_name) - self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder_2), "Lora not correctly set in text encoder 2" - ) - return pipe, denoiser - - def test_simple_inference(self): - """ - Tests a simple inference and makes sure it works as expected - """ - output_no_lora = self.get_base_pipe_output() - assert output_no_lora.shape == self.output_shape - - def test_simple_inference_with_text_lora(self): - """ - Tests a simple inference with lora attached on the text encoder - and makes sure it works as expected - """ - if not self.supports_text_encoder_loras: - pytest.skip("Skipping test as text encoder LoRAs are not currently supported.") - - components, text_lora_config, _ = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - output_no_lora = self.get_base_pipe_output() - pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config=None) - - output_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue( - not np.allclose(output_lora, output_no_lora, atol=1e-3, rtol=1e-3), "Lora should change the output" - ) - - @require_peft_version_greater("0.13.1") - @require_transformers_version_greater("4.45.2") - def test_low_cpu_mem_usage_with_loading(self): - """Tests if we can load LoRA state dict with low_cpu_mem_usage.""" - components, text_lora_config, denoiser_lora_config = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config) - - images_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - - with tempfile.TemporaryDirectory() as tmpdirname: - modules_to_save = self._get_modules_to_save(pipe, has_denoiser=True) - lora_state_dicts = self._get_lora_state_dicts(modules_to_save) - self.pipeline_class.save_lora_weights( - save_directory=tmpdirname, safe_serialization=False, **lora_state_dicts - ) - - self.assertTrue(os.path.isfile(os.path.join(tmpdirname, "pytorch_lora_weights.bin"))) - pipe.unload_lora_weights() - pipe.load_lora_weights(os.path.join(tmpdirname, "pytorch_lora_weights.bin"), low_cpu_mem_usage=False) - - for module_name, module in modules_to_save.items(): - self.assertTrue(check_if_lora_correctly_set(module), f"Lora not correctly set in {module_name}") - - images_lora_from_pretrained = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue( - np.allclose(images_lora, images_lora_from_pretrained, atol=1e-3, rtol=1e-3), - "Loading from saved checkpoints should give same results.", - ) - - # Now, check for `low_cpu_mem_usage.` - pipe.unload_lora_weights() - pipe.load_lora_weights(os.path.join(tmpdirname, "pytorch_lora_weights.bin"), low_cpu_mem_usage=True) - - for module_name, module in modules_to_save.items(): - self.assertTrue(check_if_lora_correctly_set(module), f"Lora not correctly set in {module_name}") - - images_lora_from_pretrained_low_cpu = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue( - np.allclose(images_lora_from_pretrained_low_cpu, images_lora_from_pretrained, atol=1e-3, rtol=1e-3), - "Loading from saved checkpoints with `low_cpu_mem_usage` should give same results.", - ) - - def test_simple_inference_with_text_lora_and_scale(self): - """ - Tests a simple inference with lora attached on the text encoder + scale argument - and makes sure it works as expected - """ - if not self.supports_text_encoder_loras: - pytest.skip("Skipping test as text encoder LoRAs are not currently supported.") - - attention_kwargs_name = determine_attention_kwargs_name(self.pipeline_class) - components, text_lora_config, _ = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - output_no_lora = self.get_base_pipe_output() - - pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config=None) - - output_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue( - not np.allclose(output_lora, output_no_lora, atol=1e-3, rtol=1e-3), "Lora should change the output" - ) - - attention_kwargs = {attention_kwargs_name: {"scale": 0.5}} - output_lora_scale = pipe(**inputs, generator=torch.manual_seed(0), **attention_kwargs)[0] - - self.assertTrue( - not np.allclose(output_lora, output_lora_scale, atol=1e-3, rtol=1e-3), - "Lora + scale should change the output", - ) - - attention_kwargs = {attention_kwargs_name: {"scale": 0.0}} - output_lora_0_scale = pipe(**inputs, generator=torch.manual_seed(0), **attention_kwargs)[0] - - self.assertTrue( - np.allclose(output_no_lora, output_lora_0_scale, atol=1e-3, rtol=1e-3), - "Lora + 0 scale should lead to same result as no LoRA", - ) - - def test_simple_inference_with_text_lora_fused(self): - """ - Tests a simple inference with lora attached into text encoder + fuses the lora weights into base model - and makes sure it works as expected - """ - if not self.supports_text_encoder_loras: - pytest.skip("Skipping test as text encoder LoRAs are not currently supported.") - - components, text_lora_config, _ = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - output_no_lora = self.get_base_pipe_output() - - pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config=None) - - pipe.fuse_lora() - # Fusing should still keep the LoRA layers - self.assertTrue(check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder") - - if self.has_two_text_encoders or self.has_three_text_encoders: - if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: - self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder_2), "Lora not correctly set in text encoder 2" - ) - - ouput_fused = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertFalse( - np.allclose(ouput_fused, output_no_lora, atol=1e-3, rtol=1e-3), "Fused lora should change the output" - ) - - def test_simple_inference_with_text_lora_unloaded(self): - """ - Tests a simple inference with lora attached to text encoder, then unloads the lora weights - and makes sure it works as expected - """ - components, text_lora_config, _ = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - output_no_lora = self.get_base_pipe_output() - - pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config=None) - - pipe.unload_lora_weights() - # unloading should remove the LoRA layers - self.assertFalse(check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly unloaded in text encoder") - - if self.has_two_text_encoders or self.has_three_text_encoders: - if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: - self.assertFalse( - check_if_lora_correctly_set(pipe.text_encoder_2), - "Lora not correctly unloaded in text encoder 2", - ) - - ouput_unloaded = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue( - np.allclose(ouput_unloaded, output_no_lora, atol=1e-3, rtol=1e-3), - "Fused lora should change the output", - ) - - def test_simple_inference_with_text_lora_save_load(self): - """ - Tests a simple usecase where users could use saving utilities for LoRA. - """ - if not self.supports_text_encoder_loras: - pytest.skip("Skipping test as text encoder LoRAs are not currently supported.") - - components, text_lora_config, _ = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config=None) - - images_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - - with tempfile.TemporaryDirectory() as tmpdirname: - modules_to_save = self._get_modules_to_save(pipe) - lora_state_dicts = self._get_lora_state_dicts(modules_to_save) - - self.pipeline_class.save_lora_weights( - save_directory=tmpdirname, safe_serialization=False, **lora_state_dicts - ) - - self.assertTrue(os.path.isfile(os.path.join(tmpdirname, "pytorch_lora_weights.bin"))) - pipe.unload_lora_weights() - pipe.load_lora_weights(os.path.join(tmpdirname, "pytorch_lora_weights.bin")) - - for module_name, module in modules_to_save.items(): - self.assertTrue(check_if_lora_correctly_set(module), f"Lora not correctly set in {module_name}") - - images_lora_from_pretrained = pipe(**inputs, generator=torch.manual_seed(0))[0] - - self.assertTrue( - np.allclose(images_lora, images_lora_from_pretrained, atol=1e-3, rtol=1e-3), - "Loading from saved checkpoints should give same results.", - ) - - def test_simple_inference_with_partial_text_lora(self): - """ - Tests a simple inference with lora attached on the text encoder - with different ranks and some adapters removed - and makes sure it works as expected - """ - if not self.supports_text_encoder_loras: - pytest.skip("Skipping test as text encoder LoRAs are not currently supported.") - - components, _, _ = self.get_dummy_components() - # Verify `StableDiffusionLoraLoaderMixin.load_lora_into_text_encoder` handles different ranks per module (PR#8324). - text_lora_config = LoraConfig( - r=4, - rank_pattern={self.text_encoder_target_modules[i]: i + 1 for i in range(3)}, - lora_alpha=4, - target_modules=self.text_encoder_target_modules, - init_lora_weights=False, - use_dora=False, - ) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - output_no_lora = self.get_base_pipe_output() - - pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config=None) - - state_dict = {} - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - # Gather the state dict for the PEFT model, excluding `layers.4`, to ensure `load_lora_into_text_encoder` - # supports missing layers (PR#8324). - state_dict = { - f"text_encoder.{module_name}": param - for module_name, param in get_peft_model_state_dict(pipe.text_encoder).items() - if "encoder.layers.4" not in module_name - } - - if self.has_two_text_encoders or self.has_three_text_encoders: - if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: - state_dict.update( - { - f"text_encoder_2.{module_name}": param - for module_name, param in get_peft_model_state_dict(pipe.text_encoder_2).items() - if "encoder.layers.4" not in module_name - } - ) - - output_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue( - not np.allclose(output_lora, output_no_lora, atol=1e-3, rtol=1e-3), "Lora should change the output" - ) - - # Unload lora and load it back using the pipe.load_lora_weights machinery - pipe.unload_lora_weights() - pipe.load_lora_weights(state_dict) - - output_partial_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue( - not np.allclose(output_partial_lora, output_lora, atol=1e-3, rtol=1e-3), - "Removing adapters should change the output", - ) - - def test_simple_inference_save_pretrained_with_text_lora(self): - """Tests a simple usecase where users could use saving utilities for LoRA through save_pretrained.""" - if not self.supports_text_encoder_loras: - pytest.skip("Skipping test as text encoder LoRAs are not currently supported.") - - components, text_lora_config, _ = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config=None) - images_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - - with tempfile.TemporaryDirectory() as tmpdirname: - pipe.save_pretrained(tmpdirname) - - pipe_from_pretrained = self.pipeline_class.from_pretrained(tmpdirname) - pipe_from_pretrained.to(torch_device) - - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - self.assertTrue( - check_if_lora_correctly_set(pipe_from_pretrained.text_encoder), - "Lora not correctly set in text encoder", - ) - - if self.has_two_text_encoders or self.has_three_text_encoders: - if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: - self.assertTrue( - check_if_lora_correctly_set(pipe_from_pretrained.text_encoder_2), - "Lora not correctly set in text encoder 2", - ) - - images_lora_save_pretrained = pipe_from_pretrained(**inputs, generator=torch.manual_seed(0))[0] - - self.assertTrue( - np.allclose(images_lora, images_lora_save_pretrained, atol=1e-3, rtol=1e-3), - "Loading from saved checkpoints should give same results.", - ) - - def test_simple_inference_with_text_denoiser_lora_save_load(self): - """ - Tests a simple usecase where users could use saving utilities for LoRA for Unet + text encoder - """ - components, text_lora_config, denoiser_lora_config = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config) - - images_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - - with tempfile.TemporaryDirectory() as tmpdirname: - modules_to_save = self._get_modules_to_save(pipe, has_denoiser=True) - lora_state_dicts = self._get_lora_state_dicts(modules_to_save) - self.pipeline_class.save_lora_weights( - save_directory=tmpdirname, safe_serialization=False, **lora_state_dicts - ) - - self.assertTrue(os.path.isfile(os.path.join(tmpdirname, "pytorch_lora_weights.bin"))) - pipe.unload_lora_weights() - pipe.load_lora_weights(os.path.join(tmpdirname, "pytorch_lora_weights.bin")) - - for module_name, module in modules_to_save.items(): - self.assertTrue(check_if_lora_correctly_set(module), f"Lora not correctly set in {module_name}") - - images_lora_from_pretrained = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue( - np.allclose(images_lora, images_lora_from_pretrained, atol=1e-3, rtol=1e-3), - "Loading from saved checkpoints should give same results.", - ) - - def test_simple_inference_with_text_denoiser_lora_and_scale(self): - """ - Tests a simple inference with lora attached on the text encoder + Unet + scale argument - and makes sure it works as expected - """ - attention_kwargs_name = determine_attention_kwargs_name(self.pipeline_class) - components, text_lora_config, denoiser_lora_config = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - output_no_lora = self.get_base_pipe_output() - pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config) - - output_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue( - not np.allclose(output_lora, output_no_lora, atol=1e-3, rtol=1e-3), "Lora should change the output" - ) - - attention_kwargs = {attention_kwargs_name: {"scale": 0.5}} - output_lora_scale = pipe(**inputs, generator=torch.manual_seed(0), **attention_kwargs)[0] - - self.assertTrue( - not np.allclose(output_lora, output_lora_scale, atol=1e-3, rtol=1e-3), - "Lora + scale should change the output", - ) - - attention_kwargs = {attention_kwargs_name: {"scale": 0.0}} - output_lora_0_scale = pipe(**inputs, generator=torch.manual_seed(0), **attention_kwargs)[0] - - self.assertTrue( - np.allclose(output_no_lora, output_lora_0_scale, atol=1e-3, rtol=1e-3), - "Lora + 0 scale should lead to same result as no LoRA", - ) - - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - text_encoder_root = getattr(pipe.text_encoder, "text_model", pipe.text_encoder) - self.assertTrue( - text_encoder_root.encoder.layers[0].self_attn.q_proj.scaling["default"] == 1.0, - "The scaling parameter has not been correctly restored!", - ) - - def test_simple_inference_with_text_lora_denoiser_fused(self): - """ - Tests a simple inference with lora attached into text encoder + fuses the lora weights into base model - and makes sure it works as expected - with unet - """ - components, text_lora_config, denoiser_lora_config = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - output_no_lora = self.get_base_pipe_output() - - pipe, denoiser = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config) - - pipe.fuse_lora(components=self.pipeline_class._lora_loadable_modules) - - # Fusing should still keep the LoRA layers - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - self.assertTrue(check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder") - - self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser") - - if self.has_two_text_encoders or self.has_three_text_encoders: - if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: - self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder_2), "Lora not correctly set in text encoder 2" - ) - - output_fused = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertFalse( - np.allclose(output_fused, output_no_lora, atol=1e-3, rtol=1e-3), "Fused lora should change the output" - ) - - def test_simple_inference_with_text_denoiser_lora_unloaded(self): - """ - Tests a simple inference with lora attached to text encoder and unet, then unloads the lora weights - and makes sure it works as expected - """ - components, text_lora_config, denoiser_lora_config = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - output_no_lora = self.get_base_pipe_output() - - pipe, denoiser = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config) - - pipe.unload_lora_weights() - # unloading should remove the LoRA layers - self.assertFalse(check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly unloaded in text encoder") - self.assertFalse(check_if_lora_correctly_set(denoiser), "Lora not correctly unloaded in denoiser") - - if self.has_two_text_encoders or self.has_three_text_encoders: - if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: - self.assertFalse( - check_if_lora_correctly_set(pipe.text_encoder_2), - "Lora not correctly unloaded in text encoder 2", - ) - - output_unloaded = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue( - np.allclose(output_unloaded, output_no_lora, atol=1e-3, rtol=1e-3), - "Fused lora should change the output", - ) - - def test_simple_inference_with_text_denoiser_lora_unfused( - self, expected_atol: float = 1e-3, expected_rtol: float = 1e-3 - ): - """ - Tests a simple inference with lora attached to text encoder and unet, then unloads the lora weights - and makes sure it works as expected - """ - components, text_lora_config, denoiser_lora_config = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - pipe, denoiser = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config) - - pipe.fuse_lora(components=self.pipeline_class._lora_loadable_modules) - self.assertTrue(pipe.num_fused_loras == 1, f"{pipe.num_fused_loras=}, {pipe.fused_loras=}") - output_fused_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - - pipe.unfuse_lora(components=self.pipeline_class._lora_loadable_modules) - self.assertTrue(pipe.num_fused_loras == 0, f"{pipe.num_fused_loras=}, {pipe.fused_loras=}") - output_unfused_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - - # unloading should remove the LoRA layers - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - self.assertTrue(check_if_lora_correctly_set(pipe.text_encoder), "Unfuse should still keep LoRA layers") - - self.assertTrue(check_if_lora_correctly_set(denoiser), "Unfuse should still keep LoRA layers") - - if self.has_two_text_encoders or self.has_three_text_encoders: - if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: - self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder_2), "Unfuse should still keep LoRA layers" - ) - - # Fuse and unfuse should lead to the same results - self.assertTrue( - np.allclose(output_fused_lora, output_unfused_lora, atol=expected_atol, rtol=expected_rtol), - "Fused lora should not change the output", - ) - - def test_simple_inference_with_text_denoiser_multi_adapter(self): - """ - Tests a simple inference with lora attached to text encoder and unet, attaches - multiple adapters and set them - """ - components, text_lora_config, denoiser_lora_config = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - output_no_lora = self.get_base_pipe_output() - - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") - pipe.text_encoder.add_adapter(text_lora_config, "adapter-2") - self.assertTrue(check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder") - - denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet - denoiser.add_adapter(denoiser_lora_config, "adapter-1") - denoiser.add_adapter(denoiser_lora_config, "adapter-2") - self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.") - - if self.has_two_text_encoders or self.has_three_text_encoders: - if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: - pipe.text_encoder_2.add_adapter(text_lora_config, "adapter-1") - pipe.text_encoder_2.add_adapter(text_lora_config, "adapter-2") - self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder_2), "Lora not correctly set in text encoder 2" - ) - - pipe.set_adapters("adapter-1") - output_adapter_1 = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertFalse( - np.allclose(output_no_lora, output_adapter_1, atol=1e-3, rtol=1e-3), - "Adapter outputs should be different.", - ) - - pipe.set_adapters("adapter-2") - output_adapter_2 = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertFalse( - np.allclose(output_no_lora, output_adapter_2, atol=1e-3, rtol=1e-3), - "Adapter outputs should be different.", - ) - - pipe.set_adapters(["adapter-1", "adapter-2"]) - output_adapter_mixed = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertFalse( - np.allclose(output_no_lora, output_adapter_mixed, atol=1e-3, rtol=1e-3), - "Adapter outputs should be different.", - ) - - # Fuse and unfuse should lead to the same results - self.assertFalse( - np.allclose(output_adapter_1, output_adapter_2, atol=1e-3, rtol=1e-3), - "Adapter 1 and 2 should give different results", - ) - - self.assertFalse( - np.allclose(output_adapter_1, output_adapter_mixed, atol=1e-3, rtol=1e-3), - "Adapter 1 and mixed adapters should give different results", - ) - - self.assertFalse( - np.allclose(output_adapter_2, output_adapter_mixed, atol=1e-3, rtol=1e-3), - "Adapter 2 and mixed adapters should give different results", - ) - - pipe.disable_lora() - output_disabled = pipe(**inputs, generator=torch.manual_seed(0))[0] - - self.assertTrue( - np.allclose(output_no_lora, output_disabled, atol=1e-3, rtol=1e-3), - "output with no lora and output with lora disabled should give same results", - ) - - def test_wrong_adapter_name_raises_error(self): - adapter_name = "adapter-1" - - components, text_lora_config, denoiser_lora_config = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - pipe, _ = self.add_adapters_to_pipeline( - pipe, text_lora_config, denoiser_lora_config, adapter_name=adapter_name - ) - - with self.assertRaises(ValueError) as err_context: - pipe.set_adapters("test") - - self.assertTrue("not in the list of present adapters" in str(err_context.exception)) - - # test this works. - pipe.set_adapters(adapter_name) - _ = pipe(**inputs, generator=torch.manual_seed(0))[0] - - def test_multiple_wrong_adapter_name_raises_error(self): - adapter_name = "adapter-1" - components, text_lora_config, denoiser_lora_config = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - pipe, _ = self.add_adapters_to_pipeline( - pipe, text_lora_config, denoiser_lora_config, adapter_name=adapter_name - ) - - scale_with_wrong_components = {"foo": 0.0, "bar": 0.0, "tik": 0.0} - logger = logging.get_logger("diffusers.loaders.lora_base") - logger.setLevel(30) - with CaptureLogger(logger) as cap_logger: - pipe.set_adapters(adapter_name, adapter_weights=scale_with_wrong_components) - - wrong_components = sorted(set(scale_with_wrong_components.keys())) - msg = f"The following components in `adapter_weights` are not part of the pipeline: {wrong_components}. " - self.assertTrue(msg in str(cap_logger.out)) - - # test this works. - pipe.set_adapters(adapter_name) - _ = pipe(**inputs, generator=torch.manual_seed(0))[0] - - def test_simple_inference_with_text_denoiser_block_scale(self): - """ - Tests a simple inference with lora attached to text encoder and unet, attaches - one adapter and set different weights for different blocks (i.e. block lora) - """ - components, text_lora_config, denoiser_lora_config = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - output_no_lora = self.get_base_pipe_output() - - pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") - self.assertTrue(check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder") - - denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet - denoiser.add_adapter(denoiser_lora_config) - self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.") - - if self.has_two_text_encoders or self.has_three_text_encoders: - if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: - pipe.text_encoder_2.add_adapter(text_lora_config, "adapter-1") - self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder_2), "Lora not correctly set in text encoder 2" - ) - - weights_1 = {"text_encoder": 2, "unet": {"down": 5}} - pipe.set_adapters("adapter-1", weights_1) - output_weights_1 = pipe(**inputs, generator=torch.manual_seed(0))[0] - - weights_2 = {"unet": {"up": 5}} - pipe.set_adapters("adapter-1", weights_2) - output_weights_2 = pipe(**inputs, generator=torch.manual_seed(0))[0] - - self.assertFalse( - np.allclose(output_weights_1, output_weights_2, atol=1e-3, rtol=1e-3), - "LoRA weights 1 and 2 should give different results", - ) - self.assertFalse( - np.allclose(output_no_lora, output_weights_1, atol=1e-3, rtol=1e-3), - "No adapter and LoRA weights 1 should give different results", - ) - self.assertFalse( - np.allclose(output_no_lora, output_weights_2, atol=1e-3, rtol=1e-3), - "No adapter and LoRA weights 2 should give different results", - ) - - pipe.disable_lora() - output_disabled = pipe(**inputs, generator=torch.manual_seed(0))[0] - - self.assertTrue( - np.allclose(output_no_lora, output_disabled, atol=1e-3, rtol=1e-3), - "output with no lora and output with lora disabled should give same results", - ) - - def test_simple_inference_with_text_denoiser_multi_adapter_block_lora(self): - """ - Tests a simple inference with lora attached to text encoder and unet, attaches - multiple adapters and set different weights for different blocks (i.e. block lora) - """ - components, text_lora_config, denoiser_lora_config = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - output_no_lora = self.get_base_pipe_output() - - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") - pipe.text_encoder.add_adapter(text_lora_config, "adapter-2") - self.assertTrue(check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder") - - denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet - denoiser.add_adapter(denoiser_lora_config, "adapter-1") - denoiser.add_adapter(denoiser_lora_config, "adapter-2") - self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.") - - if self.has_two_text_encoders or self.has_three_text_encoders: - if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: - pipe.text_encoder_2.add_adapter(text_lora_config, "adapter-1") - pipe.text_encoder_2.add_adapter(text_lora_config, "adapter-2") - self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder_2), "Lora not correctly set in text encoder 2" - ) - - scales_1 = {"text_encoder": 2, "unet": {"down": 5}} - scales_2 = {"unet": {"down": 5, "mid": 5}} - - pipe.set_adapters("adapter-1", scales_1) - output_adapter_1 = pipe(**inputs, generator=torch.manual_seed(0))[0] - - pipe.set_adapters("adapter-2", scales_2) - output_adapter_2 = pipe(**inputs, generator=torch.manual_seed(0))[0] - - pipe.set_adapters(["adapter-1", "adapter-2"], [scales_1, scales_2]) - output_adapter_mixed = pipe(**inputs, generator=torch.manual_seed(0))[0] - - # Fuse and unfuse should lead to the same results - self.assertFalse( - np.allclose(output_adapter_1, output_adapter_2, atol=1e-3, rtol=1e-3), - "Adapter 1 and 2 should give different results", - ) - - self.assertFalse( - np.allclose(output_adapter_1, output_adapter_mixed, atol=1e-3, rtol=1e-3), - "Adapter 1 and mixed adapters should give different results", - ) - - self.assertFalse( - np.allclose(output_adapter_2, output_adapter_mixed, atol=1e-3, rtol=1e-3), - "Adapter 2 and mixed adapters should give different results", - ) - - pipe.disable_lora() - output_disabled = pipe(**inputs, generator=torch.manual_seed(0))[0] - - self.assertTrue( - np.allclose(output_no_lora, output_disabled, atol=1e-3, rtol=1e-3), - "output with no lora and output with lora disabled should give same results", - ) - - # a mismatching number of adapter_names and adapter_weights should raise an error - with self.assertRaises(ValueError): - pipe.set_adapters(["adapter-1", "adapter-2"], [scales_1]) - - def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(self): - """Tests that any valid combination of lora block scales can be used in pipe.set_adapter""" - - def updown_options(blocks_with_tf, layers_per_block, value): - """ - Generate every possible combination for how a lora weight dict for the up/down part can be. - E.g. 2, {"block_1": 2}, {"block_1": [2,2,2]}, {"block_1": 2, "block_2": [2,2,2]}, ... - """ - num_val = value - list_val = [value] * layers_per_block - - node_opts = [None, num_val, list_val] - node_opts_foreach_block = [node_opts] * len(blocks_with_tf) - - updown_opts = [num_val] - for nodes in product(*node_opts_foreach_block): - if all(n is None for n in nodes): - continue - opt = {} - for b, n in zip(blocks_with_tf, nodes): - if n is not None: - opt["block_" + str(b)] = n - updown_opts.append(opt) - return updown_opts - - def all_possible_dict_opts(unet, value): - """ - Generate every possible combination for how a lora weight dict can be. - E.g. 2, {"unet: {"down": 2}}, {"unet: {"down": [2,2,2]}}, {"unet: {"mid": 2, "up": [2,2,2]}}, ... - """ - - down_blocks_with_tf = [i for i, d in enumerate(unet.down_blocks) if hasattr(d, "attentions")] - up_blocks_with_tf = [i for i, u in enumerate(unet.up_blocks) if hasattr(u, "attentions")] - - layers_per_block = unet.config.layers_per_block - - text_encoder_opts = [None, value] - text_encoder_2_opts = [None, value] - mid_opts = [None, value] - down_opts = [None] + updown_options(down_blocks_with_tf, layers_per_block, value) - up_opts = [None] + updown_options(up_blocks_with_tf, layers_per_block + 1, value) - - opts = [] - - for t1, t2, d, m, u in product(text_encoder_opts, text_encoder_2_opts, down_opts, mid_opts, up_opts): - if all(o is None for o in (t1, t2, d, m, u)): - continue - opt = {} - if t1 is not None: - opt["text_encoder"] = t1 - if t2 is not None: - opt["text_encoder_2"] = t2 - if all(o is None for o in (d, m, u)): - # no unet scaling - continue - opt["unet"] = {} - if d is not None: - opt["unet"]["down"] = d - if m is not None: - opt["unet"]["mid"] = m - if u is not None: - opt["unet"]["up"] = u - opts.append(opt) - - return opts - - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(self.scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") - - denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet - denoiser.add_adapter(denoiser_lora_config, "adapter-1") - - if self.has_two_text_encoders or self.has_three_text_encoders: - lora_loadable_components = self.pipeline_class._lora_loadable_modules - if "text_encoder_2" in lora_loadable_components: - pipe.text_encoder_2.add_adapter(text_lora_config, "adapter-1") - - for scale_dict in all_possible_dict_opts(pipe.unet, value=1234): - # test if lora block scales can be set with this scale_dict - if not self.has_two_text_encoders and "text_encoder_2" in scale_dict: - del scale_dict["text_encoder_2"] - - pipe.set_adapters("adapter-1", scale_dict) # test will fail if this line throws an error - - def test_simple_inference_with_text_denoiser_multi_adapter_delete_adapter(self): - """ - Tests a simple inference with lora attached to text encoder and unet, attaches - multiple adapters and set/delete them - """ - components, text_lora_config, denoiser_lora_config = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - output_no_lora = self.get_base_pipe_output() - - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") - pipe.text_encoder.add_adapter(text_lora_config, "adapter-2") - self.assertTrue(check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder") - - denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet - denoiser.add_adapter(denoiser_lora_config, "adapter-1") - denoiser.add_adapter(denoiser_lora_config, "adapter-2") - self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.") - - if self.has_two_text_encoders or self.has_three_text_encoders: - lora_loadable_components = self.pipeline_class._lora_loadable_modules - if "text_encoder_2" in lora_loadable_components: - pipe.text_encoder_2.add_adapter(text_lora_config, "adapter-1") - pipe.text_encoder_2.add_adapter(text_lora_config, "adapter-2") - self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder_2), "Lora not correctly set in text encoder 2" - ) - - pipe.set_adapters("adapter-1") - output_adapter_1 = pipe(**inputs, generator=torch.manual_seed(0))[0] - - pipe.set_adapters("adapter-2") - output_adapter_2 = pipe(**inputs, generator=torch.manual_seed(0))[0] - - pipe.set_adapters(["adapter-1", "adapter-2"]) - output_adapter_mixed = pipe(**inputs, generator=torch.manual_seed(0))[0] - - self.assertFalse( - np.allclose(output_adapter_1, output_adapter_2, atol=1e-3, rtol=1e-3), - "Adapter 1 and 2 should give different results", - ) - - self.assertFalse( - np.allclose(output_adapter_1, output_adapter_mixed, atol=1e-3, rtol=1e-3), - "Adapter 1 and mixed adapters should give different results", - ) - - self.assertFalse( - np.allclose(output_adapter_2, output_adapter_mixed, atol=1e-3, rtol=1e-3), - "Adapter 2 and mixed adapters should give different results", - ) - - pipe.delete_adapters("adapter-1") - output_deleted_adapter_1 = pipe(**inputs, generator=torch.manual_seed(0))[0] - - self.assertTrue( - np.allclose(output_deleted_adapter_1, output_adapter_2, atol=1e-3, rtol=1e-3), - "Adapter 1 and 2 should give different results", - ) - - pipe.delete_adapters("adapter-2") - output_deleted_adapters = pipe(**inputs, generator=torch.manual_seed(0))[0] - - self.assertTrue( - np.allclose(output_no_lora, output_deleted_adapters, atol=1e-3, rtol=1e-3), - "output with no lora and output with lora disabled should give same results", - ) - - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") - pipe.text_encoder.add_adapter(text_lora_config, "adapter-2") - - denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet - denoiser.add_adapter(denoiser_lora_config, "adapter-1") - denoiser.add_adapter(denoiser_lora_config, "adapter-2") - self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.") - - pipe.set_adapters(["adapter-1", "adapter-2"]) - pipe.delete_adapters(["adapter-1", "adapter-2"]) - - output_deleted_adapters = pipe(**inputs, generator=torch.manual_seed(0))[0] - - self.assertTrue( - np.allclose(output_no_lora, output_deleted_adapters, atol=1e-3, rtol=1e-3), - "output with no lora and output with lora disabled should give same results", - ) - - def test_simple_inference_with_text_denoiser_multi_adapter_weighted(self): - """ - Tests a simple inference with lora attached to text encoder and unet, attaches - multiple adapters and set them - """ - components, text_lora_config, denoiser_lora_config = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - output_no_lora = self.get_base_pipe_output() - - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") - pipe.text_encoder.add_adapter(text_lora_config, "adapter-2") - self.assertTrue(check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder") - - denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet - denoiser.add_adapter(denoiser_lora_config, "adapter-1") - denoiser.add_adapter(denoiser_lora_config, "adapter-2") - self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.") - - if self.has_two_text_encoders or self.has_three_text_encoders: - lora_loadable_components = self.pipeline_class._lora_loadable_modules - if "text_encoder_2" in lora_loadable_components: - pipe.text_encoder_2.add_adapter(text_lora_config, "adapter-1") - pipe.text_encoder_2.add_adapter(text_lora_config, "adapter-2") - self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder_2), "Lora not correctly set in text encoder 2" - ) - - pipe.set_adapters("adapter-1") - output_adapter_1 = pipe(**inputs, generator=torch.manual_seed(0))[0] - - pipe.set_adapters("adapter-2") - output_adapter_2 = pipe(**inputs, generator=torch.manual_seed(0))[0] - - pipe.set_adapters(["adapter-1", "adapter-2"]) - output_adapter_mixed = pipe(**inputs, generator=torch.manual_seed(0))[0] - - # Fuse and unfuse should lead to the same results - self.assertFalse( - np.allclose(output_adapter_1, output_adapter_2, atol=1e-3, rtol=1e-3), - "Adapter 1 and 2 should give different results", - ) - - self.assertFalse( - np.allclose(output_adapter_1, output_adapter_mixed, atol=1e-3, rtol=1e-3), - "Adapter 1 and mixed adapters should give different results", - ) - - self.assertFalse( - np.allclose(output_adapter_2, output_adapter_mixed, atol=1e-3, rtol=1e-3), - "Adapter 2 and mixed adapters should give different results", - ) - - pipe.set_adapters(["adapter-1", "adapter-2"], [0.5, 0.6]) - output_adapter_mixed_weighted = pipe(**inputs, generator=torch.manual_seed(0))[0] - - self.assertFalse( - np.allclose(output_adapter_mixed_weighted, output_adapter_mixed, atol=1e-3, rtol=1e-3), - "Weighted adapter and mixed adapter should give different results", - ) - - pipe.disable_lora() - output_disabled = pipe(**inputs, generator=torch.manual_seed(0))[0] - - self.assertTrue( - np.allclose(output_no_lora, output_disabled, atol=1e-3, rtol=1e-3), - "output with no lora and output with lora disabled should give same results", - ) - - def test_get_adapters(self): - """ - Tests a simple usecase where we attach multiple adapters and check if the results - are the expected results - """ - components, text_lora_config, denoiser_lora_config = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") - - denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet - denoiser.add_adapter(denoiser_lora_config, "adapter-1") - - adapter_names = pipe.get_active_adapters() - self.assertListEqual(adapter_names, ["adapter-1"]) - - pipe.text_encoder.add_adapter(text_lora_config, "adapter-2") - denoiser.add_adapter(denoiser_lora_config, "adapter-2") - - adapter_names = pipe.get_active_adapters() - self.assertListEqual(adapter_names, ["adapter-2"]) - - pipe.set_adapters(["adapter-1", "adapter-2"]) - self.assertListEqual(pipe.get_active_adapters(), ["adapter-1", "adapter-2"]) - - def test_get_list_adapters(self): - """ - Tests a simple usecase where we attach multiple adapters and check if the results - are the expected results - """ - components, text_lora_config, denoiser_lora_config = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - - # 1. - dicts_to_be_checked = {} - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") - dicts_to_be_checked = {"text_encoder": ["adapter-1"]} - - if self.unet_kwargs is not None: - pipe.unet.add_adapter(denoiser_lora_config, "adapter-1") - dicts_to_be_checked.update({"unet": ["adapter-1"]}) - else: - pipe.transformer.add_adapter(denoiser_lora_config, "adapter-1") - dicts_to_be_checked.update({"transformer": ["adapter-1"]}) - - self.assertDictEqual(pipe.get_list_adapters(), dicts_to_be_checked) - - # 2. - dicts_to_be_checked = {} - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - pipe.text_encoder.add_adapter(text_lora_config, "adapter-2") - dicts_to_be_checked = {"text_encoder": ["adapter-1", "adapter-2"]} - - if self.unet_kwargs is not None: - pipe.unet.add_adapter(denoiser_lora_config, "adapter-2") - dicts_to_be_checked.update({"unet": ["adapter-1", "adapter-2"]}) - else: - pipe.transformer.add_adapter(denoiser_lora_config, "adapter-2") - dicts_to_be_checked.update({"transformer": ["adapter-1", "adapter-2"]}) - - self.assertDictEqual(pipe.get_list_adapters(), dicts_to_be_checked) - - # 3. - pipe.set_adapters(["adapter-1", "adapter-2"]) - - dicts_to_be_checked = {} - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - dicts_to_be_checked = {"text_encoder": ["adapter-1", "adapter-2"]} - - if self.unet_kwargs is not None: - dicts_to_be_checked.update({"unet": ["adapter-1", "adapter-2"]}) - else: - dicts_to_be_checked.update({"transformer": ["adapter-1", "adapter-2"]}) - - self.assertDictEqual( - pipe.get_list_adapters(), - dicts_to_be_checked, - ) - - # 4. - dicts_to_be_checked = {} - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - dicts_to_be_checked = {"text_encoder": ["adapter-1", "adapter-2"]} - - if self.unet_kwargs is not None: - pipe.unet.add_adapter(denoiser_lora_config, "adapter-3") - dicts_to_be_checked.update({"unet": ["adapter-1", "adapter-2", "adapter-3"]}) - else: - pipe.transformer.add_adapter(denoiser_lora_config, "adapter-3") - dicts_to_be_checked.update({"transformer": ["adapter-1", "adapter-2", "adapter-3"]}) - - self.assertDictEqual(pipe.get_list_adapters(), dicts_to_be_checked) - - def test_simple_inference_with_text_lora_denoiser_fused_multi( - self, expected_atol: float = 1e-3, expected_rtol: float = 1e-3 - ): - """ - Tests a simple inference with lora attached into text encoder + fuses the lora weights into base model - and makes sure it works as expected - with unet and multi-adapter case - """ - components, text_lora_config, denoiser_lora_config = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") - self.assertTrue(check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder") - pipe.text_encoder.add_adapter(text_lora_config, "adapter-2") - - denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet - denoiser.add_adapter(denoiser_lora_config, "adapter-1") - self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.") - denoiser.add_adapter(denoiser_lora_config, "adapter-2") - - if self.has_two_text_encoders or self.has_three_text_encoders: - lora_loadable_components = self.pipeline_class._lora_loadable_modules - if "text_encoder_2" in lora_loadable_components: - pipe.text_encoder_2.add_adapter(text_lora_config, "adapter-1") - self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder_2), "Lora not correctly set in text encoder 2" - ) - pipe.text_encoder_2.add_adapter(text_lora_config, "adapter-2") - - # set them to multi-adapter inference mode - pipe.set_adapters(["adapter-1", "adapter-2"]) - outputs_all_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - - pipe.set_adapters(["adapter-1"]) - outputs_lora_1 = pipe(**inputs, generator=torch.manual_seed(0))[0] - - pipe.fuse_lora(components=self.pipeline_class._lora_loadable_modules, adapter_names=["adapter-1"]) - self.assertTrue(pipe.num_fused_loras == 1, f"{pipe.num_fused_loras=}, {pipe.fused_loras=}") - - # Fusing should still keep the LoRA layers so output should remain the same - outputs_lora_1_fused = pipe(**inputs, generator=torch.manual_seed(0))[0] - - self.assertTrue( - np.allclose(outputs_lora_1, outputs_lora_1_fused, atol=expected_atol, rtol=expected_rtol), - "Fused lora should not change the output", - ) - - pipe.unfuse_lora(components=self.pipeline_class._lora_loadable_modules) - self.assertTrue(pipe.num_fused_loras == 0, f"{pipe.num_fused_loras=}, {pipe.fused_loras=}") - - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - self.assertTrue(check_if_lora_correctly_set(pipe.text_encoder), "Unfuse should still keep LoRA layers") - - self.assertTrue(check_if_lora_correctly_set(denoiser), "Unfuse should still keep LoRA layers") - - if self.has_two_text_encoders or self.has_three_text_encoders: - if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: - self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder_2), "Unfuse should still keep LoRA layers" - ) - - pipe.fuse_lora(components=self.pipeline_class._lora_loadable_modules, adapter_names=["adapter-2", "adapter-1"]) - self.assertTrue(pipe.num_fused_loras == 2, f"{pipe.num_fused_loras=}, {pipe.fused_loras=}") - - # Fusing should still keep the LoRA layers - output_all_lora_fused = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue( - np.allclose(output_all_lora_fused, outputs_all_lora, atol=expected_atol, rtol=expected_rtol), - "Fused lora should not change the output", - ) - pipe.unfuse_lora(components=self.pipeline_class._lora_loadable_modules) - self.assertTrue(pipe.num_fused_loras == 0, f"{pipe.num_fused_loras=}, {pipe.fused_loras=}") - - def test_lora_scale_kwargs_match_fusion(self, expected_atol: float = 1e-3, expected_rtol: float = 1e-3): - attention_kwargs_name = determine_attention_kwargs_name(self.pipeline_class) - - for lora_scale in [1.0, 0.8]: - components, text_lora_config, denoiser_lora_config = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - output_no_lora = self.get_base_pipe_output() - - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") - self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder" - ) - - denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet - denoiser.add_adapter(denoiser_lora_config, "adapter-1") - self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.") - - if self.has_two_text_encoders or self.has_three_text_encoders: - lora_loadable_components = self.pipeline_class._lora_loadable_modules - if "text_encoder_2" in lora_loadable_components: - pipe.text_encoder_2.add_adapter(text_lora_config, "adapter-1") - self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder_2), - "Lora not correctly set in text encoder 2", - ) - - pipe.set_adapters(["adapter-1"]) - attention_kwargs = {attention_kwargs_name: {"scale": lora_scale}} - outputs_lora_1 = pipe(**inputs, generator=torch.manual_seed(0), **attention_kwargs)[0] - - pipe.fuse_lora( - components=self.pipeline_class._lora_loadable_modules, - adapter_names=["adapter-1"], - lora_scale=lora_scale, - ) - self.assertTrue(pipe.num_fused_loras == 1, f"{pipe.num_fused_loras=}, {pipe.fused_loras=}") - - outputs_lora_1_fused = pipe(**inputs, generator=torch.manual_seed(0))[0] - - self.assertTrue( - np.allclose(outputs_lora_1, outputs_lora_1_fused, atol=expected_atol, rtol=expected_rtol), - "Fused lora should not change the output", - ) - self.assertFalse( - np.allclose(output_no_lora, outputs_lora_1, atol=expected_atol, rtol=expected_rtol), - "LoRA should change the output", - ) - - @unittest.skip("This is failing for now - need to investigate") - def test_simple_inference_with_text_denoiser_lora_unfused_torch_compile(self): - """ - Tests a simple inference with lora attached to text encoder and unet, then unloads the lora weights - and makes sure it works as expected - """ - components, text_lora_config, denoiser_lora_config = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config) - - pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True) - pipe.text_encoder = torch.compile(pipe.text_encoder, mode="reduce-overhead", fullgraph=True) - - if self.has_two_text_encoders or self.has_three_text_encoders: - pipe.text_encoder_2 = torch.compile(pipe.text_encoder_2, mode="reduce-overhead", fullgraph=True) - - # Just makes sure it works. - _ = pipe(**inputs, generator=torch.manual_seed(0))[0] - - def test_logs_info_when_no_lora_keys_found(self): - # Skip text encoder check for now as that is handled with `transformers`. - components, _, _ = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - - _, _, inputs = self.get_dummy_inputs(with_generator=False) - output_no_lora = self.get_base_pipe_output() - - no_op_state_dict = {"lora_foo": torch.tensor(2.0), "lora_bar": torch.tensor(3.0)} - logger = logging.get_logger("diffusers.loaders.peft") - logger.setLevel(logging.WARNING) - - with CaptureLogger(logger) as cap_logger: - pipe.load_lora_weights(no_op_state_dict) - out_after_lora_attempt = pipe(**inputs, generator=torch.manual_seed(0))[0] - - denoiser = getattr(pipe, "unet") if self.unet_kwargs is not None else getattr(pipe, "transformer") - self.assertTrue(cap_logger.out.startswith(f"No LoRA keys associated to {denoiser.__class__.__name__}")) - self.assertTrue(np.allclose(output_no_lora, out_after_lora_attempt, atol=1e-5, rtol=1e-5)) - - # test only for text encoder - for lora_module in self.pipeline_class._lora_loadable_modules: - if "text_encoder" in lora_module: - text_encoder = getattr(pipe, lora_module) - if lora_module == "text_encoder": - prefix = "text_encoder" - elif lora_module == "text_encoder_2": - prefix = "text_encoder_2" - - logger = logging.get_logger("diffusers.loaders.lora_base") - logger.setLevel(logging.WARNING) - - with CaptureLogger(logger) as cap_logger: - self.pipeline_class.load_lora_into_text_encoder( - no_op_state_dict, network_alphas=None, text_encoder=text_encoder, prefix=prefix - ) - - self.assertTrue( - cap_logger.out.startswith(f"No LoRA keys associated to {text_encoder.__class__.__name__}") - ) - - def test_set_adapters_match_attention_kwargs(self): - """Test to check if outputs after `set_adapters()` and attention kwargs match.""" - attention_kwargs_name = determine_attention_kwargs_name(self.pipeline_class) - components, text_lora_config, denoiser_lora_config = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - output_no_lora = self.get_base_pipe_output() - pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config) - - lora_scale = 0.5 - attention_kwargs = {attention_kwargs_name: {"scale": lora_scale}} - output_lora_scale = pipe(**inputs, generator=torch.manual_seed(0), **attention_kwargs)[0] - self.assertFalse( - np.allclose(output_no_lora, output_lora_scale, atol=1e-3, rtol=1e-3), - "Lora + scale should change the output", - ) - - pipe.set_adapters("default", lora_scale) - output_lora_scale_wo_kwargs = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue( - not np.allclose(output_no_lora, output_lora_scale_wo_kwargs, atol=1e-3, rtol=1e-3), - "Lora + scale should change the output", - ) - self.assertTrue( - np.allclose(output_lora_scale, output_lora_scale_wo_kwargs, atol=1e-3, rtol=1e-3), - "Lora + scale should match the output of `set_adapters()`.", - ) - - with tempfile.TemporaryDirectory() as tmpdirname: - modules_to_save = self._get_modules_to_save(pipe, has_denoiser=True) - lora_state_dicts = self._get_lora_state_dicts(modules_to_save) - self.pipeline_class.save_lora_weights( - save_directory=tmpdirname, safe_serialization=True, **lora_state_dicts - ) - - self.assertTrue(os.path.isfile(os.path.join(tmpdirname, "pytorch_lora_weights.safetensors"))) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - pipe.load_lora_weights(os.path.join(tmpdirname, "pytorch_lora_weights.safetensors")) - - for module_name, module in modules_to_save.items(): - self.assertTrue(check_if_lora_correctly_set(module), f"Lora not correctly set in {module_name}") - - output_lora_from_pretrained = pipe(**inputs, generator=torch.manual_seed(0), **attention_kwargs)[0] - self.assertTrue( - not np.allclose(output_no_lora, output_lora_from_pretrained, atol=1e-3, rtol=1e-3), - "Lora + scale should change the output", - ) - self.assertTrue( - np.allclose(output_lora_scale, output_lora_from_pretrained, atol=1e-3, rtol=1e-3), - "Loading from saved checkpoints should give same results as attention_kwargs.", - ) - self.assertTrue( - np.allclose(output_lora_scale_wo_kwargs, output_lora_from_pretrained, atol=1e-3, rtol=1e-3), - "Loading from saved checkpoints should give same results as set_adapters().", - ) - - @parameterized.expand([4, 8, 16]) - def test_lora_adapter_metadata_is_loaded_correctly(self, lora_alpha): - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(lora_alpha=lora_alpha) - pipe = self.pipeline_class(**components) - - pipe, _ = self.add_adapters_to_pipeline( - pipe, text_lora_config=text_lora_config, denoiser_lora_config=denoiser_lora_config - ) - - with tempfile.TemporaryDirectory() as tmpdir: - modules_to_save = self._get_modules_to_save(pipe, has_denoiser=True) - lora_state_dicts = self._get_lora_state_dicts(modules_to_save) - lora_metadatas = self._get_lora_adapter_metadata(modules_to_save) - self.pipeline_class.save_lora_weights(save_directory=tmpdir, **lora_state_dicts, **lora_metadatas) - pipe.unload_lora_weights() - - out = pipe.lora_state_dict(tmpdir, return_lora_metadata=True) - if len(out) == 3: - _, _, parsed_metadata = out - elif len(out) == 2: - _, parsed_metadata = out - - denoiser_key = ( - f"{self.pipeline_class.transformer_name}" - if self.transformer_kwargs is not None - else f"{self.pipeline_class.unet_name}" - ) - self.assertTrue(any(k.startswith(f"{denoiser_key}.") for k in parsed_metadata)) - check_module_lora_metadata( - parsed_metadata=parsed_metadata, lora_metadatas=lora_metadatas, module_key=denoiser_key - ) - - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - text_encoder_key = self.pipeline_class.text_encoder_name - self.assertTrue(any(k.startswith(f"{text_encoder_key}.") for k in parsed_metadata)) - check_module_lora_metadata( - parsed_metadata=parsed_metadata, lora_metadatas=lora_metadatas, module_key=text_encoder_key - ) - - if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: - text_encoder_2_key = "text_encoder_2" - self.assertTrue(any(k.startswith(f"{text_encoder_2_key}.") for k in parsed_metadata)) - check_module_lora_metadata( - parsed_metadata=parsed_metadata, lora_metadatas=lora_metadatas, module_key=text_encoder_2_key - ) - - @parameterized.expand([4, 8, 16]) - def test_lora_adapter_metadata_save_load_inference(self, lora_alpha): - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(lora_alpha=lora_alpha) - pipe = self.pipeline_class(**components).to(torch_device) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - pipe, _ = self.add_adapters_to_pipeline( - pipe, text_lora_config=text_lora_config, denoiser_lora_config=denoiser_lora_config - ) - output_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - - with tempfile.TemporaryDirectory() as tmpdir: - modules_to_save = self._get_modules_to_save(pipe, has_denoiser=True) - lora_state_dicts = self._get_lora_state_dicts(modules_to_save) - lora_metadatas = self._get_lora_adapter_metadata(modules_to_save) - self.pipeline_class.save_lora_weights(save_directory=tmpdir, **lora_state_dicts, **lora_metadatas) - pipe.unload_lora_weights() - pipe.load_lora_weights(tmpdir) - - output_lora_pretrained = pipe(**inputs, generator=torch.manual_seed(0))[0] - - self.assertTrue( - np.allclose(output_lora, output_lora_pretrained, atol=1e-3, rtol=1e-3), "Lora outputs should match." - ) - - def test_lora_unload_add_adapter(self): - """Tests if `unload_lora_weights()` -> `add_adapter()` works.""" - components, text_lora_config, denoiser_lora_config = self.get_dummy_components() - pipe = self.pipeline_class(**components).to(torch_device) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - pipe, _ = self.add_adapters_to_pipeline( - pipe, text_lora_config=text_lora_config, denoiser_lora_config=denoiser_lora_config - ) - _ = pipe(**inputs, generator=torch.manual_seed(0))[0] - - # unload and then add. - pipe.unload_lora_weights() - pipe, _ = self.add_adapters_to_pipeline( - pipe, text_lora_config=text_lora_config, denoiser_lora_config=denoiser_lora_config - ) - _ = pipe(**inputs, generator=torch.manual_seed(0))[0] - - def test_inference_load_delete_load_adapters(self): - "Tests if `load_lora_weights()` -> `delete_adapters()` -> `load_lora_weights()` works." - components, text_lora_config, denoiser_lora_config = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - output_no_lora = self.get_base_pipe_output() - - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - pipe.text_encoder.add_adapter(text_lora_config) - self.assertTrue(check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder") - - denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet - denoiser.add_adapter(denoiser_lora_config) - self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.") - - if self.has_two_text_encoders or self.has_three_text_encoders: - lora_loadable_components = self.pipeline_class._lora_loadable_modules - if "text_encoder_2" in lora_loadable_components: - pipe.text_encoder_2.add_adapter(text_lora_config) - self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder_2), "Lora not correctly set in text encoder 2" - ) - - output_adapter_1 = pipe(**inputs, generator=torch.manual_seed(0))[0] - - with tempfile.TemporaryDirectory() as tmpdirname: - modules_to_save = self._get_modules_to_save(pipe, has_denoiser=True) - lora_state_dicts = self._get_lora_state_dicts(modules_to_save) - self.pipeline_class.save_lora_weights(save_directory=tmpdirname, **lora_state_dicts) - self.assertTrue(os.path.isfile(os.path.join(tmpdirname, "pytorch_lora_weights.safetensors"))) - - # First, delete adapter and compare. - pipe.delete_adapters(pipe.get_active_adapters()[0]) - output_no_adapter = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertFalse(np.allclose(output_adapter_1, output_no_adapter, atol=1e-3, rtol=1e-3)) - self.assertTrue(np.allclose(output_no_lora, output_no_adapter, atol=1e-3, rtol=1e-3)) - - # Then load adapter and compare. - pipe.load_lora_weights(tmpdirname) - - output_lora_loaded = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(np.allclose(output_adapter_1, output_lora_loaded, atol=1e-3, rtol=1e-3)) - - def _test_group_offloading_inference_denoiser(self, offload_type, use_stream): - from diffusers.hooks.group_offloading import _get_top_level_group_offload_hook - - onload_device = torch_device - offload_device = torch.device("cpu") - - components, text_lora_config, denoiser_lora_config = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - - denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet - denoiser.add_adapter(denoiser_lora_config) - self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.") - - with tempfile.TemporaryDirectory() as tmpdirname: - modules_to_save = self._get_modules_to_save(pipe, has_denoiser=True) - lora_state_dicts = self._get_lora_state_dicts(modules_to_save) - self.pipeline_class.save_lora_weights( - save_directory=tmpdirname, safe_serialization=True, **lora_state_dicts - ) - self.assertTrue(os.path.isfile(os.path.join(tmpdirname, "pytorch_lora_weights.safetensors"))) - - components, _, _ = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe.set_progress_bar_config(disable=None) - denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet - - pipe.load_lora_weights(os.path.join(tmpdirname, "pytorch_lora_weights.safetensors")) - check_if_lora_correctly_set(denoiser) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - # Test group offloading with load_lora_weights - denoiser.enable_group_offload( - onload_device=onload_device, - offload_device=offload_device, - offload_type=offload_type, - num_blocks_per_group=1, - use_stream=use_stream, - ) - # Place other model-level components on `torch_device`. - for _, component in pipe.components.items(): - if isinstance(component, torch.nn.Module): - component.to(torch_device) - group_offload_hook_1 = _get_top_level_group_offload_hook(denoiser) - self.assertTrue(group_offload_hook_1 is not None) - output_1 = pipe(**inputs, generator=torch.manual_seed(0))[0] - - # Test group offloading after removing the lora - pipe.unload_lora_weights() - group_offload_hook_2 = _get_top_level_group_offload_hook(denoiser) - self.assertTrue(group_offload_hook_2 is not None) - output_2 = pipe(**inputs, generator=torch.manual_seed(0))[0] # noqa: F841 - - # Add the lora again and check if group offloading works - pipe.load_lora_weights(os.path.join(tmpdirname, "pytorch_lora_weights.safetensors")) - check_if_lora_correctly_set(denoiser) - group_offload_hook_3 = _get_top_level_group_offload_hook(denoiser) - self.assertTrue(group_offload_hook_3 is not None) - output_3 = pipe(**inputs, generator=torch.manual_seed(0))[0] - - self.assertTrue(np.allclose(output_1, output_3, atol=1e-3, rtol=1e-3)) - - @parameterized.expand([("block_level", True), ("leaf_level", False), ("leaf_level", True)]) - @require_torch_accelerator - def test_group_offloading_inference_denoiser(self, offload_type, use_stream): - for cls in inspect.getmro(self.__class__): - if "test_group_offloading_inference_denoiser" in cls.__dict__ and cls is not PeftLoraLoaderMixinTests: - # Skip this test if it is overwritten by child class. We need to do this because parameterized - # materializes the test methods on invocation which cannot be overridden. - return - self._test_group_offloading_inference_denoiser(offload_type, use_stream) - - @require_torch_accelerator - def test_lora_loading_model_cpu_offload(self): - components, _, denoiser_lora_config = self.get_dummy_components() - _, _, inputs = self.get_dummy_inputs(with_generator=False) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - - denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet - denoiser.add_adapter(denoiser_lora_config) - self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.") - - output_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - - with tempfile.TemporaryDirectory() as tmpdirname: - modules_to_save = self._get_modules_to_save(pipe, has_denoiser=True) - lora_state_dicts = self._get_lora_state_dicts(modules_to_save) - self.pipeline_class.save_lora_weights( - save_directory=tmpdirname, safe_serialization=True, **lora_state_dicts - ) - # reinitialize the pipeline to mimic the inference workflow. - components, _, denoiser_lora_config = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe.enable_model_cpu_offload(device=torch_device) - pipe.load_lora_weights(tmpdirname) - denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet - self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.") - - output_lora_loaded = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(np.allclose(output_lora, output_lora_loaded, atol=1e-3, rtol=1e-3)) - - @require_torch_accelerator - def test_lora_group_offloading_delete_adapters(self): - components, _, denoiser_lora_config = self.get_dummy_components() - _, _, inputs = self.get_dummy_inputs(with_generator=False) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - - denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet - denoiser.add_adapter(denoiser_lora_config) - self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.") - - try: - with tempfile.TemporaryDirectory() as tmpdirname: - modules_to_save = self._get_modules_to_save(pipe, has_denoiser=True) - lora_state_dicts = self._get_lora_state_dicts(modules_to_save) - self.pipeline_class.save_lora_weights( - save_directory=tmpdirname, safe_serialization=True, **lora_state_dicts - ) - - components, _, _ = self.get_dummy_components() - pipe = self.pipeline_class(**components) - denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet - pipe.to(torch_device) - - # Enable Group Offloading (leaf_level for more granular testing) - apply_group_offloading( - denoiser, - onload_device=torch_device, - offload_device="cpu", - offload_type="leaf_level", - ) - - pipe.load_lora_weights(tmpdirname, adapter_name="default") - - out_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - - # Delete the adapter - pipe.delete_adapters("default") - - out_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - - self.assertFalse(np.allclose(out_lora, out_no_lora, atol=1e-3, rtol=1e-3)) - finally: - # Clean up the hooks to prevent state leak - if hasattr(denoiser, "_diffusers_hook"): - denoiser._diffusers_hook.remove_hook(_GROUP_OFFLOADING, recurse=True) diff --git a/tests/pipelines/mochi/test_mochi.py b/tests/pipelines/mochi/test_mochi.py index d3446f4fc71b..97eabf737384 100644 --- a/tests/pipelines/mochi/test_mochi.py +++ b/tests/pipelines/mochi/test_mochi.py @@ -27,12 +27,15 @@ numpy_cosine_similarity_distance, require_big_accelerator, require_torch_accelerator, + skip_mps, torch_device, ) from ..testing_utils import ( BasePipelineTesterConfig, FasterCacheTesterMixin, FirstBlockCacheTesterMixin, + LoraMemoryTesterMixin, + LoraTesterMixin, MemoryTesterMixin, PipelineTesterMixin, ) @@ -143,6 +146,16 @@ def test_vae_tiling(self, expected_diff_max: float = 0.2): ) +@skip_mps +class TestMochiPipelineLoRA(MochiPipelineTesterConfig, LoraTesterMixin): + pass + + +@skip_mps +class TestMochiPipelineLoRAMemory(MochiPipelineTesterConfig, LoraMemoryTesterMixin): + pass + + class TestMochiPipelineMemory(MochiPipelineTesterConfig, MemoryTesterMixin): """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Mochi pipeline.""" diff --git a/tests/pipelines/sana/test_sana.py b/tests/pipelines/sana/test_sana.py index 13da387e644f..000dcd38ff6d 100644 --- a/tests/pipelines/sana/test_sana.py +++ b/tests/pipelines/sana/test_sana.py @@ -33,6 +33,8 @@ ) from ..testing_utils import ( BasePipelineTesterConfig, + LoraMemoryTesterMixin, + LoraTesterMixin, MemoryTesterMixin, PipelineTesterMixin, ) @@ -191,6 +193,14 @@ def test_half_precision_inference_no_nan(self, dtype): super().test_half_precision_inference_no_nan(dtype) +class TestSanaPipelineLoRA(SanaPipelineTesterConfig, LoraTesterMixin): + pass + + +class TestSanaPipelineLoRAMemory(SanaPipelineTesterConfig, LoraMemoryTesterMixin): + pass + + class TestSanaPipelineMemory(SanaPipelineTesterConfig, MemoryTesterMixin): """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Sana pipeline.""" diff --git a/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3.py b/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3.py index 833bd01b492d..9163924f7a9a 100644 --- a/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3.py +++ b/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3.py @@ -14,6 +14,7 @@ from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler, SD3Transformer2DModel, StableDiffusion3Pipeline +from ...models.testing_utils.lora import check_if_lora_correctly_set from ...testing_utils import ( assert_tensors_close, backend_empty_cache, @@ -24,6 +25,8 @@ ) from ..testing_utils import ( BasePipelineTesterConfig, + LoraMemoryTesterMixin, + LoraTesterMixin, MemoryTesterMixin, PipelineTesterMixin, ) @@ -170,6 +173,21 @@ def test_skip_guidance_layers(self): assert output_full.shape == output_skip.shape, "Outputs should have the same shape" +class TestStableDiffusion3PipelineLoRA(StableDiffusion3PipelineTesterConfig, LoraTesterMixin): + @pytest.mark.parametrize("weight_name", ["lora_diffusers_format.safetensors", "lora_peft_format.safetensors"]) + def test_sd3_lora(self, weight_name): + pipe = self.get_pipeline().to(torch_device) + pipe.load_lora_weights("hf-internal-testing/tiny-sd3-loras", weight_name=weight_name) + assert check_if_lora_correctly_set(pipe.transformer) + self.run_pipe(pipe) + pipe.unload_lora_weights() + assert not check_if_lora_correctly_set(pipe.transformer) + + +class TestStableDiffusion3PipelineLoRAMemory(StableDiffusion3PipelineTesterConfig, LoraMemoryTesterMixin): + pass + + class TestStableDiffusion3PipelineMemory(StableDiffusion3PipelineTesterConfig, MemoryTesterMixin): """Memory optimization tests (CPU offload, group offload, layerwise casting) for the SD3 pipeline.""" diff --git a/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3_img2img.py b/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3_img2img.py index 7f69411c9c8c..df84f38b1d13 100644 --- a/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3_img2img.py +++ b/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3_img2img.py @@ -26,8 +26,12 @@ assert_tensors_close, backend_empty_cache, floats_tensor, + is_lora, + nightly, numpy_cosine_similarity_distance, require_big_accelerator, + require_peft_backend, + require_torch_accelerator, slow, torch_device, ) @@ -209,3 +213,57 @@ def test_sd3_img2img_inference(self): max_diff = numpy_cosine_similarity_distance(expected_slice.flatten(), image_slice.flatten()) assert max_diff < 1e-4, f"Outputs are not close enough, got {max_diff}" + + +@is_lora +@nightly +@require_torch_accelerator +@require_peft_backend +@require_big_accelerator +class TestStableDiffusion3Img2ImgLoRAIntegration: + pipeline_class = StableDiffusion3Img2ImgPipeline + repo_id = "stabilityai/stable-diffusion-3-medium-diffusers" + + @pytest.fixture(autouse=True) + def cleanup(self): + gc.collect() + backend_empty_cache(torch_device) + yield + gc.collect() + backend_empty_cache(torch_device) + + def get_inputs(self, device, seed=0): + init_image = load_image( + "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/bird_canny.png" + ) + if str(device).startswith("mps"): + generator = torch.manual_seed(seed) + else: + generator = torch.Generator(device="cpu").manual_seed(seed) + + return { + "prompt": "corgi", + "num_inference_steps": 2, + "guidance_scale": 5.0, + "output_type": "np", + "generator": generator, + "image": init_image, + } + + def test_sd3_img2img_lora(self): + pipe = self.pipeline_class.from_pretrained(self.repo_id, torch_dtype=torch.float16) + pipe.load_lora_weights("zwloong/sd3-lora-training-rank16-v2") + pipe.fuse_lora() + pipe.unload_lora_weights() + pipe = pipe.to(torch_device) + + inputs = self.get_inputs(torch_device) + + image = pipe(**inputs).images[0] + image_slice = image[0, -3:, -3:] + expected_slice = np.array([0.5649, 0.5405, 0.5488, 0.5688, 0.5449, 0.5513, 0.5337, 0.5107, 0.5059]) + + max_diff = numpy_cosine_similarity_distance(expected_slice.flatten(), image_slice.flatten()) + + assert max_diff < 1e-4, f"Outputs are not close enough, got {max_diff}" + pipe.unload_lora_weights() diff --git a/utils/extract_tests_from_mixin.py b/utils/extract_tests_from_mixin.py index e9b1982b8531..f218fd72a01c 100644 --- a/utils/extract_tests_from_mixin.py +++ b/utils/extract_tests_from_mixin.py @@ -71,9 +71,9 @@ def generate_pattern_for_mixins(mixin_classes: List[Type]) -> str: ] elif args.type == "lora": - from tests.lora.utils import PeftLoraLoaderMixinTests + from tests.pipelines.testing_utils import LoraMemoryTesterMixin, LoraTesterMixin, UNetLoraTesterMixin - mixin_classes = [PeftLoraLoaderMixinTests] + mixin_classes = [LoraTesterMixin, LoraMemoryTesterMixin, UNetLoraTesterMixin] pattern = generate_pattern_for_mixins(mixin_classes) print(pattern)