fix: missing AutoencoderKL lora adapter

Author: beniz

src/diffusers/models/autoencoders/autoencoder_kl.py

@@ -31,10 +31,11 @@
 )
 from ..modeling_outputs import AutoencoderKLOutput
 from ..modeling_utils import ModelMixin
+from ...loaders import PeftAdapterMixin
 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

@@ -846,6 +846,39 @@ def test_stable_diffusion_encode_sample(self, seed, expected_slice):
         assert torch_all_close(output_slice, expected_output_slice, atol=tolerance)
 
 
+    def test_lora_adapter(self):
+        
+        vae = AutoencoderKL.from_pretrained("stabilityai/sd-turbo", subfolder="vae")
+        
+        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")
+        
+        
 @slow
 class AsymmetricAutoencoderKLIntegrationTests(unittest.TestCase):
     def get_file_format(self, seed, shape):