FLUX LoRA Support #6847

RyanJDick · 2024-09-12T23:05:30Z

Summary

This PR adds support for FLUX LoRA models on both quantized and non-quantized base models.

Supported formats:

diffusers
kohya

Full changelist:

Consolidated LoRA handling code in invokeai/backend/lora
Add support for FLUX kohya and FLUX diffusers LoRA model loading
Add ability to either patch LoRAs or run as a sidecar model (the latter enables LoRAs to be applied to a wide range of quantized models).

QA Instructions

Note to reviewers: I tested everything in this checklist. Feel free to re-verify any of this, but also test any LoRAs that you have. There are many small LoRA format variations, and there's a risk of breaking one of them with this change.

FLUX LoRA

Import / probe of kohya FLUX LoRA (https://civitai.com/models/159333/pokemon-trainer-sprite-pixelart?modelVersionId=779247)
Import / probe of Diffusers FLUX LoRA (https://civitai.com/models/200255/hands-xl-sd-15-flux1-dev?modelVersionId=781855)
kohya with non-quantized base model
kohya with quantized base model (should roughly match the non-quantized case)
diffusers with non-quantized base model
diffusers with quantized base model (should roughly match the non-quantized case)
Sidecar LoRA patching speed (<0.1secs after model is loaded)
Stacking multiple fused LoRA models (i.e. on top on non-quantized model)
Stacking multiple sidecar LoRA models (i.e. on top of quantized model)

Regression Tests

SD1.5 LoRA (check output, speed and memory)
SDXL LoRA (check output, speed and memory)
USE_MODULAR_DENOISE=1 smoke test with LoRA

Test for output regression with the following LoRA formats:

LoRA
LoHA
LoKr
IA3

Checklist

The PR has a short but descriptive title, suitable for a changelog
Tests added / updated (if applicable)
Documentation added / updated (if applicable)

invokeai/backend/model_manager/probe.py

RyanJDick · 2024-09-13T18:36:06Z

Note to reviewers: Please test any LoRAs that you have (SD or FLUX). There are many small LoRA format variations, and there's a risk of breaking one of them with this change.

invokeai/backend/model_manager/load/model_loaders/lora.py

bghira · 2024-09-14T01:16:46Z

i don't see how alpha=8 would work for any peft loras that aren't also rank=8

… unit tests.

…e LoRAModelRaw class.

…SDXL lora key conversions, out of LoRAModelRaw and into LoRALoader.

…xception if an unexpected key is encountered, and add a corresponding unit test.

…X kohya LoRA format.

…ass.

…sted.

…idecarLayer.

…xup dtype handling in the sidecar layers.

… in the hopes that the latter works properly on MacOS.

…de is uglier, it turns out that the Module implementation of some ops like .to(...) is noticeably slower.

bghira · 2024-09-15T18:24:35Z

if you're going for lycoris support note that the support for this in peft is very lacking currently, with an open issue for community help request on updating the various modules

it's very easy to wrap the upstream lycoris library modules though, as @KohakuBlueleaf has added a functional API for this use case. it just wraps the Diffusers model and has .apply_to() and .merge_to() for sidecar or fused running

there's built in methods for extracting models and playing with them in other ways. i'd highly recommend switching from peft for lycoris stuff.

KohakuBlueleaf · 2024-09-15T23:20:41Z

it's very easy to wrap the upstream lycoris library modules though, as @KohakuBlueleaf has added a functional API for this use case. it just wraps the Diffusers model and has .apply_to() and .merge_to() for sidecar or fused running

.apply_to() is module API. Functional API is something designed for library to implement their own low-level wrapper.
(For example, lokr_output = layer(x) + lokr_bypass_forward(x, w1, w2a, w2b ...))

…ches(), and add unit tests for the stacked LoRA case.

RyanJDick · 2024-09-16T14:59:10Z

if you're going for lycoris support note that the support for this in peft is very lacking currently, with an open issue for community help request on updating the various modules

it's very easy to wrap the upstream lycoris library modules though, as @KohakuBlueleaf has added a functional API for this use case. it just wraps the Diffusers model and has .apply_to() and .merge_to() for sidecar or fused running

there's built in methods for extracting models and playing with them in other ways. i'd highly recommend switching from peft for lycoris stuff.

Thanks for this context. We explored using PEFT at one point, but decided against using it because it didn't meet our requirements for high patching/unpatching speeds.

I'll leave the integration of the upstream lycoris lib for future work. For now, we have our own implementation of the lycoris layers for fused execution, and partial support for sidecar execution (and easy to extend).

bghira · 2024-09-16T15:12:23Z

i agree that the patching speed of diffusers could be improved. if you have any observations/suggestions please open an issue report

invokeai/app/invocations/flux_lora_loader.py

invokeai/app/invocations/flux_denoise.py

invokeai/backend/lora/lora_patcher.py

invokeai/backend/model_manager/probe.py

github-actions bot added python PRs that change python files invocations PRs that change invocations backend PRs that change backend files python-tests PRs that change python tests labels Sep 12, 2024

RyanJDick commented Sep 13, 2024

View reviewed changes

invokeai/backend/model_manager/probe.py Show resolved Hide resolved

RyanJDick force-pushed the ryan/flux-lora-quantized branch 2 times, most recently from 9c2f173 to c3aa092 Compare September 13, 2024 15:25

RyanJDick marked this pull request as ready for review September 13, 2024 18:35

RyanJDick requested review from lstein, blessedcoolant, brandonrising, hipsterusername and psychedelicious as code owners September 13, 2024 18:35

bghira reviewed Sep 14, 2024

View reviewed changes

invokeai/backend/model_manager/load/model_loaders/lora.py Outdated Show resolved Hide resolved

RyanJDick added 14 commits September 15, 2024 04:39

Add state_dict keys for two FLUX LoRA formats to be used in unit tests.

7a80d9e

WIP - Initial logic for kohya FLUX LoRA conversion.

c41bd59

Get convert_flux_kohya_state_dict_to_invoke_format(...) working, with…

ade75b4

… unit tests.

Start moving SDXL-specific LoRA conversions out of the general-purpos…

fc380f0

…e LoRAModelRaw class.

Fix type errors in sdxl_lora_conversion_utils.py

d0d91ea

Remove unused LoRAModelRaw.name attribute.

8518ae9

Move the responsibilities of 1) state_dict loading from file, and 2) …

