'UNet2DConditionModel' object has no attribute 'peft_config'

[duduke121]

Describe the bug

After pulling the latest code from the examples/research_projects/instructpix2pix_lora repository, I encountered the following error when running the default code:
AttributeError: 'Linear' object has no attribute 'set_lora_layer'
I resolved this issue by modifying the code.
The training now runs successfully, but when I try to save the model, I encounter another error:
Traceback (most recent call last):7%|██████████████████████████████████████████████████████████████████████████████████████▊ | 4/7 [00:00<00:00, 17.90it/s] File "/home/guying/diffusers/examples/research_projects/instructpix2pix_lora/train_instruct_pix2pix_lora.py", line 1160, in <module> main() File "/home/guying/diffusers/examples/research_projects/instructpix2pix_lora/train_instruct_pix2pix_lora.py", line 1120, in main unet.save_attn_procs(args.output_dir, weights_name="pytorch_lora_weights.bin") File "/home/guying/diffusers/src/diffusers/loaders/unet.py", line 500, in save_attn_procs state_dict = get_peft_model_state_dict(self) File "/home/guying/anaconda3/envs/diffuser/lib/python3.9/site-packages/peft/utils/save_and_load.py", line 73, in get_peft_model_state_dict config = model.peft_config[adapter_name] File "/home/guying/diffusers/src/diffusers/models/modeling_utils.py", line 155, in __getattr__ return super().__getattr__(name) File "/home/guying/anaconda3/envs/diffuser/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1931, in __getattr__ raise AttributeError( AttributeError: 'UNet2DConditionModel' object has no attribute 'peft_config'

Reproduction

Pull the latest code from the repository.
Apply the code modification to train the model.
`
# # referred to https://github.com/huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image_lora.py
# unet_lora_parameters = []
# for attn_processor_name, attn_processor in unet.attn_processors.items():
# # Parse the attention module.
# attn_module = unet
# for n in attn_processor_name.split(".")[:-1]:
# attn_module = getattr(attn_module, n)
# print(type(attn_module.to_q)) # 添加调试输出

#     # Set the `lora_layer` attribute of the attention-related matrices.
#     attn_module.to_q.set_lora_layer(
#         LoRALinearLayer(
#             in_features=attn_module.to_q.in_features, out_features=attn_module.to_q.out_features, rank=args.rank
#         )
#     )
#     attn_module.to_k.set_lora_layer(
#         LoRALinearLayer(
#             in_features=attn_module.to_k.in_features, out_features=attn_module.to_k.out_features, rank=args.rank
#         )
#     )

#     attn_module.to_v.set_lora_layer(
#         LoRALinearLayer(
#             in_features=attn_module.to_v.in_features, out_features=attn_module.to_v.out_features, rank=args.rank
#         )
#     )
#     attn_module.to_out[0].set_lora_layer(
#         LoRALinearLayer(
#             in_features=attn_module.to_out[0].in_features,
#             out_features=attn_module.to_out[0].out_features,
#             rank=args.rank,
#         )
#     )

#     # Accumulate the LoRA params to optimize.
#     unet_lora_parameters.extend(attn_module.to_q.lora_layer.parameters())
#     unet_lora_parameters.extend(attn_module.to_k.lora_layer.parameters())
#     unet_lora_parameters.extend(attn_module.to_v.lora_layer.parameters())
#     unet_lora_parameters.extend(attn_module.to_out[0].lora_layer.parameters())


unet_lora_parameters = []
for attn_processor_name, attn_processor in unet.attn_processors.items():
    # Parse the attention module.
    attn_module = unet
    for n in attn_processor_name.split(".")[:-1]:
        attn_module = getattr(attn_module, n)

    # Replace the Linear layers with LoRALinearLayer
    # Creating new LoRA layers directly replacing the existing layers
    setattr(attn_module, 'to_q', LoRALinearLayer(
        in_features=attn_module.to_q.in_features, 
        out_features=attn_module.to_q.out_features, 
        rank=args.rank)
    )
    setattr(attn_module, 'to_k', LoRALinearLayer(
        in_features=attn_module.to_k.in_features, 
        out_features=attn_module.to_k.out_features, 
        rank=args.rank)
    )
    setattr(attn_module, 'to_v', LoRALinearLayer(
        in_features=attn_module.to_v.in_features, 
        out_features=attn_module.to_v.out_features, 
        rank=args.rank)
    )
    # Assuming `to_out` is a sequence and the first element is the linear layer
    setattr(attn_module.to_out, '0', LoRALinearLayer(
        in_features=attn_module.to_out[0].in_features,
        out_features=attn_module.to_out[0].out_features,
        rank=args.rank)
    )

    # Add the parameters of the newly created LoRA layers to the parameters list
    unet_lora_parameters.extend(attn_module.to_q.parameters())
    unet_lora_parameters.extend(attn_module.to_k.parameters())
    unet_lora_parameters.extend(attn_module.to_v.parameters())
    unet_lora_parameters.extend(attn_module.to_out[0].parameters())

`

