Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix Flux schedule shift and add resolution-dependent schedule shift #905

Merged
merged 2 commits into from
Aug 31, 2024

Conversation

mhirki
Copy link
Contributor

@mhirki mhirki commented Aug 31, 2024

This is based on discussion after PR #892. This fixes --flux_schedule_shift by applying the shift to the sigmas. This makes it equivalent with Kohya's code.

This pull request also adds new option --flux_schedule_auto_shift for resolution-dependent shifting based on official Flux inference code:
https://github.com/black-forest-labs/flux/blob/b4f689aaccd40de93429865793e84a734f4a6254/src/flux/sampling.py#L66

People on Discord have been trying this yesterday already and results have been positive. Previous version of the code is here (before rebasing it due refactored train.py):
mhirki@7cafb86

I tested the code using my previous configuration from this run:
https://huggingface.co/mikaelh/flux-sanna-marin-v0.4-fp8-adan2

One of the observations from my own testing was that learning rate needed to be reduced. Initially I tried running with --flux_schedule_shift=3.1582 and lowering learning rate by 2x but that ended up overfitting so I ended up using a checkpoint from 2000 steps. Then I tried --flux_schedule_auto_shift with 5x lower learning rate and it turned out well.

--flux_schedule_shift=3.1582 with 2x lower LR @ 2000 steps --flux_schedule_auto_shift with 5x lower LR @ 3000 steps
flux_1725001441_0001 flux_1725092497_0001
flux_1725001553_0001 flux_1725092526_0001
flux_1725001583_0001 flux_1725092556_0001
flux_1725001613_0001 flux_1725092586_0001
flux_1725001643_0001 flux_1725092624_0001
flux_1725001674_0001 flux_1725092663_0001
flux_1725001704_0001 flux_1725092693_0001
flux_1725001735_0001 flux_1725092723_0001

…hedule_auto_shift for resolution-dependent shifting.
helpers/training/trainer.py Outdated Show resolved Hide resolved
@bghira bghira merged commit d43c64e into bghira:main Aug 31, 2024
1 check passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants