feat: add LCM scheduler (#983)

This commit is contained in:
Wagner Bruna 2025-11-22 02:53:31 -03:00 committed by GitHub
parent 869d023416
commit 45c46779af
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 26 additions and 3 deletions

View File

@ -253,6 +253,23 @@ struct SGMUniformScheduler : SigmaScheduler {
} }
}; };
struct LCMScheduler : SigmaScheduler {
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> result;
result.reserve(n + 1);
const int original_steps = 50;
const int k = TIMESTEPS / original_steps;
for (int i = 0; i < n; i++) {
// the rounding ensures we match the training schedule of the LCM model
int index = (i * original_steps) / n;
int timestep = (original_steps - index) * k - 1;
result.push_back(t_to_sigma(timestep));
}
result.push_back(0.0f);
return result;
}
};
struct KarrasScheduler : SigmaScheduler { struct KarrasScheduler : SigmaScheduler {
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> get_sigmas(uint32_t n, float sigma_min, float sigma_max, t_to_sigma_t t_to_sigma) override {
// These *COULD* be function arguments here, // These *COULD* be function arguments here,
@ -375,6 +392,10 @@ struct Denoiser {
LOG_INFO("get_sigmas with SmoothStep scheduler"); LOG_INFO("get_sigmas with SmoothStep scheduler");
scheduler = std::make_shared<SmoothStepScheduler>(); scheduler = std::make_shared<SmoothStepScheduler>();
break; break;
case LCM_SCHEDULER:
LOG_INFO("get_sigmas with LCM scheduler");
scheduler = std::make_shared<LCMScheduler>();
break;
default: default:
LOG_INFO("get_sigmas with discrete scheduler (default)"); LOG_INFO("get_sigmas with discrete scheduler (default)");
scheduler = std::make_shared<DiscreteScheduler>(); scheduler = std::make_shared<DiscreteScheduler>();

View File

@ -107,8 +107,8 @@ Options:
compatibility issues with quantized parameters, but it usually offers faster inference compatibility issues with quantized parameters, but it usually offers faster inference
speed and, in some cases, lower memory usage. The at_runtime mode, on the other speed and, in some cases, lower memory usage. The at_runtime mode, on the other
hand, is exactly the opposite. hand, is exactly the opposite.
--scheduler denoiser sigma scheduler, one of [discrete, karras, exponential, ays, gits, smoothstep, sgm_uniform, simple], default: --scheduler denoiser sigma scheduler, one of [discrete, karras, exponential, ays, gits, smoothstep, sgm_uniform, simple, lcm],
discrete default: discrete
--skip-layers layers to skip for SLG steps (default: [7,8,9]) --skip-layers layers to skip for SLG steps (default: [7,8,9])
--high-noise-sampling-method (high noise) sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm, --high-noise-sampling-method (high noise) sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm,
ddim_trailing, tcd] default: euler for Flux/SD3/Wan, euler_a otherwise ddim_trailing, tcd] default: euler for Flux/SD3/Wan, euler_a otherwise

View File

@ -1197,7 +1197,7 @@ void parse_args(int argc, const char** argv, SDParams& params) {
on_lora_apply_mode_arg}, on_lora_apply_mode_arg},
{"", {"",
"--scheduler", "--scheduler",
"denoiser sigma scheduler, one of [discrete, karras, exponential, ays, gits, smoothstep, sgm_uniform, simple], default: discrete", "denoiser sigma scheduler, one of [discrete, karras, exponential, ays, gits, smoothstep, sgm_uniform, simple, lcm], default: discrete",
on_scheduler_arg}, on_scheduler_arg},
{"", {"",
"--skip-layers", "--skip-layers",

View File

@ -2268,6 +2268,7 @@ const char* scheduler_to_str[] = {
"sgm_uniform", "sgm_uniform",
"simple", "simple",
"smoothstep", "smoothstep",
"lcm",
}; };
const char* sd_scheduler_name(enum scheduler_t scheduler) { const char* sd_scheduler_name(enum scheduler_t scheduler) {

View File

@ -61,6 +61,7 @@ enum scheduler_t {
SGM_UNIFORM_SCHEDULER, SGM_UNIFORM_SCHEDULER,
SIMPLE_SCHEDULER, SIMPLE_SCHEDULER,
SMOOTHSTEP_SCHEDULER, SMOOTHSTEP_SCHEDULER,
LCM_SCHEDULER,
SCHEDULER_COUNT SCHEDULER_COUNT
}; };