VAE lora with AutoencoderKL triggers object has no attribute 'add_adapter' since 0.26.0

[beniz]

Hello!

Finetuning a VAE with a lora, using AutoencoderDL class, is somehow broken starting from 0.26.0. I am not filing a bug since I'm not sure whether there's actually a new appropriate way of doing this.

We're doing:

vae = AutoencoderKL.from_pretrained("stabilityai/sd-turbo", subfolder="vae")

target_modules_vae = [
            "conv1",
            "conv2",
             ...
]

vae_lora_config = LoraConfig(
            r=lora_rank_vae,
            init_lora_weights="gaussian",
            target_modules=target_modules_vae,
        )
vae.add_adapter(vae_lora_config, adapter_name="vae_skip")

The error is
AttributeError: 'AutoencoderKL' object has no attribute 'add_adapter'

FTR, this issue is reported in a couple places (but might be 'niche'):

• pix2pixTurbo: "AttributeError: 'AutoencoderKL' object has no attribute 'add_adapter'" radames/Real-Time-Latent-Consistency-Model#38
• 'AutoencoderKL' object has no attribute 'add_adapter' GaParmar/img2img-turbo#64

From looking at the diffusers codebase, there was a lora refactor between 0.25.1 and 0.26.0 (between Jan 17 2024 and Jan 31 2024) but I couldn't quite find nor understand where the critical change was and how to mitigate it.

Most importantly, maybe the above is not the proper way of finetuning a VAE with lora anymore, and I would be grateful for any help on getting it to work again.

(Our usage is custom img2img models, that do require a ft of the VAE)

[a-r-r-o-w]

cc @sayakpaul

[sayakpaul]

I checked the git blame here: https://github.com/huggingface/diffusers/blame/main/src/diffusers/models/autoencoders/autoencoder_kl.py#L37

I don't see any LoRA refactoring affecting that class in the first. Could you point me to a version here the class ever support LoRA loading?

[beniz]

Hi @sayakpaul thanks I actually did the same before posting and since I couldn't find something obvious, I've been looking at how the add_adapter was inherited.
Would you have any lead on this I can follow ?

0.25.1 does support the loading. If nothing obvious arises I could try a per commit test until breaks. Or maybe it's something with pedt dependencies ?

[sayakpaul]

If you could provide a minimal reproducible snippet where you can load a LoRA into the VAE and the library versions, I can look into it. Would that work?

[beniz]

Thanks @sayakpaul . The snippet below works with 0.25.1 and not with 0.26.0

from diffusers import AutoencoderKL
from peft import LoraConfig

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_skip")

yields:

Traceback (most recent call last):
  File "/path/to/test_vae_adapter.py", line 30, in <module>
    vae.add_adapter(vae_lora_config, adapter_name="vae_skip")
  File "/home/me/.local/lib/python3.10/site-packages/diffusers/models/modeling_utils.py", line 220, in __getattr__
    return super().__getattr__(name)
  File "/home/me/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1709, in __getattr__
    raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
AttributeError: 'AutoencoderKL' object has no attribute 'add_adapter'

[sayakpaul]

Thanks so much for proividing the easy reproducer -- much helpful! I see what's happening.

In v0.25.1, we had PEFT-related methods in the ModelMixin class:

diffusers/src/diffusers/models/modeling_utils.py

Line 306 in 1909952

def add_adapter(self, adapter_config, adapter_name: str = "default") -> None:

Which was really an anti-pattern and which is why we created PeftAdapterMixin:

diffusers/src/diffusers/loaders/peft.py

Line 40 in fddbab7

class PeftAdapterMixin:

However, we should have added this class as one of the parents of this mentioned VAE but it's not really common to do LoRAs on VAEs. So, we applied the necessary changes to the classes we knew the community to do LoRAs on. In fact, this is probably the first VAE + LoRA-related discussion.

Thankfully, the easy fix would be just to add PeftAdapterMixin of the AutoencoderKL class and hopefully, everything should be back to normal for you:
https://github.com/huggingface/diffusers/blob/fddbab79932eedf1a78041ef38c47df80ab84c90/src/diffusers/models/autoencoders/autoencoder_kl.py#L37C7-L37C20

Does this make sense? If so, would you maybe like to open a PR?

[beniz]

@sayakpaul perfect, thanks for confirming this (missing mixin + niche vae ft), I can do the PR.
Will reflect here once I have it.

[beniz]

Hey @sayakpaul the PR is here: #9807
The mini code check now passes fine with the patch, thanks for figuring it out so quickly :)