From 7025526f64d1d1e1e7dc60aff7f1aa8b5aa6566a Mon Sep 17 00:00:00 2001 From: Emmanuel Benazera Date: Wed, 30 Oct 2024 07:39:19 +0000 Subject: [PATCH 1/2] fix: missing AutoencoderKL lora adapter --- .../models/autoencoders/autoencoder_kl.py | 3 +- tests/models/autoencoders/test_models_vae.py | 37 +++++++++++++++++++ 2 files changed, 39 insertions(+), 1 deletion(-) diff --git a/src/diffusers/models/autoencoders/autoencoder_kl.py b/src/diffusers/models/autoencoders/autoencoder_kl.py index 99a7da4a0b6f..9036c027a535 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl.py @@ -17,6 +17,7 @@ import torch.nn as nn from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import PeftAdapterMixin from ...loaders.single_file_model import FromOriginalModelMixin from ...utils import deprecate from ...utils.accelerate_utils import apply_forward_hook @@ -34,7 +35,7 @@ from .vae import Decoder, DecoderOutput, DiagonalGaussianDistribution, Encoder -class AutoencoderKL(ModelMixin, ConfigMixin, FromOriginalModelMixin): +class AutoencoderKL(ModelMixin, ConfigMixin, FromOriginalModelMixin, PeftAdapterMixin): r""" A VAE model with KL loss for encoding images into latents and decoding latent representations into images. diff --git a/tests/models/autoencoders/test_models_vae.py b/tests/models/autoencoders/test_models_vae.py index 0188f9121ae0..f8516edb2614 100644 --- a/tests/models/autoencoders/test_models_vae.py +++ b/tests/models/autoencoders/test_models_vae.py @@ -21,6 +21,10 @@ from datasets import load_dataset from parameterized import parameterized + +if is_peft_available(): + from peft import LoraConfig + from diffusers import ( AsymmetricAutoencoderKL, AutoencoderKL, @@ -36,6 +40,7 @@ backend_empty_cache, enable_full_determinism, floats_tensor, + is_peft_available, load_hf_numpy, require_torch_accelerator, require_torch_accelerator_with_fp16, @@ -299,6 +304,38 @@ def test_output_pretrained(self): self.assertTrue(torch_all_close(output_slice, expected_output_slice, rtol=1e-2)) + @require_peft_backend + def test_lora_adapter(self): + init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() + vae = self.model_class(**init_dict) + + target_modules_vae = [ + "conv1", + "conv2", + "conv_in", + "conv_shortcut", + "conv", + "conv_out", + "skip_conv_1", + "skip_conv_2", + "skip_conv_3", + "skip_conv_4", + "to_k", + "to_q", + "to_v", + "to_out.0", + ] + vae_lora_config = LoraConfig( + r=16, + init_lora_weights="gaussian", + target_modules=target_modules_vae, + ) + + vae.add_adapter(vae_lora_config, adapter_name="vae_lora") + active_lora = vae.active_adapters() + self.assertTrue(len(active_lora) == 1) + self.assertTrue(active_lora[0] == "vae_lora") + class AsymmetricAutoencoderKLTests(ModelTesterMixin, UNetTesterMixin, unittest.TestCase): model_class = AsymmetricAutoencoderKL From ecda35049032da4a541850c1e9b780f1ec537896 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Fri, 29 Nov 2024 16:57:44 +0530 Subject: [PATCH 2/2] fix --- tests/models/autoencoders/test_models_vae.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/tests/models/autoencoders/test_models_vae.py b/tests/models/autoencoders/test_models_vae.py index d0320a1935a5..d475160cc796 100644 --- a/tests/models/autoencoders/test_models_vae.py +++ b/tests/models/autoencoders/test_models_vae.py @@ -21,10 +21,6 @@ from datasets import load_dataset from parameterized import parameterized - -if is_peft_available(): - from peft import LoraConfig - from diffusers import ( AsymmetricAutoencoderKL, AutoencoderKL, @@ -42,6 +38,7 @@ floats_tensor, is_peft_available, load_hf_numpy, + require_peft_backend, require_torch_accelerator, require_torch_accelerator_with_fp16, require_torch_gpu, @@ -55,6 +52,10 @@ from ..test_modeling_common import ModelTesterMixin, UNetTesterMixin +if is_peft_available(): + from peft import LoraConfig + + enable_full_determinism()