feat: add normal alias for discrete scheduler (#1724)

examples/cli/CMakeLists.txt

@@ -15,7 +15,7 @@ target_include_directories(${TARGET} PRIVATE
     "${PROJECT_SOURCE_DIR}/src"
 )
 install(TARGETS ${TARGET} RUNTIME)
-target_link_libraries(${TARGET} PRIVATE stable-diffusion ${CMAKE_THREAD_LIBS_INIT})
+target_link_libraries(${TARGET} PRIVATE stable-diffusion zip ${CMAKE_THREAD_LIBS_INIT})
 if(SD_WEBP)
     target_compile_definitions(${TARGET} PRIVATE SD_USE_WEBP)
     target_link_libraries(${TARGET} PRIVATE webp libwebpmux)

examples/common/common.cpp

@@ -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",
@@ -960,7 +960,7 @@ ArgOptions SDGenerationParams::get_options() {
          &hires_upscaler},
         {"",
          "--extra-sample-args",
-         "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",
+         "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",
          (int)',',
          &extra_sample_args},
         {"",
@@ -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, flux], alias: normal=discrete, default: model-specific",
          on_scheduler_arg},
         {"",
          "--sigmas",

examples/server/routes_sdapi.cpp

@@ -438,6 +438,9 @@ void register_sdapi_endpoints(httplib::Server& svr, ServerRuntime& rt) {
         scheduler_names.push_back("default");
         for (int i = 0; i < SCHEDULER_COUNT; i++) {
             scheduler_names.push_back(sd_scheduler_name((scheduler_t)i));
+            if (i == DISCRETE_SCHEDULER) {
+                scheduler_names.push_back("normal");
+            }
         }
         json r = json::array();
         for (auto name : scheduler_names) {

examples/server/routes_sdcpp.cpp

@@ -219,6 +219,9 @@ static json make_capabilities_json(ServerRuntime& runtime) {
 
     for (int i = 0; i < SCHEDULER_COUNT; ++i) {
         schedulers.push_back(sd_scheduler_name((scheduler_t)i));
+        if (i == DISCRETE_SCHEDULER) {
+            schedulers.push_back("normal");
+        }
     }
 
     {

include/stable-diffusion.h

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

src/model/vae/wan_vae.hpp

@@ -1119,6 +1119,10 @@ namespace WAN {
             GGML_ASSERT(b == 1);
             GGML_ASSERT(decode_only == false);
 
+            if (x->ne[2] > 1 && is_2D) {
+                LOG_WARN("Using 2D VAE to encode video, expect poor results");
+            }
+
             clear_cache();
 
             if (wan2_2) {
@@ -1138,7 +1142,8 @@ namespace WAN {
                     auto in = ggml_ext_slice(ctx->ggml_ctx, x, 2, 0, 1);  // [b*c, 1, h, w]
                     out     = encoder->forward(ctx, in, b, _enc_feat_map, _enc_conv_idx, i);
                 } else {
-                    auto in   = ggml_ext_slice(ctx->ggml_ctx, x, 2, 1 + 4 * (i - 1), 1 + 4 * i);  // [b*c, 4, h, w]
+                    // if is_2D, drop 3 out of 4 frames
+                    auto in   = ggml_ext_slice(ctx->ggml_ctx, x, 2, 1 + 4 * (i - 1), (is_2D ? 1 - 3 : 1) + 4 * i);  // [b*c, 4, h, w]
                     auto out_ = encoder->forward(ctx, in, b, _enc_feat_map, _enc_conv_idx, i);
                     out       = ggml_concat(ctx->ggml_ctx, out, out_, 2);
                 }
@@ -1161,6 +1166,10 @@ namespace WAN {
             // z: [b*c, t, h, w]
             GGML_ASSERT(b == 1);
 
+            if (z->ne[2] > 1 && is_2D) {
+                LOG_WARN("Using 2D VAE to decode video, expect poor results");
+            }
+
             clear_cache();
 
             auto decoder = std::dynamic_pointer_cast<Decoder3d>(blocks["decoder"]);
@@ -1185,6 +1194,12 @@ namespace WAN {
                     auto in   = ggml_ext_slice(ctx->ggml_ctx, x, 2, i, i + 1);  // [b*c, 1, h, w]
                     auto out_ = decoder->forward(ctx, in, b, _feat_map, _conv_idx, i);
                     out       = ggml_concat(ctx->ggml_ctx, out, out_, 2);
+                    if (is_2D) {
+                        // repeat frames to avoid mismatch
+                        for (int j = 0; j < 4 - 1; j++) {
+                            out = ggml_concat(ctx->ggml_ctx, out, out_, 2);
+                        }
+                    }
                 }
             }
             if (wan2_2) {

src/runtime/denoiser.hpp

@@ -559,6 +559,122 @@ 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/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;
+
+    explicit FluxScheduler(int image_seq_len, const char* extra_sample_args = nullptr)
+        : image_seq_len(image_seq_len) {
+        parse_extra_sample_args(extra_sample_args);
+    }
+
+    void parse_extra_sample_args(const char* extra_sample_args) {
+        for (const auto& [key, value] : parse_key_value_args(extra_sample_args, "flux scheduler arg")) {
+            if (key == "base_shift") {
+                if (!parse_strict_float(value, base_shift)) {
+                    LOG_WARN("ignoring invalid flux scheduler arg '%s=%s'", key.c_str(), value.c_str());
+                }
+            } else if (key == "max_shift") {
+                if (!parse_strict_float(value, max_shift)) {
+                    LOG_WARN("ignoring invalid flux scheduler arg '%s=%s'", key.c_str(), value.c_str());
+                }
+            }
+        }
+    }
+
+    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;
+        return static_cast<float>(image_seq_len) * m + 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_mu();
+        LOG_DEBUG("Flux 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 {
+                sigmas.push_back(flux_time_shift(mu, 1.0f, t));
+            }
+        }
+
+        sigmas[n] = 0.0f;
+        return sigmas;
+    }
+};
+
+// 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 +940,16 @@ 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;
+            }
+            case FLUX_SCHEDULER: {
+                LOG_INFO("get_sigmas with Flux scheduler");
+                scheduler = std::make_shared<FluxScheduler>(image_seq_len, extra_sample_args);
+                break;
+            }
             default:
                 LOG_INFO("get_sigmas with discrete scheduler (default)");
                 scheduler = std::make_shared<DiscreteScheduler>();
@@ -988,10 +1114,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 +1129,7 @@ struct FluxFlowDenoiser : public DiscreteFlowDenoiser {
 
 struct SefiFlowDenoiser;
 
-struct Flux2FlowDenoiser : public FluxFlowDenoiser {
+struct SefiFlowDenoiser : 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 {
     static constexpr int kNumTrainTimesteps = 1000;
     static constexpr int kSemChannels       = 16;
     static constexpr int kTotalChannels     = 144;

src/stable-diffusion.cpp

@@ -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,8 @@ const char* scheduler_to_str[] = {
     "bong_tangent",
     "ltx2",
     "logit_normal",
+    "flux2",
+    "flux",
 };
 
 const char* sd_scheduler_name(enum scheduler_t scheduler) {
@@ -2576,6 +2572,9 @@ const char* sd_scheduler_name(enum scheduler_t scheduler) {
 }
 
 enum scheduler_t str_to_scheduler(const char* str) {
+    if (!strcmp(str, "normal")) {
+        return DISCRETE_SCHEDULER;
+    }
     for (int i = 0; i < SCHEDULER_COUNT; i++) {
         if (!strcmp(str, scheduler_to_str[i])) {
             return (enum scheduler_t)i;
@@ -2590,7 +2589,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 +3165,10 @@ 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_flux(sd_ctx->sd->version)) {
+        return FLUX_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)) {