I tried the following bash command to launch the training:
`export MODEL_ID="timbrooks/instruct-pix2pix"
export DATASET_ID="keaiguogongzhu/relight-dataset"
export OUTPUT_DIR="/home/guying/relight_nerf/instruct-nerf2nerf/instructPix2Pix-relighting_4"

accelerate launch /home/guying/diffusers/examples/research_projects/instructpix2pix_lora/train_instruct_pix2pix_lora.py
--pretrained_model_name_or_path=$MODEL_ID
--dataset_name=$DATASET_ID
--enable_xformers_memory_efficient_attention
--resolution=256 --random_flip
--train_batch_size=4 --gradient_accumulation_steps=2 --gradient_checkpointing
--max_train_steps=100
--checkpointing_steps=50 --checkpoints_total_limit=5
--learning_rate=5e-05 --lr_warmup_steps=0
--seed=42
--rank=4
--output_dir=$OUTPUT_DIR
--report_to=wandb `

Logs

No response

System Info

• 🤗 Diffusers version: 0.32.0.dev0
• Platform: Linux-5.15.0-118-generic-x86_64-with-glibc2.31
• Running on Google Colab?: No
• Python version: 3.9.20
• PyTorch version (GPU?): 2.5.1+cu124 (True)
• Flax version (CPU?/GPU?/TPU?): not installed (NA)
• Jax version: not installed
• JaxLib version: not installed
• Huggingface_hub version: 0.26.2
• Transformers version: 4.46.3
• Accelerate version: 1.1.1
• PEFT version: 0.13.2
• Bitsandbytes version: 0.44.1
• Safetensors version: 0.4.5
• xFormers version: 0.0.28.post3
• Accelerator: NVIDIA GeForce RTX 3080 Ti Laptop GPU, 16384 MiB
• Using GPU in script?:
• Using distributed or parallel set-up in script?:

Who can help?

@sayakpaul Hi, sorry to bother you, but I’ve run into a problem with my modifications and am hoping you could assist me in resolving it. I really appreciate any help you can provide!

[sayakpaul]

Cc: @Aiden-Frost.

This is happening because we have long deprecated non-PEFT backend for all things LoRA. So, you should either downgrade to the diffusers version that supported this or adjust the script accordingly to use PEFT.

I think the second option is even better and you could refer to https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/train_dreambooth_lora_sdxl.py as an example to see what kind of changes are required.

I know this is a bit involved but I think it could make the script better in the long run.

Let us know.

[Aiden-Frost]

Sure, I'll be able to work on this. I will keep you posted.

[duduke121]

Thank you for your previous advice! I have made the following modifications to the code to enable saving .bin files after training. The specific implementation is as follows:
`
from peft import LoraConfig, get_peft_model
unet_lora_parameters = []

for attn_processor_name, attn_processor in unet.attn_processors.items():
    attn_module = unet
    for n in attn_processor_name.split(".")[:-1]:
        attn_module = getattr(attn_module, n)

    # 检查并筛选出支持 LoRA 的层（Linear）
    target_modules = []
    for submodule_name, submodule in attn_module.named_modules():
        if isinstance(submodule, torch.nn.Linear):
            target_modules.append(submodule)

    # 为每个 Linear 层配置 LoRA
    if target_modules:
        lora_config = LoraConfig(
            r=args.rank,
            lora_alpha=args.rank,
            target_modules=["to_q", "to_k", "to_v", "to_out.0"],
            lora_dropout=0.1,
        )

        # 使用 PEFT 获取 LoRA 层
        peft_model = get_peft_model(attn_module, lora_config)

        # 确保 LoRA 配置被绑定
        assert hasattr(peft_model, 'peft_config'), "LoRA config is not correctly attached to the model"

        # 保存 LoRA 权重的函数
        def save_lora_weights(model, output_dir, weights_name="pytorch_lora_weights.bin"):
            lora_weights = model.state_dict()
            torch.save(lora_weights, os.path.join(output_dir, weights_name))

        # 保存 LoRA 权重
        # save_lora_weights(peft_model, args.output_dir)
        # 手动保存 LoRA 权重
        save_lora_weights(peft_model, args.output_dir, weights_name="manual_lora_weights.bin")

        # 收集 LoRA 参数
        for name, param in peft_model.named_parameters():
            unet_lora_parameters.append(param)

        # 检查 peft_model 是否已绑定 LoRA 配置
        # print(hasattr(peft_model, 'peft_config'))  # True 如果 peft_config 已绑定
        # print(getattr(peft_model, 'peft_config', None))  # 打印 peft_config，查看其内容

        # 如果 'peft_config' 不存在，手动绑定
        if not hasattr(unet, 'peft_config'):
            unet.peft_config = { 'default': lora_config }

        # # 使用 PEFT 模型保存
        # if hasattr(peft_model, 'save_attn_procs'):
        #     peft_model.save_attn_procs(args.output_dir, weights_name="pytorch_lora_weights.bin")
        # else:
        #     print("PEFT model does not have 'save_attn_procs' method.")

# Create the pipeline using the trained modules and save it.
accelerator.wait_for_everyone()
if accelerator.is_main_process:
unet = accelerator.unwrap_model(unet)
if args.use_ema:
ema_unet.copy_to(unet.parameters())

    pipeline = StableDiffusionInstructPix2PixPipeline.from_pretrained(
        args.pretrained_model_name_or_path,
        text_encoder=accelerator.unwrap_model(text_encoder),
        vae=accelerator.unwrap_model(vae),
        unet=unet,
        revision=args.revision,
        variant=args.variant,
    )
    # store only LORA layers
    # unet.save_attn_procs(args.output_dir)
    unet.save_attn_procs(args.output_dir, weights_name="pytorch_lora_weights.bin")

`

