[Flux Dreambooth lora] add latent caching #9160

linoytsaban · 2024-08-12T14:56:17Z

adds latent caching and relevant tests
changes default upcasting of transformer trained layers at the end of training

HuggingFaceDocBuilderDev · 2024-08-12T15:02:35Z

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.

sayakpaul

LGTM, let's add a test too and a note in the README?

Also, how much memory does it save? Do we have a ballpark?

examples/dreambooth/README_flux.md

sayakpaul · 2024-08-13T09:59:36Z

examples/dreambooth/train_dreambooth_lora_flux.py

+            del vae
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+                gc.collect()


This doesn't have to be conditioned on the availability of CUDA, no?

hmm I think maybe not, but for some reason we've used this condition in most other places too

it should be - there's one for mps to call too. i think there should be a utility helper for it

Good point. Let's have it as is for now. I am working on a small utility for cleaning models and retaining accelerator memory.

@linoytsaban possible to use?

diffusers/src/diffusers/training_utils.py

Line 263 in 48e3635

def clear_objs_and_retain_memory(objs: List[Any]):

https://github.com/ostris/ai-toolkit/blob/9ee1ef2a0a2a9a02b92d114a95f21312e5906e54/toolkit/samplers/custom_flowmatch_sampler.py#L95

linoytsaban · 2024-08-21T12:45:20Z

@sayakpaul wdyt about the change I made for the upcasting in the end?
it makes a significant different when training on A100 and quality is still good, I tested it with my yarn art lora

bghira · 2024-08-21T12:54:29Z

saving weights in fp32 just feels unnecessary with Flux. most people are sharing fp8 weights around instead of even bf16, but simpletuner saves its weights in the dtype it was trained in, which users expected

sayakpaul

@linoytsaban thanks for the changes.

My major comment is on reusing stuff from the custom scheduler as necessary instead of copy-pasting it fully. LMK what you think.

sayakpaul · 2024-08-22T01:27:13Z

examples/dreambooth/train_dreambooth_lora_flux.py

+# CustomFlowMatchEulerDiscreteScheduler was taken from ostris ai-toolkit trainer:
+# https://github.com/ostris/ai-toolkit/blob/9ee1ef2a0a2a9a02b92d114a95f21312e5906e54/toolkit/samplers/custom_flowmatch_sampler.py#L95
+class CustomFlowMatchEulerDiscreteScheduler(FlowMatchEulerDiscreteScheduler):


@linoytsaban I would just re-use the relevant parts from the original CustomFlowMatchEulerDiscreteScheduler here rather copy-pasting it entirely. For example, we don't need the get_simas() method 'cause we already have one inside the script.

Also option to use this scheduler related changes should be made configurable IMO.

sayakpaul · 2024-08-22T01:28:42Z

examples/dreambooth/train_dreambooth_lora_flux.py

+            del vae
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+                gc.collect()


Good point. Let's have it as is for now. I am working on a small utility for cleaning models and retaining accelerator memory.

examples/dreambooth/README_flux.md

sayakpaul · 2024-08-22T01:30:44Z

@sayakpaul wdyt about the change I made for the upcasting in the end? it makes a significant different when training on A100 and quality is still good, I tested it with my yarn art lora

Perhaps we could make this configurable and note from the README?

…efore_saving

…ce#9160

examples/dreambooth/README_flux.md

sayakpaul · 2024-09-14T07:48:53Z

examples/dreambooth/test_dreambooth_lora_flux.py

            )
            self.assertTrue(starts_with_expected_prefix)

+    def test_dreambooth_lora_latent_caching(self):


sayakpaul

Excellent work! Very minor suggestions.

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>

sayakpaul

LGTM, thanks for this work, @linoytsaban!

sayakpaul · 2024-09-15T06:31:43Z

@linoytsaban feel free to merge after the tests pass.

* add ostris trainer to README & add cache latents of vae * add ostris trainer to README & add cache latents of vae * style * readme * add test for latent caching * add ostris noise scheduler https://github.com/ostris/ai-toolkit/blob/9ee1ef2a0a2a9a02b92d114a95f21312e5906e54/toolkit/samplers/custom_flowmatch_sampler.py#L95 * style * fix import * style * fix tests * style * --change upcasting of transformer? * update readme according to main * keep only latent caching * add configurable param for final saving of trained layers- --upcast_before_saving * style * Update examples/dreambooth/README_flux.md Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> * Update examples/dreambooth/README_flux.md Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> * use clear_objs_and_retain_memory from utilities * style --------- Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>

