[LoRA] Improve warning messages when LoRA loading becomes a no-op

[sayakpaul]

What does this PR do?

As discussed with Dhruv, this PR improves our logging for LoRAs when the specified state dict isn't an appropriate one, leaving the load_lora_weights(), load_lora_into_text_encoder(), and load_lora_adapter() methods essentially a no-op.

#9950 mentions this problem, and I think it's better to at least let the user know that the specified LoRA state dict is either not a correct one (which we already do) or is an ineffective one (this PR).

This PR also standardizes a bit how we log these kinds of info across our LoRA loading utilities.

Some comments are in-line.

• Run the integration tests for SD LoRA
• Run the integration tests for SDXL LoRA
• Run the integration tests for Flux LoRA (FluxControlLoRA integration tests will pass after [LoRA] fix: lora unloading when using expanded Flux LoRAs. #10397)
• Run the integration tests for SD3 LoRA

[sayakpaul]

We're handling all of this within the load_lora_adapter() method now which I think is more appropriate as:

1. It takes care of logging when users try to load LoRA into a model via pipe.load_lora_weights().
2. Users try to load LoRAs directly into a model with the load_lora_adapter() method (with something like unet.load_lora_adapter().

Helps to avoid duplication. I have run the integration tests, too and nothing is breaking due to this.

[sayakpaul]

Similar philosophy as explained above.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[sayakpaul]

In case prefix is not None and no prefix matched state dict keys are found, we log from the load_lora_adapter() method.

This way, we cover for both load_lora_weights() which is pipeline-level and load_lora_adapter() which is model-level.

[hlky]

Thanks @sayakpaul

[sayakpaul]

Ran the integration tests, too. They are passing. Will merge once the CI is green.

[BenjaminBossan]

@sayakpaul This PR broke LoRA loading. It seems like after calling load_lora_adapter, there is no LoRA being loaded. Reproducer:

import os
import tempfile

import torch

from peft import LoraConfig
from peft.tuners.lora import LoraLayer
from diffusers import UNet2DConditionModel
from diffusers.utils.testing_utils import floats_tensor

torch_device = 0

class TestIt:
    def get_small_unet(self):
        # from diffusers UNet2DConditionModelTests
        init_dict = {
            "block_out_channels": (4, 8),
            "norm_num_groups": 4,
            "down_block_types": ("CrossAttnDownBlock2D", "DownBlock2D"),
            "up_block_types": ("UpBlock2D", "CrossAttnUpBlock2D"),
            "cross_attention_dim": 8,
            "attention_head_dim": 2,
            "out_channels": 4,
            "in_channels": 4,
            "layers_per_block": 1,
            "sample_size": 16,
        }
        model = UNet2DConditionModel(**init_dict)
        return model.to(torch_device)

    def test_it(self):
        lora_config = LoraConfig(
            target_modules=["to_q"],
            init_lora_weights=False,
        )
        unet = self.get_small_unet()
        unet.add_adapter(lora_config, adapter_name="adapter0")
        num_lora_before = sum(isinstance(m, LoraLayer) for m in unet.modules())
        assert num_lora_before > 0

        with tempfile.TemporaryDirectory() as tmp_dirname:
            unet.save_lora_adapter(tmp_dirname, safe_serialization=True, adapter_name="adapter0")
            del unet

            torch.manual_seed(0)
            unet = self.get_small_unet()
            file_name = os.path.join(tmp_dirname, "pytorch_lora_weights.safetensors")
            unet.load_lora_adapter(file_name, safe_serialization=True, adapter_name="adapter0")
            num_lora_after = sum(isinstance(m, LoraLayer) for m in unet.modules())

            # this fails, no LoRA layer is loaded
            assert num_lora_after > 0

The previous commit, 0703ce88008b2765ef6636c6e5cb013d227c42ca, still works. I think this is quite urgent.

[sayakpaul]

Use:

unet.load_lora_adapter(file_name, safe_serialization=True, prefix=None, adapter_name="adapter0")

This is reflected in the tests:

diffusers/tests/models/test_modeling_common.py

Line 1094 in 733b44a

model.load_lora_adapter(tmpdir, prefix=None, use_safetensors=True)

[BenjaminBossan]

Thanks for the quick reply. Isn't this a backwards incompatible change on a public API? At least for me this was surprising, especially since there is no error message or warning.

[sayakpaul]

Can you check if you're hitting?

diffusers/src/diffusers/loaders/peft.py

Line 357 in 733b44a

logger.info(

If not, then there's something wrong that I can look into fixing.

[BenjaminBossan]

there's something wrong that I can look into

When I load without prefix=None, I hit this code path:

"No LoRA keys associated to UNet2DConditionModel found with the prefix='transformer'. This is safe to ignore if LoRA state dict didn't originally have any UNet2DConditionModel related params. Open an issue if you think it's unexpected: https://github.com/huggingface/diffusers/issues/new"

[sayakpaul]

Earlier you said

At least for me this was surprising, especially since there is no error message or warning.

But you are indeed hitting the right message here. Would you suggest turning this into a warning?

[BenjaminBossan]

Since logger.info is not shown by default, I had no clue about this message. At least to me, when I call

unet.load_lora_adapter(file_name, safe_serialization=True, adapter_name="adapter0")

but nothing happens, it would be useful to get a warning. Maybe there is a good reason not to have a warning because there are legitimate use cases, but it feels dangerous to me. As mentioned, everything appears to work on the surface, but the outputs are not what I expect, so this could lead to hard to debug silent errors in code that used to work.

Also, I don't find the message to be very helpful, since the solution you provided is not mentioned. But perhaps in the bigger context of things, this is obvious?

[sayakpaul]

#11044