feat: add Flux2 scheduler (#1722)

src/runtime/denoiser.hpp

@@ -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;