upstream-archive/stable-diffusion.cpp
1
0
Fork 0
mirror of https://github.com/leejet/stable-diffusion.cpp.git synced 2026-07-23 11:20:53 -05:00
Code Issues Packages Projects Releases Wiki Activity

feat: add Flux2 scheduler (#1722)

Browse Source
This commit is contained in:
leejet
2026-06-29 23:17:31 +08:00
committed by GitHub
parent 3ec374aecc
commit 61a637bfd5
4 changed files with 71 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
examples/common/common.cpp
View File

@@ -653,7 +653,7 @@ ArgOptions SDContextParams::get_options() {
on_sampler_rng_arg},
{"",
"--prediction",
"prediction type override, one of [eps, v, edm_v, sd3_flow, flux_flow, flux2_flow]",
"prediction type override, one of [eps, v, edm_v, sd3_flow, flux_flow, sefi_flow]",
on_prediction_arg},
{"",
"--lora-apply-mode",
@@ -1475,7 +1475,7 @@ ArgOptions SDGenerationParams::get_options() {
on_high_noise_sample_method_arg},
{"",
"--scheduler",
"denoiser sigma scheduler, one of [discrete, karras, exponential, ays, gits, smoothstep, sgm_uniform, simple, kl_optimal, lcm, bong_tangent, ltx2, logit_normal], default: model-specific",
"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",
on_scheduler_arg},
{"",
"--sigmas",

2
include/stable-diffusion.h
View File

@@ -71,6 +71,7 @@ enum scheduler_t {
BONG_TANGENT_SCHEDULER,
LTX2_SCHEDULER,
LOGIT_NORMAL_SCHEDULER,
FLUX2_SCHEDULER,
SCHEDULER_COUNT
};
@@ -80,7 +81,6 @@ enum prediction_t {
EDM_V_PRED,
FLOW_PRED,
FLUX_FLOW_PRED,
FLUX2_FLOW_PRED,
SEFI_FLOW_PRED,
PREDICTION_COUNT
};

100
src/runtime/denoiser.hpp
View File

@@ -559,6 +559,63 @@ struct LTX2Scheduler : SigmaScheduler {
}
};
inline float flux_time_shift(float mu, float sigma, float t) {
return ::expf(mu) / (::expf(mu) + ::powf((1.0f / t - 1.0f), sigma));
}
// https://github.com/black-forest-labs/flux2/blob/main/src/flux2/sampling.py#L244
struct Flux2Scheduler : SigmaScheduler {
int image_seq_len = 0;
explicit Flux2Scheduler(int image_seq_len)
: image_seq_len(image_seq_len) {}
static float compute_empirical_mu(int image_seq_len, uint32_t num_steps) {
const float a1 = 8.73809524e-05f;
const float b1 = 1.89833333f;
const float a2 = 0.00016927f;
const float b2 = 0.45666666f;
if (image_seq_len > 4300) {
return a2 * image_seq_len + b2;
}
float m_200 = a2 * image_seq_len + b2;
float m_10 = a1 * image_seq_len + b1;
float a = (m_200 - m_10) / 190.0f;
float b = m_200 - 200.0f * a;
return a * num_steps + b;
}
std::vector<float> get_sigmas(uint32_t n, float /*sigma_min*/, float /*sigma_max*/, t_to_sigma_t /*t_to_sigma*/) override {
std::vector<float> sigmas;
sigmas.reserve(n + 1);
float mu = compute_empirical_mu(image_seq_len, n);
LOG_DEBUG("Flux2 scheduler: image_seq_len=%d, steps=%u, mu=%.3f", image_seq_len, n, mu);
if (n == 0) {
sigmas.push_back(1.0f);
return sigmas;
}
for (uint32_t i = 0; i <= n; ++i) {
float t = 1.0f - static_cast<float>(i) / static_cast<float>(n);
if (t <= 0.0f) {
sigmas.push_back(0.0f);
} else if (t >= 1.0f) {
sigmas.push_back(1.0f);
} else {
sigmas.push_back(flux_time_shift(mu, 1.0f, t));
}
}
sigmas[n] = 0.0f;
return sigmas;
}
};
/*
* Logit-Normal Scheduler
* Based on: https://github.com/ideogram-oss/ideogram4/blob/main/src/ideogram4/scheduler.py
@@ -824,6 +881,11 @@ struct Denoiser {
scheduler = std::make_shared<LogitNormalScheduler>(image_seq_len, extra_sample_args);
break;
}
case FLUX2_SCHEDULER: {
LOG_INFO("get_sigmas with Flux2 scheduler");
scheduler = std::make_shared<Flux2Scheduler>(image_seq_len);
break;
}
default:
LOG_INFO("get_sigmas with discrete scheduler (default)");
scheduler = std::make_shared<DiscreteScheduler>();
@@ -988,10 +1050,6 @@ struct DiscreteFlowDenoiser : public Denoiser {
}
};
inline float flux_time_shift(float mu, float sigma, float t) {
return ::expf(mu) / (::expf(mu) + ::powf((1.0f / t - 1.0f), sigma));
}
struct FluxFlowDenoiser : public DiscreteFlowDenoiser {
FluxFlowDenoiser() = default;
@@ -1007,39 +1065,7 @@ struct FluxFlowDenoiser : public DiscreteFlowDenoiser {
struct SefiFlowDenoiser;
struct Flux2FlowDenoiser : public FluxFlowDenoiser {
Flux2FlowDenoiser() = default;
float compute_empirical_mu(uint32_t n, int image_seq_len) {
const float a1 = 8.73809524e-05f;
const float b1 = 1.89833333f;
const float a2 = 0.00016927f;
const float b2 = 0.45666666f;
if (image_seq_len > 4300) {
float mu = a2 * image_seq_len + b2;
return mu;
}
float m_200 = a2 * image_seq_len + b2;
float m_10 = a1 * image_seq_len + b1;
float a = (m_200 - m_10) / 190.0f;
float b = m_200 - 200.0f * a;
float mu = a * n + b;
return mu;
}
std::vector<float> get_sigmas(uint32_t n, int image_seq_len, scheduler_t scheduler_type, SDVersion version, const char* extra_sample_args = nullptr) override {
float mu = compute_empirical_mu(n, image_seq_len);
LOG_DEBUG("Flux2FlowDenoiser: set shift to %.3f", mu);
set_shift(mu);
return Denoiser::get_sigmas(n, image_seq_len, scheduler_type, version, extra_sample_args);
}
};
struct SefiFlowDenoiser : public Flux2FlowDenoiser {
struct SefiFlowDenoiser : public FluxFlowDenoiser {
static constexpr int kNumTrainTimesteps = 1000;
static constexpr int kSemChannels = 16;
static constexpr int kTotalChannels = 144;

13
src/stable-diffusion.cpp
View File

@@ -1274,6 +1274,7 @@ public:
default_flow_shift = 3.f;
}
} else if (sd_version_is_flux(version) ||
sd_version_is_flux2(version) ||
sd_version_is_longcat(version) ||
sd_version_is_lens(version) ||
sd_version_is_ltxav(version) ||
@@ -1298,8 +1299,6 @@ public:
}
} else if (sd_version_is_sefi_image(version)) {
pred_type = SEFI_FLOW_PRED;
} else if (sd_version_is_flux2(version)) {
pred_type = FLUX2_FLOW_PRED;
} else {
pred_type = EPS_PRED;
}
@@ -1332,11 +1331,6 @@ public:
denoiser = std::make_shared<FluxFlowDenoiser>();
break;
}
case FLUX2_FLOW_PRED: {
LOG_INFO("running in Flux2 FLOW mode");
denoiser = std::make_shared<Flux2FlowDenoiser>();
break;
}
case SEFI_FLOW_PRED: {
LOG_INFO("running in SeFi-Image dual-time FLOW mode");
denoiser = std::make_shared<SefiFlowDenoiser>();
@@ -2566,6 +2560,7 @@ const char* scheduler_to_str[] = {
"bong_tangent",
"ltx2",
"logit_normal",
"flux2",
};
const char* sd_scheduler_name(enum scheduler_t scheduler) {
@@ -2590,7 +2585,7 @@ const char* prediction_to_str[] = {
"edm_v",
"sd3_flow",
"flux_flow",
"flux2_flow",
"sefi_flow",
};
const char* sd_prediction_name(enum prediction_t prediction) {
@@ -3166,6 +3161,8 @@ enum scheduler_t sd_get_default_scheduler(const sd_ctx_t* sd_ctx, enum sample_me
return LCM_SCHEDULER;
} else if (sample_method == DDIM_TRAILING_SAMPLE_METHOD) {
return SIMPLE_SCHEDULER;
} else if (sd_ctx != nullptr && sd_ctx->sd != nullptr && sd_version_is_flux2(sd_ctx->sd->version)) {
return FLUX2_SCHEDULER;
} else if (sd_ctx != nullptr && sd_ctx->sd != nullptr && sd_version_is_ltxav(sd_ctx->sd->version)) {
return LTX2_SCHEDULER;
} else if (sd_ctx != nullptr && sd_ctx->sd != nullptr && sd_version_is_ideogram4(sd_ctx->sd->version)) {