Krea 2 LoRA DreamBooth trainer#14046
Conversation
Adds examples/dreambooth/train_dreambooth_lora_krea2.py and the Krea2LoraLoaderMixin (Krea2Pipeline now inherits it) so LoRA adapters can be saved/loaded for the Krea 2 transformer. The transformer/pipeline forward now accept attention_kwargs for runtime LoRA scaling. Co-Authored-By: linoytsaban <linoy.tsaban@gmail.com>
|
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. |
|
Note on the latest commit: after porting the pipeline to Fix: I inlined the two small pack/unpack ops directly in the training loop ( Alternative (maybe in a separate PR): restoring |
|
@bot /style |
|
Style fix runs successfully without any file modified. |

What does this do?
Adds LoRA support for Krea 2 (now merged in #14045):
examples/dreambooth/train_dreambooth_lora_krea2.py— DreamBooth LoRA trainer with the authors' recommended defaults (rank/alpha 32, target layers, lr 3e-4 constant), fp8 support, validation, plusREADME_krea2.mdandrequirements_krea2.txt.Krea2LoraLoaderMixininloaders/lora_pipeline.py;Krea2Pipelinenow inherits it so LoRA adapters can be saved/loaded.Krea2Transformer2DModel.forwardand the pipeline__call__acceptattention_kwargsfor runtime LoRA scaling.tests/lora/test_lora_layers_krea2.py.Ported from the Krea 2 staging repo and rebased onto current
main.