Skip to content

Commit 57e19fa

Browse files
leejetleejet
authored
feat: add Flux scheduler (#1723)
1 parent 61a637b commit 57e19fa

4 files changed

Lines changed: 70 additions & 2 deletions

File tree

‎examples/common/common.cpp‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -960,7 +960,7 @@ ArgOptions SDGenerationParams::get_options() {
960960
&hires_upscaler},
961961
{"",
962962
"--extra-sample-args",
963-
"extra sampler/scheduler/guidance args, key=value list. CFG supports guidance_schedule; APG supports apg_eta, apg_momentum, apg_norm_threshold, apg_norm_threshold_smoothing; SLG supports slg_uncond; lcm supports noise_clip_std, noise_scale_start, noise_scale_end; ltx2 supports max_shift, base_shift, stretch, terminal; euler_ge supports gamma;; logit_normal supports mu, std, logsnr_min, logsnr_max, resolution_aware",
963+
"extra sampler/scheduler/guidance args, key=value list. CFG supports guidance_schedule; APG supports apg_eta, apg_momentum, apg_norm_threshold, apg_norm_threshold_smoothing; SLG supports slg_uncond; lcm supports noise_clip_std, noise_scale_start, noise_scale_end; flux supports base_shift, max_shift; ltx2 supports max_shift, base_shift, stretch, terminal; euler_ge supports gamma;; logit_normal supports mu, std, logsnr_min, logsnr_max, resolution_aware",
964964
(int)',',
965965
&extra_sample_args},
966966
{"",
@@ -1475,7 +1475,7 @@ ArgOptions SDGenerationParams::get_options() {
14751475
on_high_noise_sample_method_arg},
14761476
{"",
14771477
"--scheduler",
1478-
"denoiser sigma scheduler, one of [discrete, karras, exponential, ays, gits, smoothstep, sgm_uniform, simple, kl_optimal, lcm, bong_tangent, ltx2, logit_normal, flux2], default: model-specific",
1478+
"denoiser sigma scheduler, one of [discrete, karras, exponential, ays, gits, smoothstep, sgm_uniform, simple, kl_optimal, lcm, bong_tangent, ltx2, logit_normal, flux2, flux], default: model-specific",
14791479
on_scheduler_arg},
14801480
{"",
14811481
"--sigmas",

‎include/stable-diffusion.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,7 @@ enum scheduler_t {
7272
LTX2_SCHEDULER,
7373
LOGIT_NORMAL_SCHEDULER,
7474
FLUX2_SCHEDULER,
75+
FLUX_SCHEDULER,
7576
SCHEDULER_COUNT
7677
};
7778

‎src/runtime/denoiser.hpp‎

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -563,6 +563,65 @@ inline float flux_time_shift(float mu, float sigma, float t) {
563563
return ::expf(mu) / (::expf(mu) + ::powf((1.0f / t - 1.0f), sigma));
564564
}
565565

566+
// https://github.com/black-forest-labs/flux/blob/main/src/flux/sampling.py#L289
567+
struct FluxScheduler : SigmaScheduler {
568+
int image_seq_len = 0;
569+
float base_shift = 0.5f;
570+
float max_shift = 1.15f;
571+
572+
explicit FluxScheduler(int image_seq_len, const char* extra_sample_args = nullptr)
573+
: image_seq_len(image_seq_len) {
574+
parse_extra_sample_args(extra_sample_args);
575+
}
576+
577+
void parse_extra_sample_args(const char* extra_sample_args) {
578+
for (const auto& [key, value] : parse_key_value_args(extra_sample_args, "flux scheduler arg")) {
579+
if (key == "base_shift") {
580+
if (!parse_strict_float(value, base_shift)) {
581+
LOG_WARN("ignoring invalid flux scheduler arg '%s=%s'", key.c_str(), value.c_str());
582+
}
583+
} else if (key == "max_shift") {
584+
if (!parse_strict_float(value, max_shift)) {
585+
LOG_WARN("ignoring invalid flux scheduler arg '%s=%s'", key.c_str(), value.c_str());
586+
}
587+
}
588+
}
589+
}
590+
591+
float compute_mu() const {
592+
constexpr float base_shift_anchor = 256.0f;
593+
constexpr float max_shift_anchor = 4096.0f;
594+
float m = (max_shift - base_shift) / (max_shift_anchor - base_shift_anchor);
595+
float b = base_shift - m * base_shift_anchor;
596+
return static_cast<float>(image_seq_len) * m + b;
597+
}
598+
599+
std::vector<float> get_sigmas(uint32_t n, float /*sigma_min*/, float /*sigma_max*/, t_to_sigma_t /*t_to_sigma*/) override {
600+
std::vector<float> sigmas;
601+
sigmas.reserve(n + 1);
602+
603+
float mu = compute_mu();
604+
LOG_DEBUG("Flux scheduler: image_seq_len=%d, steps=%u, mu=%.3f", image_seq_len, n, mu);
605+
606+
if (n == 0) {
607+
sigmas.push_back(1.0f);
608+
return sigmas;
609+
}
610+
611+
for (uint32_t i = 0; i <= n; ++i) {
612+
float t = 1.0f - static_cast<float>(i) / static_cast<float>(n);
613+
if (t <= 0.0f) {
614+
sigmas.push_back(0.0f);
615+
} else {
616+
sigmas.push_back(flux_time_shift(mu, 1.0f, t));
617+
}
618+
}
619+
620+
sigmas[n] = 0.0f;
621+
return sigmas;
622+
}
623+
};
624+
566625
// https://github.com/black-forest-labs/flux2/blob/main/src/flux2/sampling.py#L244
567626
struct Flux2Scheduler : SigmaScheduler {
568627
int image_seq_len = 0;
@@ -886,6 +945,11 @@ struct Denoiser {
886945
scheduler = std::make_shared<Flux2Scheduler>(image_seq_len);
887946
break;
888947
}
948+
case FLUX_SCHEDULER: {
949+
LOG_INFO("get_sigmas with Flux scheduler");
950+
scheduler = std::make_shared<FluxScheduler>(image_seq_len, extra_sample_args);
951+
break;
952+
}
889953
default:
890954
LOG_INFO("get_sigmas with discrete scheduler (default)");
891955
scheduler = std::make_shared<DiscreteScheduler>();

‎src/stable-diffusion.cpp‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2561,6 +2561,7 @@ const char* scheduler_to_str[] = {
25612561
"ltx2",
25622562
"logit_normal",
25632563
"flux2",
2564+
"flux",
25642565
};
25652566

25662567
const char* sd_scheduler_name(enum scheduler_t scheduler) {
@@ -3161,6 +3162,8 @@ enum scheduler_t sd_get_default_scheduler(const sd_ctx_t* sd_ctx, enum sample_me
31613162
return LCM_SCHEDULER;
31623163
} else if (sample_method == DDIM_TRAILING_SAMPLE_METHOD) {
31633164
return SIMPLE_SCHEDULER;
3165+
} else if (sd_ctx != nullptr && sd_ctx->sd != nullptr && sd_version_is_flux(sd_ctx->sd->version)) {
3166+
return FLUX_SCHEDULER;
31643167
} else if (sd_ctx != nullptr && sd_ctx->sd != nullptr && sd_version_is_flux2(sd_ctx->sd->version)) {
31653168
return FLUX2_SCHEDULER;
31663169
} else if (sd_ctx != nullptr && sd_ctx->sd != nullptr && sd_version_is_ltxav(sd_ctx->sd->version)) {

0 commit comments

Comments
 (0)