04b37e6

…SDXL lora key conversions, out of LoRAModelRaw and into LoRALoader.

Update convert_flux_kohya_state_dict_to_invoke_format() to raise an e…

7b5befa

…xception if an unexpected key is encountered, and add a corresponding unit test.

Add utility function for detecting whether a state_dict is in the FLU…

00e5686

…X kohya LoRA format.

Get probing of FLUX LoRA kohya models working.

db61ec4

WIP - add invocations to support FLUX LORAs.

01a15b4

WIP

50c9410

Fixup FLUX LoRA unit tests.

92b8477

Rename flux_kohya_lora_conversion_utils.py

cf9f30c

RyanJDick added 17 commits September 15, 2024 04:39

WIP - adding LoRA sidecar layers

049ce18

WIP - LoRA sidecar layers.

3e12ac9

Bug fixes to get LoRA sidecar patching working for the first time.

f5f8944

WIP - Implement sidecar LoRA layers using functional API.

45bc8fc

Get diffusers FLUX LoRA working as sidecar patch on quantized model.

10c3c61

Assume LoRA alpha=8 for FLUX diffusers PEFT LoRAs.

81fbaf2

Update all lycoris layer types to use the new torch.nn.Module base cl…

9438ea6

…ass.

Add links to test models for loha, lokr, ia3.

5bb0c79

Fixup unit tests.

7ce41bf

Remove LoRA conv sidecar layers until they are needed and properly te…

ae41651

…sted.

Minor cleanup and documentation updates.

61d3d56

Add unit tests for LoRALinearSidecarLayer and ConcatenatedLoRALinearS…

ba3ba3c

…idecarLayer.

Add unit tests for LoRAPatcher.apply_lora_sidecar_patches(...) and fi…

02f27c7

…xup dtype handling in the sidecar layers.

Delete duplicate file that was accidentally kept during rebase.

9466824

Replace 'torch.device("meta")' with 'accelerate.init_empty_weights()'…

b1cf5e9

… in the hopes that the latter works properly on MacOS.

Revert change of make all LoRA layers torch.nn.Module's. While the co…

78efed4

…de is uglier, it turns out that the Module implementation of some ops like .to(...) is noticeably slower.

Add bias to LoRA sidecar layer unit tests.

d51f2c5

hipsterusername force-pushed the ryan/flux-lora-quantized branch from 08e2b76 to d51f2c5 Compare September 15, 2024 01:39

RyanJDick added 2 commits September 16, 2024 13:57

Assume alpha=rank for FLUX diffusers PEFT LoRA models.

e88d3cf

Fix bug when applying multiple LoRA models via apply_lora_sidecar_pat…

2934e31

…ches(), and add unit tests for the stacked LoRA case.

brandonrising reviewed Sep 18, 2024

View reviewed changes

brandonrising approved these changes Sep 18, 2024

View reviewed changes

Merge branch 'main' into ryan/flux-lora-quantized

3d6f60f

RyanJDick enabled auto-merge September 18, 2024 17:23

RyanJDick merged commit d6da2aa into main Sep 18, 2024
14 checks passed

RyanJDick deleted the ryan/flux-lora-quantized branch September 18, 2024 17:49

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

FLUX LoRA Support #6847

FLUX LoRA Support #6847

RyanJDick commented Sep 12, 2024 •

edited

Loading

RyanJDick commented Sep 13, 2024

bghira commented Sep 14, 2024

bghira commented Sep 15, 2024

KohakuBlueleaf commented Sep 15, 2024

RyanJDick commented Sep 16, 2024

bghira commented Sep 16, 2024

FLUX LoRA Support #6847

FLUX LoRA Support #6847

Conversation

RyanJDick commented Sep 12, 2024 • edited Loading

Summary

QA Instructions

Checklist

RyanJDick commented Sep 13, 2024

bghira commented Sep 14, 2024

bghira commented Sep 15, 2024

KohakuBlueleaf commented Sep 15, 2024

RyanJDick commented Sep 16, 2024

bghira commented Sep 16, 2024

RyanJDick commented Sep 12, 2024 •

edited

Loading