fix: missing AutoencoderKL lora adapter

Author: beniz

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.
 

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