move flash-flow sampling options into LCM extra args

examples/common/common.cpp

@@ -807,6 +807,10 @@ ArgOptions SDGenerationParams::get_options() {
          "Latent (antialiased), Latent (bicubic), Latent (bicubic antialiased), or a model name "
          "under --hires-upscalers-dir (default: Latent)",
          &hires_upscaler},
+        {"",
+         "--extra-sample-args",
+         "extra sampler args, key=value list. Currently lcm supports noise_clip_std, noise_scale_start, noise_scale_end",
+         &extra_sample_args},
     };
 
     options.int_options = {
@@ -1607,6 +1611,7 @@ bool SDGenerationParams::from_json_str(
 
     auto parse_sample_params_json = [&](const json& sample_json,
                                         sd_sample_params_t& target_params,
+                                        std::string& target_extra_sample_args,
                                         std::vector<int>& target_skip_layers,
                                         std::vector<float>* target_custom_sigmas) {
         if (sample_json.contains("sample_steps") && sample_json["sample_steps"].is_number_integer()) {
@@ -1621,6 +1626,9 @@ bool SDGenerationParams::from_json_str(
         if (sample_json.contains("flow_shift") && sample_json["flow_shift"].is_number()) {
             target_params.flow_shift = sample_json["flow_shift"];
         }
+        if (sample_json.contains("extra_sample_args") && sample_json["extra_sample_args"].is_string()) {
+            target_extra_sample_args = sample_json["extra_sample_args"].get<std::string>();
+        }
         if (target_custom_sigmas != nullptr &&
             sample_json.contains("custom_sigmas") &&
             sample_json["custom_sigmas"].is_array()) {
@@ -1668,11 +1676,12 @@ bool SDGenerationParams::from_json_str(
     };
 
     if (j.contains("sample_params") && j["sample_params"].is_object()) {
-        parse_sample_params_json(j["sample_params"], sample_params, skip_layers, &custom_sigmas);
+        parse_sample_params_json(j["sample_params"], sample_params, extra_sample_args, skip_layers, &custom_sigmas);
     }
     if (j.contains("high_noise_sample_params") && j["high_noise_sample_params"].is_object()) {
         parse_sample_params_json(j["high_noise_sample_params"],
                                  high_noise_sample_params,
+                                 high_noise_extra_sample_args,
                                  high_noise_skip_layers,
                                  nullptr);
     }
@@ -2099,6 +2108,8 @@ sd_img_gen_params_t SDGenerationParams::to_sd_img_gen_params_t() {
     high_noise_sample_params.guidance.slg.layer_count = high_noise_skip_layers.size();
     sample_params.custom_sigmas                       = custom_sigmas.empty() ? nullptr : custom_sigmas.data();
     sample_params.custom_sigmas_count                 = static_cast<int>(custom_sigmas.size());
+    sample_params.extra_sample_args                   = extra_sample_args.empty() ? nullptr : extra_sample_args.c_str();
+    high_noise_sample_params.extra_sample_args        = high_noise_extra_sample_args.empty() ? nullptr : high_noise_extra_sample_args.c_str();
     cache_params.scm_mask                             = scm_mask.empty() ? nullptr : scm_mask.c_str();
 
     sd_pm_params_t pm_params = {
@@ -2168,6 +2179,8 @@ sd_vid_gen_params_t SDGenerationParams::to_sd_vid_gen_params_t() {
     high_noise_sample_params.guidance.slg.layer_count = high_noise_skip_layers.size();
     sample_params.custom_sigmas                       = custom_sigmas.empty() ? nullptr : custom_sigmas.data();
     sample_params.custom_sigmas_count                 = static_cast<int>(custom_sigmas.size());
+    sample_params.extra_sample_args                   = extra_sample_args.empty() ? nullptr : extra_sample_args.c_str();
+    high_noise_sample_params.extra_sample_args        = high_noise_extra_sample_args.empty() ? nullptr : high_noise_extra_sample_args.c_str();
     cache_params.scm_mask                             = scm_mask.empty() ? nullptr : scm_mask.c_str();
 
     params.loras                    = lora_vec.empty() ? nullptr : lora_vec.data();
@@ -2306,6 +2319,7 @@ static json build_sampling_metadata_json(const sd_sample_params_t& sample_params
         {"eta", sample_params.eta},
         {"shifted_timestep", sample_params.shifted_timestep},
         {"flow_shift", sample_params.flow_shift},
+        {"extra_sample_args", safe_json_string(sample_params.extra_sample_args)},
         {"guidance",
          {
              {"txt_cfg", sample_params.guidance.txt_cfg},
@@ -2497,6 +2511,9 @@ std::string get_image_params(const SDContextParams& ctx_params,
     }
     parameter_string += "Guidance: " + std::to_string(gen_params.sample_params.guidance.distilled_guidance) + ", ";
     parameter_string += "Eta: " + std::to_string(gen_params.sample_params.eta) + ", ";
+    if (!gen_params.extra_sample_args.empty()) {
+        parameter_string += "Extra sample args: " + gen_params.extra_sample_args + ", ";
+    }
     parameter_string += "Seed: " + std::to_string(seed) + ", ";
     parameter_string += "Size: " + std::to_string(gen_params.get_resolved_width()) + "x" + std::to_string(gen_params.get_resolved_height()) + ", ";
     parameter_string += "Model: " + sd_basename(ctx_params.model_path) + ", ";

examples/common/common.h

@@ -168,6 +168,8 @@ struct SDGenerationParams {
 
     sd_sample_params_t sample_params;
     sd_sample_params_t high_noise_sample_params;
+    std::string extra_sample_args;
+    std::string high_noise_extra_sample_args;
     std::vector<int> skip_layers            = {7, 8, 9};
     std::vector<int> high_noise_skip_layers = {7, 8, 9};
 

include/stable-diffusion.h

@@ -37,7 +37,6 @@ enum rng_type_t {
 
 enum sample_method_t {
     EULER_SAMPLE_METHOD,
-    EULER_FLOW_FLASH_SAMPLE_METHOD,
     EULER_A_SAMPLE_METHOD,
     HEUN_SAMPLE_METHOD,
     DPM2_SAMPLE_METHOD,
@@ -239,6 +238,7 @@ typedef struct {
     float* custom_sigmas;
     int custom_sigmas_count;
     float flow_shift;
+    const char* extra_sample_args;
 } sd_sample_params_t;
 
 typedef struct {

src/denoiser.hpp

@@ -2,6 +2,7 @@
 #define __DENOISER_HPP__
 
 #include <cmath>
+#include <string>
 #include <utility>
 
 #include "ggml_extend.hpp"
@@ -867,49 +868,6 @@ static sd::Tensor<float> sample_euler_flow(denoise_cb_t model,
     return x;
 }
 
-static sd::Tensor<float> sample_euler_flow_flash(denoise_cb_t model,
-                                                 sd::Tensor<float> x,
-                                                 const std::vector<float>& sigmas,
-                                                 std::shared_ptr<RNG> rng,
-                                                 float eta) {
-    constexpr float noise_clip_std = 2.5f;
-    float s_noise                  = eta;
-    int steps                      = static_cast<int>(sigmas.size()) - 1;
-    for (int i = 0; i < steps; i++) {
-        float sigma       = sigmas[i];
-        float sigma_next  = sigmas[i + 1];
-        auto denoised_opt = model(x, sigma, i + 1);
-        if (denoised_opt.empty()) {
-            return {};
-        }
-        sd::Tensor<float> denoised = std::move(denoised_opt);
-        if (sigma_next == 0.0f) {
-            x = std::move(denoised);
-            continue;
-        }
-        auto noise = sd::Tensor<float>::randn_like(x, rng);
-        if (noise_clip_std > 0.0f && noise.numel() > 0) {
-            double mean = 0.0;
-            for (int64_t j = 0; j < noise.numel(); ++j) {
-                mean += static_cast<double>(noise[j]);
-            }
-            mean /= static_cast<double>(noise.numel());
-
-            double variance = 0.0;
-            for (int64_t j = 0; j < noise.numel(); ++j) {
-                double centered = static_cast<double>(noise[j]) - mean;
-                variance += centered * centered;
-            }
-            variance /= static_cast<double>(noise.numel());
-
-            float clip_val = noise_clip_std * static_cast<float>(std::sqrt(variance));
-            noise          = sd::ops::clamp(noise, -clip_val, clip_val);
-        }
-        x = sigma_next * noise * s_noise + (1.0f - sigma_next) * denoised;
-    }
-    return x;
-}
-
 static sd::Tensor<float> sample_euler(denoise_cb_t model,
                                       sd::Tensor<float> x,
                                       const std::vector<float>& sigmas) {
@@ -1191,7 +1149,80 @@ static sd::Tensor<float> sample_lcm(denoise_cb_t model,
                                     sd::Tensor<float> x,
                                     const std::vector<float>& sigmas,
                                     std::shared_ptr<RNG> rng,
-                                    bool is_flow_denoiser) {
+                                    bool is_flow_denoiser,
+                                    const char* extra_sample_args = nullptr) {
+    struct LCMSampleArgs {
+        float noise_clip_std    = 0.0f;
+        float noise_scale_start = 1.0f;
+        float noise_scale_end   = 1.0f;
+    };
+
+    auto trim = [](std::string value) -> std::string {
+        const char* whitespace = " \t\r\n";
+        size_t begin           = value.find_first_not_of(whitespace);
+        if (begin == std::string::npos) {
+            return "";
+        }
+        size_t end = value.find_last_not_of(whitespace);
+        return value.substr(begin, end - begin + 1);
+    };
+
+    LCMSampleArgs args;
+    if (extra_sample_args != nullptr && extra_sample_args[0] != '\0') {
+        std::string raw(extra_sample_args);
+        size_t start                   = 0;
+        bool noise_scale_end_was_set   = false;
+        bool noise_scale_start_was_set = false;
+        auto parse_arg                 = [&](const std::string& item) {
+            std::string token = trim(item);
+            if (token.empty()) {
+                return;
+            }
+            size_t eq = token.find('=');
+            if (eq == std::string::npos) {
+                LOG_WARN("ignoring invalid lcm extra sample arg '%s'", token.c_str());
+                return;
+            }
+
+            std::string key   = trim(token.substr(0, eq));
+            std::string value = trim(token.substr(eq + 1));
+            float parsed      = 0.0f;
+            try {
+                size_t consumed = 0;
+                parsed          = std::stof(value, &consumed);
+                if (trim(value.substr(consumed)).size() != 0) {
+                    LOG_WARN("ignoring invalid lcm extra sample arg '%s'", token.c_str());
+                    return;
+                }
+            } catch (const std::exception&) {
+                LOG_WARN("ignoring invalid lcm extra sample arg '%s'", token.c_str());
+                return;
+            }
+
+            if (key == "noise_clip_std") {
+                args.noise_clip_std = parsed;
+            } else if (key == "noise_scale_start") {
+                args.noise_scale_start    = parsed;
+                noise_scale_start_was_set = true;
+            } else if (key == "noise_scale_end") {
+                args.noise_scale_end    = parsed;
+                noise_scale_end_was_set = true;
+            } else {
+                LOG_WARN("ignoring unknown lcm extra sample arg '%s'", key.c_str());
+            }
+        };
+
+        for (size_t pos = 0; pos <= raw.size(); ++pos) {
+            if (pos == raw.size() || raw[pos] == ',' || raw[pos] == ';') {
+                parse_arg(raw.substr(start, pos - start));
+                start = pos + 1;
+            }
+        }
+        if (noise_scale_start_was_set && !noise_scale_end_was_set) {
+            args.noise_scale_end = args.noise_scale_start;
+        }
+    }
+
     int steps = static_cast<int>(sigmas.size()) - 1;
     for (int i = 0; i < steps; i++) {
         auto denoised_opt = model(x, sigmas[i], i + 1);
@@ -1203,7 +1234,27 @@ static sd::Tensor<float> sample_lcm(denoise_cb_t model,
             if (is_flow_denoiser) {
                 x *= (1 - sigmas[i + 1]);
             }
-            x += sd::Tensor<float>::randn_like(x, rng) * sigmas[i + 1];
+            auto noise = sd::Tensor<float>::randn_like(x, rng);
+            if (args.noise_clip_std > 0.0f && noise.numel() > 0) {
+                double mean = 0.0;
+                for (int64_t j = 0; j < noise.numel(); ++j) {
+                    mean += static_cast<double>(noise[j]);
+                }
+                mean /= static_cast<double>(noise.numel());
+
+                double variance = 0.0;
+                for (int64_t j = 0; j < noise.numel(); ++j) {
+                    double centered = static_cast<double>(noise[j]) - mean;
+                    variance += centered * centered;
+                }
+                variance /= static_cast<double>(noise.numel());
+
+                float clip_val = args.noise_clip_std * static_cast<float>(std::sqrt(variance));
+                noise          = sd::ops::clamp(noise, -clip_val, clip_val);
+            }
+            float t           = steps > 1 ? static_cast<float>(i) / static_cast<float>(steps - 1) : 0.0f;
+            float noise_scale = args.noise_scale_start + (args.noise_scale_end - args.noise_scale_start) * t;
+            x += noise * (sigmas[i + 1] * noise_scale);
         }
     }
     return x;
@@ -1699,10 +1750,9 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
                                             std::vector<float> sigmas,
                                             std::shared_ptr<RNG> rng,
                                             float eta,
-                                            bool is_flow_denoiser) {
+                                            bool is_flow_denoiser,
+                                            const char* extra_sample_args) {
     switch (method) {
-        case EULER_FLOW_FLASH_SAMPLE_METHOD:
-            return sample_euler_flow_flash(model, std::move(x), sigmas, rng, eta);
         case EULER_A_SAMPLE_METHOD:
             if (is_flow_denoiser)
                 return sample_euler_flow(model, std::move(x), sigmas, rng, eta);
@@ -1724,7 +1774,7 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
         case DPMPP2Mv2_SAMPLE_METHOD:
             return sample_dpmpp_2m_v2(model, std::move(x), sigmas);
         case LCM_SAMPLE_METHOD:
-            return sample_lcm(model, std::move(x), sigmas, rng, is_flow_denoiser);
+            return sample_lcm(model, std::move(x), sigmas, rng, is_flow_denoiser, extra_sample_args);
         case IPNDM_SAMPLE_METHOD:
             return sample_ipndm(model, std::move(x), sigmas);
         case IPNDM_V_SAMPLE_METHOD:

src/stable-diffusion.cpp

@@ -60,7 +60,6 @@ const char* model_version_to_str[] = {
 
 const char* sampling_methods_str[] = {
     "Euler",
-    "Euler Flow Flash",
     "Euler A",
     "Heun",
     "DPM2",
@@ -1601,6 +1600,7 @@ public:
                              int shifted_timestep,
                              sample_method_t method,
                              bool is_flow_denoiser,
+                             const char* extra_sample_args,
                              const std::vector<float>& sigmas,
                              int start_merge_step,
                              const std::vector<sd::Tensor<float>>& ref_latents,
@@ -1809,7 +1809,7 @@ public:
             return denoised;
         };
 
-        auto x0_opt = sample_k_diffusion(method, denoise, x_t, sigmas, sampler_rng, eta, is_flow_denoiser);
+        auto x0_opt = sample_k_diffusion(method, denoise, x_t, sigmas, sampler_rng, eta, is_flow_denoiser, extra_sample_args);
         if (x0_opt.empty()) {
             LOG_ERROR("Diffusion model sampling failed");
             if (control_net) {
@@ -1981,7 +1981,6 @@ enum rng_type_t str_to_rng_type(const char* str) {
 
 const char* sample_method_to_str[] = {
     "euler",
-    "euler_flow_flash",
     "euler_a",
     "heun",
     "dpm2",
@@ -2301,6 +2300,7 @@ void sd_sample_params_init(sd_sample_params_t* sample_params) {
     sample_params->custom_sigmas               = nullptr;
     sample_params->custom_sigmas_count         = 0;
     sample_params->flow_shift                  = INFINITY;
+    sample_params->extra_sample_args           = nullptr;
 }
 
 char* sd_sample_params_to_str(const sd_sample_params_t* sample_params) {
@@ -2322,7 +2322,8 @@ char* sd_sample_params_to_str(const sd_sample_params_t* sample_params) {
              "sample_steps: %d, "
              "eta: %.2f, "
              "shifted_timestep: %d, "
-             "flow_shift: %.2f)",
+             "flow_shift: %.2f, "
+             "extra_sample_args: %s)",
              sample_params->guidance.txt_cfg,
              std::isfinite(sample_params->guidance.img_cfg)
                  ? sample_params->guidance.img_cfg
@@ -2337,7 +2338,8 @@ char* sd_sample_params_to_str(const sd_sample_params_t* sample_params) {
              sample_params->sample_steps,
              sample_params->eta,
              sample_params->shifted_timestep,
-             sample_params->flow_shift);
+             sample_params->flow_shift,
+             SAFE_STR(sample_params->extra_sample_args));
 
     return buf;
 }
@@ -2770,6 +2772,8 @@ struct GenerationRequest {
 struct SamplePlan {
     enum sample_method_t sample_method            = SAMPLE_METHOD_COUNT;
     enum sample_method_t high_noise_sample_method = SAMPLE_METHOD_COUNT;
+    const char* extra_sample_args                 = nullptr;
+    const char* high_noise_extra_sample_args      = nullptr;
     float eta                                     = 0.f;
     float high_noise_eta                          = 0.f;
     int sample_steps                              = 0;
@@ -2782,22 +2786,25 @@ struct SamplePlan {
     SamplePlan(sd_ctx_t* sd_ctx,
                const sd_img_gen_params_t* sd_img_gen_params,
                const GenerationRequest& request) {
-        sample_method = sd_img_gen_params->sample_params.sample_method;
+        sample_method     = sd_img_gen_params->sample_params.sample_method;
-        eta           = sd_img_gen_params->sample_params.eta;
+        extra_sample_args = sd_img_gen_params->sample_params.extra_sample_args;
-        sample_steps  = sd_img_gen_params->sample_params.sample_steps;
+        eta               = sd_img_gen_params->sample_params.eta;
+        sample_steps      = sd_img_gen_params->sample_params.sample_steps;
         resolve(sd_ctx, &request, &sd_img_gen_params->sample_params);
     }
 
     SamplePlan(sd_ctx_t* sd_ctx,
                const sd_vid_gen_params_t* sd_vid_gen_params,
                const GenerationRequest& request) {
-        sample_method = sd_vid_gen_params->sample_params.sample_method;
+        sample_method     = sd_vid_gen_params->sample_params.sample_method;
-        eta           = sd_vid_gen_params->sample_params.eta;
+        extra_sample_args = sd_vid_gen_params->sample_params.extra_sample_args;
-        sample_steps  = sd_vid_gen_params->sample_params.sample_steps;
+        eta               = sd_vid_gen_params->sample_params.eta;
+        sample_steps      = sd_vid_gen_params->sample_params.sample_steps;
         if (sd_ctx->sd->high_noise_diffusion_model) {
-            high_noise_sample_steps  = sd_vid_gen_params->high_noise_sample_params.sample_steps;
+            high_noise_sample_steps      = sd_vid_gen_params->high_noise_sample_params.sample_steps;
-            high_noise_sample_method = sd_vid_gen_params->high_noise_sample_params.sample_method;
+            high_noise_sample_method     = sd_vid_gen_params->high_noise_sample_params.sample_method;
-            high_noise_eta           = sd_vid_gen_params->high_noise_sample_params.eta;
+            high_noise_extra_sample_args = sd_vid_gen_params->high_noise_sample_params.extra_sample_args;
+            high_noise_eta               = sd_vid_gen_params->high_noise_sample_params.eta;
         }
         moe_boundary = sd_vid_gen_params->moe_boundary;
         resolve(sd_ctx, &request, &sd_vid_gen_params->sample_params);
@@ -3456,6 +3463,7 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
                                                    request.shifted_timestep,
                                                    plan.sample_method,
                                                    sd_ctx->sd->is_flow_denoiser(),
+                                                   plan.extra_sample_args,
                                                    plan.sigmas,
                                                    plan.start_merge_step,
                                                    latents.ref_latents,
@@ -3581,6 +3589,7 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
                                                             request.shifted_timestep,
                                                             plan.sample_method,
                                                             sd_ctx->sd->is_flow_denoiser(),
+                                                            plan.extra_sample_args,
                                                             hires_sigma_sched,
                                                             plan.start_merge_step,
                                                             latents.ref_latents,
@@ -3945,6 +3954,7 @@ SD_API sd_image_t* generate_video(sd_ctx_t* sd_ctx, const sd_vid_gen_params_t* s
                                                            request.shifted_timestep,
                                                            plan.high_noise_sample_method,
                                                            sd_ctx->sd->is_flow_denoiser(),
+                                                           plan.high_noise_extra_sample_args,
                                                            high_noise_sigmas,
                                                            -1,
                                                            std::vector<sd::Tensor<float>>{},
@@ -3987,6 +3997,7 @@ SD_API sd_image_t* generate_video(sd_ctx_t* sd_ctx, const sd_vid_gen_params_t* s
                                                         sd_vid_gen_params->sample_params.shifted_timestep,
                                                         plan.sample_method,
                                                         sd_ctx->sd->is_flow_denoiser(),
+                                                        plan.extra_sample_args,
                                                         plan.sigmas,
                                                         -1,
                                                         std::vector<sd::Tensor<float>>{},