[WIP][Training] Flux Control LoRA training script #10130

sayakpaul · 2024-12-05T13:39:36Z

Will be updated later.

This PR has both LoRA and a non-LoRA version for training Control Flux.

~~Number of changed files are more than expected because the flux-control-lora branch in merged into this PR branch.~~

…x rank calculation assumption

Co-authored-by: Sayak Paul <[email protected]>

This reverts commit 73cfc51.

a-r-r-o-w · 2024-12-11T07:35:53Z

@sayakpaul I think some things like the model card are yet to be updated, and there's probably more things on your mind. Once you mark this ready for review, I can take a deeper look - everything looks great already though! Thanks for involving me in the discussions and looking forward to starting my own training run soon!

sayakpaul · 2024-12-11T07:40:11Z

I think some things like the model card are yet to be updated

I think I have the model cards already updated. But let me know if you think some things should be added/updated.

Marking this PR as ready to review.

sayakpaul · 2024-12-11T11:08:30Z

@apolinario @linoytsaban do you think it could make sense to show users how to derive a LoRA from the full Control fine-tune like how the BFL folks did?

a-r-r-o-w

Awesome, LGTM! Just some minor comments

examples/control-lora/README.md

a-r-r-o-w · 2024-12-11T17:18:04Z

examples/control-lora/train_control_flux.py

+    return pixel_latents.to(weight_dtype)
+
+
+def log_validation(flux_transformer, args, accelerator, weight_dtype, step, is_final_validation=False):


Very clean!

a-r-r-o-w · 2024-12-11T17:18:47Z

examples/control-lora/train_control_flux.py

+# controlnet-{repo_id}
+
+These are Control weights trained on {base_model} with new type of conditioning.
+{img_str}


Oh, this is what I missed earlier when I though model card was incomplete. The validation images will be displayed, awesome!

examples/control-lora/train_control_flux.py

a-r-r-o-w · 2024-12-11T17:23:33Z

examples/control-lora/train_control_flux.py

+                sigmas = get_sigmas(timesteps, n_dim=pixel_latents.ndim, dtype=pixel_latents.dtype)
+                noisy_model_input = (1.0 - sigmas) * pixel_latents + sigmas * noise
+                # Concatenate across channels.
+                # Question: Should we concatenate before adding noise?


I think what you have currently is correct. Noisy latent + clean control latent

examples/control-lora/train_control_lora_flux.py

a-r-r-o-w · 2024-12-11T17:32:37Z

examples/control-lora/train_control_lora_flux.py

+            " Please use `huggingface-cli login` to authenticate with the Hub."
+        )
+    if args.use_lora_bias and args.gaussian_init_lora:
+        raise ValueError("`gaussian` LoRA init scheme isn't supported when `use_lora_bias` is True.")


Just wondering why this is the case. We could set initialization to gaussian before the lora bias changes IIRC, no?

So, we do:

LoraConfig(lora_bias=True, init_lora_weights="gaussian", ...)

it errors out. Maybe a question for @BenjaminBossan.

As discussed earlier, it just hasn't been implemented yet. The question is: If a user wants Gaussian init for the LoRA weights, how do we initialize the LoRA bias? If we have an answer to that, we can enable Gaussian init for lora_bias.

examples/control-lora/train_control_lora_flux.py

Co-authored-by: Aryan <[email protected]>

sayakpaul · 2024-12-12T06:33:41Z

@a-r-r-o-w I have addressed your comments, thank you!

If you could review the changes and the README mods (made in 02016ca), that would be helpful!

a-r-r-o-w

Let's get the GPUs burning! LoRA testing scripts can be added in a follow-up PR if required but not very important

sayakpaul · 2024-12-12T09:51:50Z

Thanks, I think also okay to not add tests just yet. Based on the usage, we can always revisit.

Adenialzz · 2024-12-15T12:42:02Z

Hi, thanks for your implementation.

I'm curious why we set model offloading as default and don't have the option to turn it off. In my testing, this can significantly reduce training speed ( ~5x in my case). I recommend making model offloading optional in args

sayakpaul · 2024-12-15T13:46:17Z

@Adenialzz, thanks for your comment. Feel free to open a PR :)

Adenialzz · 2024-12-15T14:34:48Z

Sure. Here #10225 .

* update * add * update * add control-lora conversion script; make flux loader handle norms; fix rank calculation assumption * control lora updates * remove copied-from * create separate pipelines for flux control * make fix-copies * update docs * add tests * fix * Apply suggestions from code review Co-authored-by: Sayak Paul <[email protected]> * remove control lora changes * apply suggestions from review * Revert "remove control lora changes" This reverts commit 73cfc51. * update * update * improve log messages * updates. * updates * support register_config. * fix * fix * fix * updates * updates * updates * fix-copies * fix * apply suggestions from review * add tests * remove conversion script; enable on-the-fly conversion * bias -> lora_bias. * fix-copies * peft.py * fix lora conversion * changes Co-authored-by: a-r-r-o-w <[email protected]> * fix-copies * updates for tests * fix * alpha_pattern. * add a test for varied lora ranks and alphas. * revert changes in num_channels_latents = self.transformer.config.in_channels // 8 * revert moe * add a sanity check on unexpected keys when loading norm layers. * contro lora. * fixes * fixes * fixes * tests * reviewer feedback * fix * proper peft version for lora_bias * fix-copies * updates * updates * updates * remove debug code * update docs * integration tests * nis * fuse and unload. * fix * add slices. * more updates. * button up readme * train() * add full fine-tuning version. * fixes * Apply suggestions from code review Co-authored-by: Aryan <[email protected]> * set_grads_to_none remove. * readme --------- Co-authored-by: Aryan <[email protected]> Co-authored-by: yiyixuxu <[email protected]> Co-authored-by: a-r-r-o-w <[email protected]>

