Skip to content

Krea 2 LoRA DreamBooth trainer#14046

Merged
DN6 merged 5 commits into
mainfrom
krea2-lora
Jun 23, 2026
Merged

Krea 2 LoRA DreamBooth trainer#14046
DN6 merged 5 commits into
mainfrom
krea2-lora

Conversation

@apolinario

Copy link
Copy Markdown
Collaborator

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, plus README_krea2.md and requirements_krea2.txt.
  • Krea2LoraLoaderMixin in loaders/lora_pipeline.py; Krea2Pipeline now inherits it so LoRA adapters can be saved/loaded.
  • Krea2Transformer2DModel.forward and the pipeline __call__ accept attention_kwargs for runtime LoRA scaling.
  • LoRA loader tests in tests/lora/test_lora_layers_krea2.py.

Ported from the Krea 2 staging repo and rebased onto current main.

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

Copy link
Copy Markdown

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.

@yiyixuxu yiyixuxu requested a review from sayakpaul June 23, 2026 00:11
@linoytsaban

Copy link
Copy Markdown
Collaborator
image

@linoytsaban

Copy link
Copy Markdown
Collaborator

Note on the latest commit: after porting the pipeline to main, training crashed because the trainer calls Krea2Pipeline._pack_latents(...) and Krea2Pipeline._unpack_latents(...) at the class level, and these are now instance methods (they use self.patch_size / self.vae_scale_factor) - for context they were @staticmethod before the merge iirc. So the first positional argument bound to self, shifting the rest - _pack_latents raised TypeError

Fix: I inlined the two small pack/unpack ops directly in the training loop (patch_size=2), to work around it.

Alternative (maybe in a separate PR): restoring @staticmethod on _pack_latents/_unpack_latents (as in QwenImagePipeline, which they're # Copied from).

@DN6

DN6 commented Jun 23, 2026

Copy link
Copy Markdown
Collaborator

@bot /style

@github-actions github-actions Bot added the documentation Improvements or additions to documentation label Jun 23, 2026
@github-actions

github-actions Bot commented Jun 23, 2026

Copy link
Copy Markdown
Contributor

Style fix runs successfully without any file modified.

@DN6 DN6 merged commit 3993de5 into main Jun 23, 2026
36 of 37 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

documentation Improvements or additions to documentation examples loaders lora models pipelines size/L PR with diff > 200 LOC tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants