Skip to content

Conversation

sayakpaul
Copy link
Member

@sayakpaul sayakpaul commented Dec 27, 2024

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:

...
pipeline.load_lora_weights("black-forest-labs/FLUX.1-Canny-dev-lora")
pipeline.fuse_lora()
pipeline.unload_lora_weights()
...

This would revert the transformer config 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 the unload_lora_weights() of FluxLoraLoaderMixin and suggesting to add a note about it in the docs.

@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.

Copy link
Contributor

@a-r-r-o-w a-r-r-o-w left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 <[email protected]>

# We override this here account for `_transformer_norm_layers`.
def unload_lora_weights(self):
def unload_lora_weights(self, reset_to_overwritten_params=False):
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not too fixated on the name. Open to suggestions.

Also open to tackling the problem in better ways.

@sayakpaul
Copy link
Member Author

@a-r-r-o-w docs and docstrings done.

Would maybe check if an error is raised when trying to run inference with reset_to_overwritten_params=True

I am not sure what is meant here.

@yiyixuxu yiyixuxu merged commit d9d94e1 into main Jan 6, 2025
15 checks passed
@yiyixuxu yiyixuxu deleted the fix-unloading-expanded-flux branch January 6, 2025 18:35
DN6 pushed a commit that referenced this pull request Jan 15, 2025
* fix: lora unloading when using expanded Flux LoRAs.

* fix argument name.

Co-authored-by: a-r-r-o-w <[email protected]>

* docs.

---------

Co-authored-by: a-r-r-o-w <[email protected]>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants