[LoRA] fix: lora unloading when using expanded Flux LoRAs. #10397
Conversation
|
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. |
There was a problem hiding this comment.
Maybe a docstring for what reset_to_overwritten_params does and why we added it would be helpful for someone looking at it in future, but no problem if it doesn't make sense to.
The test looks good. Would maybe check if an error is raised when trying to run inference with reset_to_overwritten_params=True. Off to YiYi
Co-authored-by: a-r-r-o-w <contact.aryanvs@gmail.com>
| @@ -2278,15 +2278,15 @@ def unfuse_lora(self, components: List[str] = ["transformer", "text_encoder"], * | |||
| super().unfuse_lora(components=components) | |||
|
|
|||
| # We override this here account for `_transformer_norm_layers`. | |||
| def unload_lora_weights(self): | |||
| def unload_lora_weights(self, reset_to_overwritten_params=False): | |||
There was a problem hiding this comment.
Not too fixated on the name. Open to suggestions.
Also open to tackling the problem in better ways.
|
@a-r-r-o-w docs and docstrings done.
I am not sure what is meant here. |
* fix: lora unloading when using expanded Flux LoRAs. * fix argument name. Co-authored-by: a-r-r-o-w <contact.aryanvs@gmail.com> * docs. --------- Co-authored-by: a-r-r-o-w <contact.aryanvs@gmail.com>
What does this PR do?
#10206 added support for properly unloading Control LoRA weights from Flux and allowing the pipeline to be used with
from_pipe().However, this led to the following problem. Say, if we do:
This would revert the
transformerconfig back to the original (in_channels=64), rendering this line to output 8 (whereas it should be 16). This is fine for use cases #10206 served but it's not effective for the above use case.This went uncaught because we didn't have a proper test for it and I didn't run the integration tests for Control LoRA. Sorry about that.
This PR fixes this. To do that, I am adding a flag
reset_to_overwritten_params(not fixated on the name) to theunload_lora_weights()ofFluxLoraLoaderMixinand suggesting to add a note about it in the docs.