a-r-r-o-w and others added 30 commits November 21, 2024 06:53

update

2829679

Merge branch 'main' into flux-new

be67dbd

add

f56ffb1

update

7e4df06

Merge remote-tracking branch 'origin/flux-fill-yiyi' into flux-new

9ea52da

add control-lora conversion script; make flux loader handle norms; fi…

217e90c

…x rank calculation assumption

control lora updates

b4f1cbf

remove copied-from

414b30b

create separate pipelines for flux control

6b02ac2

make fix-copies

3169bf5

update docs

f7f006d

add tests

8bb940e

fix

9e615fd

Merge branch 'main' into flux-new

6d168db

Apply suggestions from code review

89fd970

Co-authored-by: Sayak Paul <[email protected]>

remove control lora changes

73cfc51

apply suggestions from review

c94966f

Revert "remove control lora changes"

cfe13e7

This reverts commit 73cfc51.

update

0c959a7

update

6ef2c8b

improve log messages

42970ee

Merge branch 'main' into flux-control-lora

2ec93ba

Merge branch 'main' into flux-control-lora

993f3d3

updates.

6523fa6

updates

81ab40b

Merge branch 'main' into flux-control-lora

4432e73

Merge branch 'main' into flux-control-lora

0f747c0

Merge branch 'flux-control-lora' into sayak-flux-control-lora

6d0c6dc

support register_config.

1633619

fix

b9039b1

sayakpaul added 7 commits December 7, 2024 10:43

add slices.

0b83deb

Merge branch 'flux-control-lora' into flux-control-lora-training-script

4d10b33

more updates.

b06376f

button up readme

7d63a2a

train()

9d90c51

add full fine-tuning version.

f46330b

fixes

f188e80

yiyixuxu requested a review from a-r-r-o-w December 10, 2024 18:33

Merge branch 'main' into flux-control-lora-training-script

6c1a60f

sayakpaul marked this pull request as ready for review December 11, 2024 07:40

Merge branch 'main' into flux-control-lora-training-script

2610e6a

a-r-r-o-w approved these changes Dec 11, 2024

View reviewed changes

Merge branch 'main' into flux-control-lora-training-script

ecc16ed

sayakpaul commented Dec 12, 2024

View reviewed changes

examples/control-lora/train_control_lora_flux.py Outdated Show resolved Hide resolved

sayakpaul and others added 3 commits December 12, 2024 11:25

Apply suggestions from code review

7dfe378

Co-authored-by: Aryan <[email protected]>

set_grads_to_none remove.

f1d9550

readme

02016ca

a-r-r-o-w approved these changes Dec 12, 2024

View reviewed changes

Merge branch 'main' into flux-control-lora-training-script

64814c9

sayakpaul merged commit 8170dc3 into main Dec 12, 2024
12 checks passed

sayakpaul deleted the flux-control-lora-training-script branch December 12, 2024 10:05

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[WIP][Training] Flux Control LoRA training script #10130

[WIP][Training] Flux Control LoRA training script #10130

sayakpaul commented Dec 5, 2024 •

edited

Loading

a-r-r-o-w commented Dec 11, 2024

sayakpaul commented Dec 11, 2024

sayakpaul commented Dec 11, 2024

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

a-r-r-o-w Dec 11, 2024

a-r-r-o-w Dec 11, 2024

a-r-r-o-w Dec 11, 2024

a-r-r-o-w Dec 11, 2024

sayakpaul Dec 12, 2024

BenjaminBossan Dec 16, 2024

sayakpaul commented Dec 12, 2024

a-r-r-o-w left a comment •

edited

Loading

sayakpaul commented Dec 12, 2024

Adenialzz commented Dec 15, 2024 •

edited

Loading

sayakpaul commented Dec 15, 2024

Adenialzz commented Dec 15, 2024

		return pixel_latents.to(weight_dtype)


		def log_validation(flux_transformer, args, accelerator, weight_dtype, step, is_final_validation=False):

[WIP][Training] Flux Control LoRA training script #10130

[WIP][Training] Flux Control LoRA training script #10130

Conversation

sayakpaul commented Dec 5, 2024 • edited Loading

a-r-r-o-w commented Dec 11, 2024

sayakpaul commented Dec 11, 2024

sayakpaul commented Dec 11, 2024

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

Choose a reason for hiding this comment

a-r-r-o-w Dec 11, 2024

Choose a reason for hiding this comment

a-r-r-o-w Dec 11, 2024

Choose a reason for hiding this comment

a-r-r-o-w Dec 11, 2024

Choose a reason for hiding this comment

a-r-r-o-w Dec 11, 2024

Choose a reason for hiding this comment

sayakpaul Dec 12, 2024

Choose a reason for hiding this comment

BenjaminBossan Dec 16, 2024

Choose a reason for hiding this comment

sayakpaul commented Dec 12, 2024

a-r-r-o-w left a comment • edited Loading

Choose a reason for hiding this comment

sayakpaul commented Dec 12, 2024

Adenialzz commented Dec 15, 2024 • edited Loading

sayakpaul commented Dec 15, 2024

Adenialzz commented Dec 15, 2024

sayakpaul commented Dec 5, 2024 •

edited

Loading

a-r-r-o-w left a comment •

edited

Loading

Adenialzz commented Dec 15, 2024 •

edited

Loading