diff --git a/src/runtime/denoiser.hpp b/src/runtime/denoiser.cpp similarity index 93% rename from src/runtime/denoiser.hpp rename to src/runtime/denoiser.cpp index f58f42b55..3ad8bfd94 100644 --- a/src/runtime/denoiser.hpp +++ b/src/runtime/denoiser.cpp @@ -1,5 +1,4 @@ -#ifndef __SD_RUNTIME_DENOISER_HPP__ -#define __SD_RUNTIME_DENOISER_HPP__ +#include "denoiser.h" #include #include @@ -12,17 +11,11 @@ #include #include "core/ggml_extend.hpp" -#include "core/tensor.hpp" #include "runtime/gits_noise.h" -#include "runtime/guidance.h" -/*================================================= CompVisDenoiser ==================================================*/ // Ref: https://github.com/crowsonkb/k-diffusion/blob/master/k_diffusion/external.py -#define TIMESTEPS 1000 -#define FLUX_TIMESTEPS 1000 - struct SigmaScheduler { typedef std::function t_to_sigma_t; @@ -1030,98 +1023,120 @@ struct LogitNormalScheduler : SigmaScheduler { } }; -struct Denoiser { - virtual float sigma_min() = 0; - virtual float sigma_max() = 0; - virtual float sigma_to_t(float sigma) = 0; - virtual float t_to_sigma(float t) = 0; - virtual std::vector get_scalings(float sigma) = 0; - virtual sd::Tensor noise_scaling(float sigma, - const sd::Tensor& noise, - const sd::Tensor& latent) = 0; - virtual sd::Tensor inverse_noise_scaling(float sigma, - const sd::Tensor& latent) = 0; - virtual float noise_level_to_sigma(float noise_level) = 0; - - virtual std::vector get_sigmas(uint32_t n, int image_seq_len, scheduler_t scheduler_type, SDVersion version, const char* extra_sample_args = nullptr) { - auto bound_t_to_sigma = std::bind(&Denoiser::t_to_sigma, this, std::placeholders::_1); - std::shared_ptr scheduler; - switch (scheduler_type) { - case DISCRETE_SCHEDULER: - LOG_INFO("get_sigmas with discrete scheduler"); - scheduler = std::make_shared(); - break; - case KARRAS_SCHEDULER: - LOG_INFO("get_sigmas with Karras scheduler"); - scheduler = std::make_shared(); - break; - case BETA_SCHEDULER: - LOG_INFO("get_sigmas with Beta scheduler"); - scheduler = std::make_shared(extra_sample_args); - break; - case EXPONENTIAL_SCHEDULER: - LOG_INFO("get_sigmas exponential scheduler"); - scheduler = std::make_shared(); - break; - case AYS_SCHEDULER: - LOG_INFO("get_sigmas with Align-Your-Steps scheduler"); - scheduler = std::make_shared(version); - break; - case GITS_SCHEDULER: - LOG_INFO("get_sigmas with GITS scheduler"); - scheduler = std::make_shared(); - break; - case SGM_UNIFORM_SCHEDULER: - LOG_INFO("get_sigmas with SGM Uniform scheduler"); - scheduler = std::make_shared(); - break; - case SIMPLE_SCHEDULER: - LOG_INFO("get_sigmas with Simple scheduler"); - scheduler = std::make_shared(); - break; - case SMOOTHSTEP_SCHEDULER: - LOG_INFO("get_sigmas with SmoothStep scheduler"); - scheduler = std::make_shared(); - break; - case BONG_TANGENT_SCHEDULER: - LOG_INFO("get_sigmas with bong_tangent scheduler"); - scheduler = std::make_shared(); - break; - case KL_OPTIMAL_SCHEDULER: - LOG_INFO("get_sigmas with KL Optimal scheduler"); - scheduler = std::make_shared(); - break; - case LCM_SCHEDULER: - LOG_INFO("get_sigmas with LCM scheduler"); - scheduler = std::make_shared(); - break; - case LTX2_SCHEDULER: - LOG_INFO("get_sigmas with LTX2 scheduler"); - scheduler = std::make_shared(image_seq_len, extra_sample_args); - break; - case LOGIT_NORMAL_SCHEDULER: { - LOG_INFO("get_sigmas with Logit-Normal scheduler"); - scheduler = std::make_shared(image_seq_len, extra_sample_args); - break; - } - case FLUX2_SCHEDULER: { - LOG_INFO("get_sigmas with Flux2 scheduler"); - scheduler = std::make_shared(image_seq_len); - break; - } - case FLUX_SCHEDULER: { - LOG_INFO("get_sigmas with Flux scheduler"); - scheduler = std::make_shared(image_seq_len, extra_sample_args); - break; - } - default: - LOG_INFO("get_sigmas with discrete scheduler (default)"); - scheduler = std::make_shared(); - break; +std::vector Denoiser::get_sigmas(uint32_t n, int image_seq_len, scheduler_t scheduler_type, SDVersion version, const char* extra_sample_args) { + auto bound_t_to_sigma = std::bind(&Denoiser::t_to_sigma, this, std::placeholders::_1); + std::shared_ptr scheduler; + switch (scheduler_type) { + case DISCRETE_SCHEDULER: + LOG_INFO("get_sigmas with discrete scheduler"); + scheduler = std::make_shared(); + break; + case KARRAS_SCHEDULER: + LOG_INFO("get_sigmas with Karras scheduler"); + scheduler = std::make_shared(); + break; + case BETA_SCHEDULER: + LOG_INFO("get_sigmas with Beta scheduler"); + scheduler = std::make_shared(extra_sample_args); + break; + case EXPONENTIAL_SCHEDULER: + LOG_INFO("get_sigmas exponential scheduler"); + scheduler = std::make_shared(); + break; + case AYS_SCHEDULER: + LOG_INFO("get_sigmas with Align-Your-Steps scheduler"); + scheduler = std::make_shared(version); + break; + case GITS_SCHEDULER: + LOG_INFO("get_sigmas with GITS scheduler"); + scheduler = std::make_shared(); + break; + case SGM_UNIFORM_SCHEDULER: + LOG_INFO("get_sigmas with SGM Uniform scheduler"); + scheduler = std::make_shared(); + break; + case SIMPLE_SCHEDULER: + LOG_INFO("get_sigmas with Simple scheduler"); + scheduler = std::make_shared(); + break; + case SMOOTHSTEP_SCHEDULER: + LOG_INFO("get_sigmas with SmoothStep scheduler"); + scheduler = std::make_shared(); + break; + case BONG_TANGENT_SCHEDULER: + LOG_INFO("get_sigmas with bong_tangent scheduler"); + scheduler = std::make_shared(); + break; + case KL_OPTIMAL_SCHEDULER: + LOG_INFO("get_sigmas with KL Optimal scheduler"); + scheduler = std::make_shared(); + break; + case LCM_SCHEDULER: + LOG_INFO("get_sigmas with LCM scheduler"); + scheduler = std::make_shared(); + break; + case LTX2_SCHEDULER: + LOG_INFO("get_sigmas with LTX2 scheduler"); + scheduler = std::make_shared(image_seq_len, extra_sample_args); + break; + case LOGIT_NORMAL_SCHEDULER: { + LOG_INFO("get_sigmas with Logit-Normal scheduler"); + scheduler = std::make_shared(image_seq_len, extra_sample_args); + break; + } + case FLUX2_SCHEDULER: { + LOG_INFO("get_sigmas with Flux2 scheduler"); + scheduler = std::make_shared(image_seq_len); + break; + } + case FLUX_SCHEDULER: { + LOG_INFO("get_sigmas with Flux scheduler"); + scheduler = std::make_shared(image_seq_len, extra_sample_args); + break; } - return scheduler->get_sigmas(n, sigma_min(), sigma_max(), bound_t_to_sigma); + default: + LOG_INFO("get_sigmas with discrete scheduler (default)"); + scheduler = std::make_shared(); + break; } -}; + return scheduler->get_sigmas(n, sigma_min(), sigma_max(), bound_t_to_sigma); +} + +void Denoiser::refresh_compvis_denoiser(const std::vector& file_alphas_cumprod) { + (void)file_alphas_cumprod; +} + +scheduler_t Denoiser::get_default_scheduler() { + return SCHEDULER_COUNT; +} + +bool Denoiser::is_flow_denoiser() { + return false; +} + +void Denoiser::set_shift(float shift) { + (void)shift; +} + +std::vector Denoiser::get_timesteps(int step) { + (void)step; + return std::vector(); +} + +static void calculate_alphas_cumprod(float* alphas_cumprod, + float linear_start = 0.00085f, + float linear_end = 0.0120f, + int timesteps = TIMESTEPS) { + float ls_sqrt = sqrtf(linear_start); + float le_sqrt = sqrtf(linear_end); + float amount = le_sqrt - ls_sqrt; + float product = 1.0f; + for (int i = 0; i < timesteps; i++) { + float beta = ls_sqrt + amount * ((float)i / (timesteps - 1)); + product *= 1.0f - powf(beta, 2.0f); + alphas_cumprod[i] = product; + } +} struct CompVisDenoiser : public Denoiser { float sigmas[TIMESTEPS]; @@ -1193,6 +1208,19 @@ struct CompVisDenoiser : public Denoiser { float noise_level_to_sigma(float noise_level) override { return noise_level / (1.0f - noise_level); } + + void refresh_compvis_denoiser(const std::vector& file_alphas_cumprod) override { + std::vector alphas_cumprod(TIMESTEPS); + if (file_alphas_cumprod.size() == TIMESTEPS) { + alphas_cumprod = file_alphas_cumprod; + } else { + calculate_alphas_cumprod(alphas_cumprod.data()); + } + for (int i = 0; i < TIMESTEPS; i++) { + this->sigmas[i] = std::sqrt((1 - alphas_cumprod[i]) / alphas_cumprod[i]); + this->log_sigmas[i] = std::log(this->sigmas[i]); + } + } }; struct CompVisVDenoiser : public CompVisDenoiser { @@ -1227,6 +1255,10 @@ struct EDMVDenoiser : public CompVisVDenoiser { float sigma_max() override { return max_sigma; } + + scheduler_t get_default_scheduler() override { + return EXPONENTIAL_SCHEDULER; + } }; inline float time_snr_shift(float alpha, float t) { @@ -1243,7 +1275,11 @@ struct DiscreteFlowDenoiser : public Denoiser { set_shift(shift); } - void set_shift(float shift) { + bool is_flow_denoiser() override { + return true; + } + + void set_shift(float shift) override { this->shift = shift; } @@ -1299,8 +1335,6 @@ struct FluxFlowDenoiser : public DiscreteFlowDenoiser { } }; -struct SefiFlowDenoiser; - struct SefiFlowDenoiser : public FluxFlowDenoiser { static constexpr int kNumTrainTimesteps = 1000; static constexpr int kSemChannels = 16; @@ -1373,6 +1407,15 @@ struct SefiFlowDenoiser : public FluxFlowDenoiser { n, timestep_shift_alpha, delta_t); return tex_sigmas; } + + std::vector get_timesteps(int step) override { + int sched_idx = step > 0 ? step - 1 : 0; + if (sched_idx >= static_cast(tex_timesteps.size())) { + sched_idx = static_cast(tex_timesteps.size()) - 1; + } + return {sem_timesteps[sched_idx], + tex_timesteps[sched_idx]}; + } }; // MiniT2I predicts x0 directly and integrates a linear flow ODE: @@ -1442,7 +1485,26 @@ struct MiniT2IFlowDenoiser : public Denoiser { } }; -typedef std::function&, float, int)> denoise_cb_t; +std::shared_ptr make_denoiser(enum prediction_t pred_type) { + switch (pred_type) { + case EPS_PRED: + return std::make_shared(); + case V_PRED: + return std::make_shared(); + case EDM_V_PRED: + return std::make_shared(); + case FLOW_PRED: + return std::make_shared(); + case FLUX_FLOW_PRED: + return std::make_shared(); + case SEFI_FLOW_PRED: + return std::make_shared(); + case MINIT2I_FLOW_PRED: + return std::make_shared(); + default: + return nullptr; + } +} static std::pair get_ancestral_step(float sigma_from, float sigma_to, @@ -2767,15 +2829,15 @@ static sd::Tensor sample_gradient_estimation(denoise_cb_t model, } // k diffusion reverse ODE: dx = (x - D(x;\sigma)) / \sigma dt; \sigma(t) = t -static sd::Tensor sample_k_diffusion(sample_method_t method, - denoise_cb_t model, - sd::Tensor x, - std::vector sigmas, - std::shared_ptr rng, - float eta, - bool is_flow_denoiser, - const char* extra_sample_args, - std::shared_ptr denoiser_for_dispatch = nullptr) { +sd::Tensor sample_k_diffusion(sample_method_t method, + denoise_cb_t model, + sd::Tensor x, + const std::vector& sigmas, + std::shared_ptr rng, + float eta, + bool is_flow_denoiser, + const char* extra_sample_args, + std::shared_ptr denoiser_for_dispatch) { if (denoiser_for_dispatch) { if (auto sefi = std::dynamic_pointer_cast(denoiser_for_dispatch)) { return sample_sefi_euler(sefi.get(), model, std::move(x)); @@ -2833,5 +2895,3 @@ static sd::Tensor sample_k_diffusion(sample_method_t method, return {}; } } - -#endif // __SD_RUNTIME_DENOISER_HPP__ diff --git a/src/runtime/denoiser.h b/src/runtime/denoiser.h new file mode 100644 index 000000000..507ecf108 --- /dev/null +++ b/src/runtime/denoiser.h @@ -0,0 +1,51 @@ +#ifndef __SD_RUNTIME_DENOISER_H__ +#define __SD_RUNTIME_DENOISER_H__ + +#include "stable-diffusion.h" + +#include "src/core/rng.hpp" +#include "src/core/tensor.hpp" +#include "src/model.h" +#include "src/runtime/guidance.h" + +#include +#include + +const int TIMESTEPS = 1000; + +struct Denoiser { + virtual float sigma_min() = 0; + virtual float sigma_max() = 0; + virtual float sigma_to_t(float sigma) = 0; + virtual float t_to_sigma(float t) = 0; + virtual std::vector get_scalings(float sigma) = 0; + virtual sd::Tensor noise_scaling(float sigma, + const sd::Tensor& noise, + const sd::Tensor& latent) = 0; + virtual sd::Tensor inverse_noise_scaling(float sigma, + const sd::Tensor& latent) = 0; + virtual float noise_level_to_sigma(float noise_level) = 0; + + virtual std::vector get_sigmas(uint32_t n, int image_seq_len, scheduler_t scheduler_type, SDVersion version, const char* extra_sample_args = nullptr); + virtual void refresh_compvis_denoiser(const std::vector& file_alphas_cumprod); + virtual scheduler_t get_default_scheduler(); + virtual bool is_flow_denoiser(); + virtual void set_shift(float shift); + virtual std::vector get_timesteps(int step); +}; + +std::shared_ptr make_denoiser(prediction_t pred_type); + +typedef std::function&, float, int)> denoise_cb_t; + +sd::Tensor sample_k_diffusion(sample_method_t method, + denoise_cb_t model, + sd::Tensor x, + const std::vector& sigmas, + std::shared_ptr rng, + float eta, + bool is_flow_denoiser, + const char* extra_sample_args, + std::shared_ptr denoiser_for_dispatch = nullptr); + +#endif // __SD_RUNTIME_DENOISER_H__ diff --git a/src/runtime/easycache.hpp b/src/runtime/easycache.hpp index 75ae3ddad..b91b59a63 100644 --- a/src/runtime/easycache.hpp +++ b/src/runtime/easycache.hpp @@ -9,7 +9,7 @@ #include "core/ggml_extend.hpp" #include "core/tensor.hpp" #include "runtime/condition_cache_utils.hpp" -#include "runtime/denoiser.hpp" +#include "runtime/denoiser.h" struct EasyCacheConfig { bool enabled = false; diff --git a/src/runtime/sample-cache.h b/src/runtime/sample-cache.h index cb0fe7bc5..cdf827509 100644 --- a/src/runtime/sample-cache.h +++ b/src/runtime/sample-cache.h @@ -7,7 +7,7 @@ #include "core/util.h" #include "model.h" #include "runtime/cache_dit.hpp" -#include "runtime/denoiser.hpp" +#include "runtime/denoiser.h" #include "runtime/easycache.hpp" #include "runtime/spectrum.hpp" #include "runtime/ucache.hpp" diff --git a/src/runtime/ucache.hpp b/src/runtime/ucache.hpp index 187e1e789..a68232ac0 100644 --- a/src/runtime/ucache.hpp +++ b/src/runtime/ucache.hpp @@ -9,7 +9,7 @@ #include "core/ggml_extend.hpp" #include "core/tensor.hpp" #include "runtime/condition_cache_utils.hpp" -#include "runtime/denoiser.hpp" +#include "runtime/denoiser.h" struct UCacheConfig { bool enabled = false; diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index 5c34d4042..19d9777a2 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -56,7 +56,7 @@ #include "model/vae/tae.hpp" #include "model/vae/vae.hpp" #include "model/vae/wan_vae.hpp" -#include "runtime/denoiser.hpp" +#include "runtime/denoiser.h" #include "runtime/guidance.h" #include "runtime/sample-cache.h" #include "upscaler.h" @@ -163,21 +163,6 @@ static bool sd_version_supports_img_cfg(SDVersion version, bool has_ref_images) (has_ref_images && sd_version_supports_ref_latent_img_cfg(version)); } -void calculate_alphas_cumprod(float* alphas_cumprod, - float linear_start = 0.00085f, - float linear_end = 0.0120f, - int timesteps = TIMESTEPS) { - float ls_sqrt = sqrtf(linear_start); - float le_sqrt = sqrtf(linear_end); - float amount = le_sqrt - ls_sqrt; - float product = 1.0f; - for (int i = 0; i < timesteps; i++) { - float beta = ls_sqrt + amount * ((float)i / (timesteps - 1)); - product *= 1.0f - powf(beta, 2.0f); - alphas_cumprod[i] = product; - } -} - static float get_cache_reuse_threshold(const sd_cache_params_t& params) { float reuse_threshold = params.reuse_threshold; if (reuse_threshold == INFINITY) { @@ -255,7 +240,7 @@ class StableDiffusionGGML { std::shared_ptr model_manager; - std::shared_ptr denoiser = std::make_shared(); + std::shared_ptr denoiser = make_denoiser(EPS_PRED); std::vector file_alphas_cumprod; StableDiffusionGGML() = default; @@ -655,23 +640,6 @@ class StableDiffusionGGML { } } - void refresh_compvis_denoiser_sigmas() { - auto comp_vis_denoiser = std::dynamic_pointer_cast(denoiser); - if (!comp_vis_denoiser) { - return; - } - std::vector alphas_cumprod(TIMESTEPS); - if (file_alphas_cumprod.size() == TIMESTEPS) { - alphas_cumprod = file_alphas_cumprod; - } else { - calculate_alphas_cumprod(alphas_cumprod.data()); - } - for (int i = 0; i < TIMESTEPS; i++) { - comp_vis_denoiser->sigmas[i] = std::sqrt((1 - alphas_cumprod[i]) / alphas_cumprod[i]); - comp_vis_denoiser->log_sigmas[i] = std::log(comp_vis_denoiser->sigmas[i]); - } - } - void load_alphas_cumprod(ModelLoader& model_loader) { file_alphas_cumprod.clear(); @@ -1794,35 +1762,35 @@ class StableDiffusionGGML { break; case V_PRED: LOG_INFO("running in v-prediction mode"); - denoiser = std::make_shared(); + denoiser = make_denoiser(pred_type); break; case EDM_V_PRED: LOG_INFO("running in v-prediction EDM mode"); - denoiser = std::make_shared(); + denoiser = make_denoiser(pred_type); break; case FLOW_PRED: { if (sd_version_is_ltxav(version)) { LOG_INFO("running in LTXAV FLOW mode"); - denoiser = std::make_shared(); + denoiser = make_denoiser(FLUX_FLOW_PRED); } else { LOG_INFO("running in FLOW mode"); - denoiser = std::make_shared(); + denoiser = make_denoiser(pred_type); } break; } case FLUX_FLOW_PRED: { LOG_INFO("running in Flux FLOW mode"); - denoiser = std::make_shared(); + denoiser = make_denoiser(pred_type); break; } case SEFI_FLOW_PRED: { LOG_INFO("running in SeFi-Image dual-time FLOW mode"); - denoiser = std::make_shared(); + denoiser = make_denoiser(pred_type); break; } case MINIT2I_FLOW_PRED: { LOG_INFO("running in MiniT2I FLOW mode"); - denoiser = std::make_shared(); + denoiser = make_denoiser(pred_type); break; } default: { @@ -1831,7 +1799,7 @@ class StableDiffusionGGML { } } - refresh_compvis_denoiser_sigmas(); + denoiser->refresh_compvis_denoiser(file_alphas_cumprod); } return true; @@ -2162,13 +2130,9 @@ class StableDiffusionGGML { const sd::Tensor& init_latent, const sd::Tensor& denoise_mask, int step) { - if (auto sefi_denoiser = std::dynamic_pointer_cast(denoiser)) { - int sched_idx = step > 0 ? step - 1 : 0; - if (sched_idx >= static_cast(sefi_denoiser->tex_timesteps.size())) { - sched_idx = static_cast(sefi_denoiser->tex_timesteps.size()) - 1; - } - return {sefi_denoiser->sem_timesteps[sched_idx], - sefi_denoiser->tex_timesteps[sched_idx]}; + auto denoiser_timesteps = denoiser->get_timesteps(step); + if (!denoiser_timesteps.empty()) { + return denoiser_timesteps; } if (diffusion_model->get_desc() == "Wan2.2-TI2V-5B") { int64_t frame_count = init_latent.shape()[2]; @@ -3017,18 +2981,14 @@ class StableDiffusionGGML { } void set_flow_shift(float flow_shift = INFINITY) { - auto flow_denoiser = std::dynamic_pointer_cast(denoiser); - if (flow_denoiser) { - if (flow_shift == INFINITY) { - flow_shift = default_flow_shift; - } - flow_denoiser->set_shift(flow_shift); + if (flow_shift == INFINITY) { + flow_shift = default_flow_shift; } + denoiser->set_shift(flow_shift); } bool is_flow_denoiser() { - auto flow_denoiser = std::dynamic_pointer_cast(denoiser); - return !!flow_denoiser; + return denoiser->is_flow_denoiser(); } std::string get_default_ref_image_preset(SDVersion version) const { @@ -3882,9 +3842,9 @@ enum sample_method_t sd_get_default_sample_method(const sd_ctx_t* sd_ctx) { enum scheduler_t sd_get_default_scheduler(const sd_ctx_t* sd_ctx, enum sample_method_t sample_method) { if (sd_ctx != nullptr && sd_ctx->sd != nullptr) { - auto edm_v_denoiser = std::dynamic_pointer_cast(sd_ctx->sd->denoiser); - if (edm_v_denoiser) { - return EXPONENTIAL_SCHEDULER; + scheduler_t denoiser_default = sd_ctx->sd->denoiser->get_default_scheduler(); + if (denoiser_default != SCHEDULER_COUNT) { + return denoiser_default; } } if (sample_method == LCM_SAMPLE_METHOD || sample_method == TCD_SAMPLE_METHOD) {