…ence (#9434) * add ostris trainer to README & add cache latents of vae * add ostris trainer to README & add cache latents of vae * style * readme * add test for latent caching * add ostris noise scheduler https://github.com/ostris/ai-toolkit/blob/9ee1ef2a0a2a9a02b92d114a95f21312e5906e54/toolkit/samplers/custom_flowmatch_sampler.py#L95 * style * fix import * style * fix tests * style * --change upcasting of transformer? * update readme according to main * add pivotal tuning for CLIP * fix imports, encode_prompt call,add TextualInversionLoaderMixin to FluxPipeline for inference * TextualInversionLoaderMixin support for FluxPipeline for inference * move changes to advanced flux script, revert canonical * add latent caching to canonical script * revert changes to canonical script to keep it separate from #9160 * revert changes to canonical script to keep it separate from #9160 * style * remove redundant line and change code block placement to align with logic * add initializer_token arg * add transformer frac for range support from pure textual inversion to the orig pivotal tuning * support pure textual inversion - wip * adjustments to support pure textual inversion and transformer optimization in only part of the epochs * fix logic when using initializer token * fix pure_textual_inversion_condition * fix ti/pivotal loading of last validation run * remove embeddings loading for ti in final training run (to avoid adding huggingface hub dependency) * support pivotal for t5 * adapt pivotal for T5 encoder * adapt pivotal for T5 encoder and support in flux pipeline * t5 pivotal support + support fo pivotal for clip only or both * fix param chaining * fix param chaining * README first draft * readme * readme * readme * style * fix import * style * add fix from #9419 * add to readme, change function names * te lr changes * readme * change concept tokens logic * fix indices * change arg name * style * dummy test * revert dummy test * reorder pivoting * add warning in case the token abstraction is not the instance prompt * experimental - wip - specific block training * fix documentation and token abstraction processing * remove transformer block specification feature (for now) * style * fix copies * fix indexing issue when --initializer_concept has different amounts * add if TextualInversionLoaderMixin to all flux pipelines * style * fix import * fix imports * address review comments - remove necessary prints & comments, use pin_memory=True, use free_memory utils, unify warning and prints * style * logger info fix * make lora target modules configurable and change the default * make lora target modules configurable and change the default * style * make lora target modules configurable and change the default, add notes to readme * style * add tests * style * fix repo id * add updated requirements for advanced flux * fix indices of t5 pivotal tuning embeddings * fix path in test * remove `pin_memory` * fix filename of embedding * fix filename of embedding --------- Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> Co-authored-by: YiYi Xu <yixu310@gmail.com>

* add ostris trainer to README & add cache latents of vae * add ostris trainer to README & add cache latents of vae * style * readme * add test for latent caching * add ostris noise scheduler https://github.com/ostris/ai-toolkit/blob/9ee1ef2a0a2a9a02b92d114a95f21312e5906e54/toolkit/samplers/custom_flowmatch_sampler.py#L95 * style * fix import * style * fix tests * style * --change upcasting of transformer? * update readme according to main * keep only latent caching * add configurable param for final saving of trained layers- --upcast_before_saving * style * Update examples/dreambooth/README_flux.md Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> * Update examples/dreambooth/README_flux.md Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> * use clear_objs_and_retain_memory from utilities * style --------- Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>

…ence (#9434) * add ostris trainer to README & add cache latents of vae * add ostris trainer to README & add cache latents of vae * style * readme * add test for latent caching * add ostris noise scheduler https://github.com/ostris/ai-toolkit/blob/9ee1ef2a0a2a9a02b92d114a95f21312e5906e54/toolkit/samplers/custom_flowmatch_sampler.py#L95 * style * fix import * style * fix tests * style * --change upcasting of transformer? * update readme according to main * add pivotal tuning for CLIP * fix imports, encode_prompt call,add TextualInversionLoaderMixin to FluxPipeline for inference * TextualInversionLoaderMixin support for FluxPipeline for inference * move changes to advanced flux script, revert canonical * add latent caching to canonical script * revert changes to canonical script to keep it separate from #9160 * revert changes to canonical script to keep it separate from #9160 * style * remove redundant line and change code block placement to align with logic * add initializer_token arg * add transformer frac for range support from pure textual inversion to the orig pivotal tuning * support pure textual inversion - wip * adjustments to support pure textual inversion and transformer optimization in only part of the epochs * fix logic when using initializer token * fix pure_textual_inversion_condition * fix ti/pivotal loading of last validation run * remove embeddings loading for ti in final training run (to avoid adding huggingface hub dependency) * support pivotal for t5 * adapt pivotal for T5 encoder * adapt pivotal for T5 encoder and support in flux pipeline * t5 pivotal support + support fo pivotal for clip only or both * fix param chaining * fix param chaining * README first draft * readme * readme * readme * style * fix import * style * add fix from #9419 * add to readme, change function names * te lr changes * readme * change concept tokens logic * fix indices * change arg name * style * dummy test * revert dummy test * reorder pivoting * add warning in case the token abstraction is not the instance prompt * experimental - wip - specific block training * fix documentation and token abstraction processing * remove transformer block specification feature (for now) * style * fix copies * fix indexing issue when --initializer_concept has different amounts * add if TextualInversionLoaderMixin to all flux pipelines * style * fix import * fix imports * address review comments - remove necessary prints & comments, use pin_memory=True, use free_memory utils, unify warning and prints * style * logger info fix * make lora target modules configurable and change the default * make lora target modules configurable and change the default * style * make lora target modules configurable and change the default, add notes to readme * style * add tests * style * fix repo id * add updated requirements for advanced flux * fix indices of t5 pivotal tuning embeddings * fix path in test * remove `pin_memory` * fix filename of embedding * fix filename of embedding --------- Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> Co-authored-by: YiYi Xu <yixu310@gmail.com>

linoytsaban added 3 commits August 12, 2024 17:30

add ostris trainer to README & add cache latents of vae

90686c2

add ostris trainer to README & add cache latents of vae

7b12ed2

style

17dca18

linoytsaban marked this pull request as ready for review August 12, 2024 15:28

Merge branch 'main' into dreambooth-lora

de24a4f

linoytsaban requested a review from sayakpaul August 13, 2024 07:06

sayakpaul reviewed Aug 13, 2024

View reviewed changes

linoytsaban and others added 7 commits August 13, 2024 17:10

readme

8b314e9

Merge branch 'main' into dreambooth-lora

a59b063

add test for latent caching

df54cd8

add ostris noise scheduler

e0e0319

https://github.com/ostris/ai-toolkit/blob/9ee1ef2a0a2a9a02b92d114a95f21312e5906e54/toolkit/samplers/custom_flowmatch_sampler.py#L95

style

18aa369

fix import

f97d53d

style

0156bec

linoytsaban changed the title ~~[Flux Dreambooth lora] add to readme + latent caching~~ [Flux Dreambooth lora] add custom scheduler + latent caching Aug 14, 2024

linoytsaban and others added 5 commits August 14, 2024 11:59

fix tests

c4c2c48

style

d514c7b

Merge branch 'main' into dreambooth-lora

7ee6041

--change upcasting of transformer?

d5c2a36

Merge branch 'main' into dreambooth-lora

e760cda

Merge branch 'main' into dreambooth-lora

f78ba77

sayakpaul reviewed Aug 22, 2024

View reviewed changes

sayakpaul and others added 4 commits August 22, 2024 07:00

Merge branch 'main' into dreambooth-lora

1b19593

update readme according to main

fbacbb5

Merge branch 'main' into dreambooth-lora

23f0636

Merge branch 'main' into dreambooth-lora

51c7667

linoytsaban mentioned this pull request Sep 13, 2024

[Flux] Add advanced training script + support textual inversion inference #9434

Merged

keep only latent caching

feae3dc

linoytsaban changed the title ~~[Flux Dreambooth lora] add custom scheduler + latent caching~~ [Flux Dreambooth lora] add latent caching Sep 13, 2024

linoytsaban and others added 3 commits September 13, 2024 18:42

add configurable param for final saving of trained layers- --upcast_b…

b53ae0b

…efore_saving

Merge branch 'main' into dreambooth-lora

79e5234

style

5cdb4f5

linoytsaban requested a review from sayakpaul September 13, 2024 16:47

linoytsaban added a commit to linoytsaban/diffusers that referenced this pull request Sep 13, 2024

revert changes to canonical script to keep it separate from huggingfa…

2bb4ce1

…ce#9160

linoytsaban added a commit to linoytsaban/diffusers that referenced this pull request Sep 13, 2024

revert changes to canonical script to keep it separate from huggingfa…

dc9be5b

…ce#9160

sayakpaul reviewed Sep 14, 2024

View reviewed changes

examples/dreambooth/README_flux.md Outdated Show resolved Hide resolved

sayakpaul reviewed Sep 14, 2024

View reviewed changes

examples/dreambooth/README_flux.md Outdated Show resolved Hide resolved

sayakpaul reviewed Sep 14, 2024

View reviewed changes

sayakpaul approved these changes Sep 14, 2024

View reviewed changes

linoytsaban and others added 4 commits September 14, 2024 20:57

Update examples/dreambooth/README_flux.md

e047ae2

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>

Update examples/dreambooth/README_flux.md

a882c41

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>

use clear_objs_and_retain_memory from utilities

75058d7

Merge branch 'main' into dreambooth-lora

d61868e

sayakpaul approved these changes Sep 15, 2024

View reviewed changes

style

88c0275

linoytsaban merged commit 37e3603 into huggingface:main Sep 15, 2024
8 checks passed

linoytsaban deleted the dreambooth-lora branch September 15, 2024 15:41

[Flux Dreambooth lora] add latent caching #9160

[Flux Dreambooth lora] add latent caching #9160

Uh oh!

Conversation

linoytsaban commented Aug 12, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

HuggingFaceDocBuilderDev commented Aug 12, 2024

Uh oh!

sayakpaul left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

linoytsaban commented Aug 21, 2024

Uh oh!

bghira commented Aug 21, 2024

Uh oh!

sayakpaul left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

sayakpaul commented Aug 22, 2024

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sayakpaul left a comment

Choose a reason for hiding this comment

Uh oh!

sayakpaul left a comment

Choose a reason for hiding this comment

Uh oh!

sayakpaul commented Sep 15, 2024

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

linoytsaban commented Aug 12, 2024 •

edited

Loading