The training process completed successfully, and the generated weight files can be loaded during inference without issues.

It seems that the model saves and loads the LoRA weights as intended, and the inference results meet expectations.
However, I still have some questions and would appreciate your guidance on the following points:

1. Does this implementation fully align with the project's design logic, particularly in terms of saving and attaching the LoRA weights?
2. Are there any suggestions for more standard or optimized ways to achieve this, especially regarding integration with PEFT?
3. Regarding future compatibility: the deprecation warning for the load_attn_procs method—are there any recommended alternatives for future-proofing this implementation?

Please let me know if there are any issues or improvements I can make to this approach. Thank you！

[sayakpaul]

Thank you very much!

Does this implementation fully align with the project's design logic, particularly in terms of saving and attaching the LoRA weights?

You can refer to this script to see how LoRA is attached and saved.

Attachement:

diffusers/examples/dreambooth/train_dreambooth_lora_sdxl.py

Lines 1204 to 1207 in 98d0cd5

	# now we will add new LoRA weights to the attention layers
	unet_target_modules = ["to_k", "to_q", "to_v", "to_out.0"]
	unet_lora_config = get_lora_config(rank=args.rank, use_dora=args.use_dora, target_modules=unet_target_modules)
	unet.add_adapter(unet_lora_config)

We don't have to go over the attention processors for this.

Saving:

diffusers/examples/dreambooth/train_dreambooth_lora_sdxl.py

Lines 1907 to 1929 in 98d0cd5

	unet = unwrap_model(unet)
	unet = unet.to(torch.float32)
	unet_lora_layers = convert_state_dict_to_diffusers(get_peft_model_state_dict(unet))
	if args.train_text_encoder:
	text_encoder_one = unwrap_model(text_encoder_one)
	text_encoder_lora_layers = convert_state_dict_to_diffusers(
	get_peft_model_state_dict(text_encoder_one.to(torch.float32))
	)
	text_encoder_two = unwrap_model(text_encoder_two)
	text_encoder_2_lora_layers = convert_state_dict_to_diffusers(
	get_peft_model_state_dict(text_encoder_two.to(torch.float32))
	)
	else:
	text_encoder_lora_layers = None
	text_encoder_2_lora_layers = None
	StableDiffusionXLPipeline.save_lora_weights(
	save_directory=args.output_dir,
	unet_lora_layers=unet_lora_layers,
	text_encoder_lora_layers=text_encoder_lora_layers,
	text_encoder_2_lora_layers=text_encoder_2_lora_layers,
	)

We can also use save_lora_adapter(). I think this also answers 2.

Regarding future compatibility: the deprecation warning for the load_attn_procs method—are there any recommended alternatives for future-proofing this implementation?

Exactly as the deprecation message says. Prefer using load_lora_adapter(). But for the purpose of training, I don't think that's needed as also evident from the above training example. But if you want to give load_lora_adapter() and save_lora_adapter() a try, then here's your reference:

diffusers/tests/models/test_modeling_common.py

Line 916 in 98d0cd5

def test_save_load_lora_adapter(self, use_dora=False):

LMK in case anything is unclear.

[Aiden-Frost]

Hey @duduke121 @sayakpaul

I have created a PR: #10205
Do let me know if the script works or any changes required.
I have tested on Google Colab with T4 GPU.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.