Fix LoRA loading for flattened CLIP text encoders - #14872
BeyondBirthday07 wants to merge 1 commit into
Conversation
| # CLIPTextModel is flattened in recent transformers versions, while older | ||
| # LoRA checkpoints can still contain the 'text_model.' prefix. | ||
| if not hasattr(text_encoder, "text_model"): | ||
| state_dict = {k.removeprefix("text_model."): v for k, v in state_dict.items()} |
There was a problem hiding this comment.
I guess this should work but the fix is very specific, I wonder if there is a more general solution.
There was a problem hiding this comment.
But this is also a very transformers-specific quirk, IMO. So, it seems okay to me. Do you have any ideas?
There was a problem hiding this comment.
Yeah i thought about making it more general too, but this looks pretty specific to the transformers CLIPTextModel flattening, so I kept the fix narrow.
| "text_encoder.text_model.encoder.layers.0.self_attn.q_proj.lora_B.weight": lora_b, | ||
| } | ||
|
|
||
| StableDiffusionLoraLoaderMixin.load_lora_into_text_encoder( |
There was a problem hiding this comment.
IMO, a test should be closer to how a user would load such a checkpoint via a pipeline, using the mixin class directly is not normally done. If the maintainers are fine with this, I won't object though.
There was a problem hiding this comment.
Yes, I agree with you, @BenjaminBossan. However, internally, the call gets dispatched to load_lora_into_text_encoder from load_lora_weights(). So, seems okay to me as a minimal and self-contained use case.
There was a problem hiding this comment.
Yeah, I kept the test at the loader level since load_lora_weights() ends up calling the same path anyway.
What does this PR do?
Fixes #14860
CLIPTextModelis flattened in newer transformers versions, so its module names no longer start withtext_model.. Older LoRA checkpoints can still contain this prefix.Because of that, the text encoder rank lookup does not find the LoRA keys, leaving
rankempty and causing anIndexError.This strips the
text_model.prefix when loading into a flattened text encoder.Added a regression test for this case. I also tested the fix with the LoRA checkpoint shared in #14860.
No new dependencies.
Tests
Before submitting
self-reviewskill on the diff?documentation guidelines, and
here are tips on formatting docstrings.
Who can review?
Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.
@sayakpaul @BenjaminBossan