diff --git a/examples/common/common.cpp b/examples/common/common.cpp index d4c8a72b..bacb4eda 100644 --- a/examples/common/common.cpp +++ b/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& target_skip_layers, std::vector* 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(); + } 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(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(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) + ", "; diff --git a/examples/common/common.h b/examples/common/common.h index f87293f3..badaa875 100644 --- a/examples/common/common.h +++ b/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 skip_layers = {7, 8, 9}; std::vector high_noise_skip_layers = {7, 8, 9}; diff --git a/include/stable-diffusion.h b/include/stable-diffusion.h index 3596c20c..37e05dae 100644 --- a/include/stable-diffusion.h +++ b/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 { diff --git a/src/denoiser.hpp b/src/denoiser.hpp index 444cc45b..16691088 100644 --- a/src/denoiser.hpp +++ b/src/denoiser.hpp @@ -2,6 +2,7 @@ #define __DENOISER_HPP__ #include +#include #include #include "ggml_extend.hpp" @@ -867,49 +868,6 @@ static sd::Tensor sample_euler_flow(denoise_cb_t model, return x; } -static sd::Tensor sample_euler_flow_flash(denoise_cb_t model, - sd::Tensor x, - const std::vector& sigmas, - std::shared_ptr rng, - float eta) { - constexpr float noise_clip_std = 2.5f; - float s_noise = eta; - int steps = static_cast(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 denoised = std::move(denoised_opt); - if (sigma_next == 0.0f) { - x = std::move(denoised); - continue; - } - auto noise = sd::Tensor::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(noise[j]); - } - mean /= static_cast(noise.numel()); - - double variance = 0.0; - for (int64_t j = 0; j < noise.numel(); ++j) { - double centered = static_cast(noise[j]) - mean; - variance += centered * centered; - } - variance /= static_cast(noise.numel()); - - float clip_val = noise_clip_std * static_cast(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 sample_euler(denoise_cb_t model, sd::Tensor x, const std::vector& sigmas) { @@ -1191,7 +1149,80 @@ static sd::Tensor sample_lcm(denoise_cb_t model, sd::Tensor x, const std::vector& sigmas, std::shared_ptr 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(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 sample_lcm(denoise_cb_t model, if (is_flow_denoiser) { x *= (1 - sigmas[i + 1]); } - x += sd::Tensor::randn_like(x, rng) * sigmas[i + 1]; + auto noise = sd::Tensor::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(noise[j]); + } + mean /= static_cast(noise.numel()); + + double variance = 0.0; + for (int64_t j = 0; j < noise.numel(); ++j) { + double centered = static_cast(noise[j]) - mean; + variance += centered * centered; + } + variance /= static_cast(noise.numel()); + + float clip_val = args.noise_clip_std * static_cast(std::sqrt(variance)); + noise = sd::ops::clamp(noise, -clip_val, clip_val); + } + float t = steps > 1 ? static_cast(i) / static_cast(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 sample_k_diffusion(sample_method_t method, std::vector sigmas, std::shared_ptr 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 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: diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index b97d1f47..ec5868b9 100644 --- a/src/stable-diffusion.cpp +++ b/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& sigmas, int start_merge_step, const std::vector>& 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; - eta = sd_img_gen_params->sample_params.eta; - sample_steps = sd_img_gen_params->sample_params.sample_steps; + sample_method = sd_img_gen_params->sample_params.sample_method; + extra_sample_args = sd_img_gen_params->sample_params.extra_sample_args; + 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; - eta = sd_vid_gen_params->sample_params.eta; - sample_steps = sd_vid_gen_params->sample_params.sample_steps; + sample_method = sd_vid_gen_params->sample_params.sample_method; + extra_sample_args = sd_vid_gen_params->sample_params.extra_sample_args; + 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_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_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_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>{}, @@ -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>{},