diff --git a/src/runtime/denoiser.hpp b/src/runtime/denoiser.hpp index 416859e6..15c12608 100644 --- a/src/runtime/denoiser.hpp +++ b/src/runtime/denoiser.hpp @@ -728,12 +728,20 @@ inline float flux_time_shift(float mu, float sigma, float t) { // https://github.com/black-forest-labs/flux/blob/main/src/flux/sampling.py#L289 struct FluxScheduler : SigmaScheduler { - int image_seq_len = 0; - float base_shift = 0.5f; - float max_shift = 1.15f; + int image_seq_len = 0; + int base_image_seq_len = 256; + int max_image_seq_len = 4096; + float base_shift = 0.5f; + float max_shift = 1.15f; + float shift_terminal = 0.0f; - explicit FluxScheduler(int image_seq_len, const char* extra_sample_args = nullptr) + FluxScheduler(int image_seq_len, SDVersion version, const char* extra_sample_args = nullptr) : image_seq_len(image_seq_len) { + if (version == VERSION_QWEN_IMAGE_2_1) { + max_image_seq_len = 8192; + max_shift = 0.9f; + shift_terminal = 0.02f; + } parse_extra_sample_args(extra_sample_args); } @@ -752,10 +760,8 @@ struct FluxScheduler : SigmaScheduler { } float compute_mu() const { - constexpr float base_shift_anchor = 256.0f; - constexpr float max_shift_anchor = 4096.0f; - float m = (max_shift - base_shift) / (max_shift_anchor - base_shift_anchor); - float b = base_shift - m * base_shift_anchor; + float m = (max_shift - base_shift) / static_cast(max_image_seq_len - base_image_seq_len); + float b = base_shift - m * static_cast(base_image_seq_len); return static_cast(image_seq_len) * m + b; } @@ -764,7 +770,7 @@ struct FluxScheduler : SigmaScheduler { sigmas.reserve(n + 1); float mu = compute_mu(); - LOG_VERBOSE("Flux scheduler: image_seq_len=%d, steps=%u, mu=%.3f", image_seq_len, n, mu); + LOG_VERBOSE("Flux scheduler: image_seq_len=%d, steps=%u, mu=%.3f, shift_terminal=%.3f", image_seq_len, n, mu, shift_terminal); if (n == 0) { sigmas.push_back(1.0f); @@ -780,6 +786,16 @@ struct FluxScheduler : SigmaScheduler { } } + if (shift_terminal > 0.0f && n > 1) { + // The terminal shift applies to the last model evaluation, not the final zero sigma. + float scale_factor = (1.0f - sigmas[n - 1]) / (1.0f - shift_terminal); + if (std::isfinite(scale_factor) && scale_factor > 0.0f) { + for (uint32_t i = 0; i < n; ++i) { + sigmas[i] = 1.0f - (1.0f - sigmas[i]) / scale_factor; + } + } + } + sigmas[n] = 0.0f; return sigmas; } @@ -1178,7 +1194,7 @@ struct Denoiser { } case FLUX_SCHEDULER: { LOG_INFO("get_sigmas with Flux scheduler"); - scheduler = std::make_shared(image_seq_len, extra_sample_args); + scheduler = std::make_shared(image_seq_len, version, extra_sample_args); break; } default: