From 10262f86046f8fa308ab44ea064f4a3827005b94 Mon Sep 17 00:00:00 2001 From: Wagner Bruna Date: Fri, 21 Aug 2026 00:17:13 -0300 Subject: [PATCH] feat: support Brownian tree noise in all noise injection samplers --- examples/common/common.cpp | 2 +- src/runtime/denoiser.hpp | 384 +++++++++++++++++++------------------ 2 files changed, 195 insertions(+), 191 deletions(-) diff --git a/examples/common/common.cpp b/examples/common/common.cpp index 35812157c..84c8ed0f4 100644 --- a/examples/common/common.cpp +++ b/examples/common/common.cpp @@ -1008,7 +1008,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; flux supports base_shift, max_shift; ltx2 supports max_shift, base_shift, stretch, terminal; euler_ge supports gamma; beta scheduler supports alpha, beta; logit_normal supports mu, std, logsnr_min, logsnr_max, resolution_aware; lms supports lms_max_order, lms_shift, lms_divisions", + "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; beta scheduler supports alpha, beta; logit_normal supports mu, std, logsnr_min, logsnr_max, resolution_aware; lms supports lms_max_order, lms_shift, lms_divisions; ancestral and lcm samplers support noise_sampler with value iid (default for most samplers) or brownian_tree", (int)',', &extra_sample_args}, {"", diff --git a/src/runtime/denoiser.hpp b/src/runtime/denoiser.hpp index e4d9af020..6dec2c542 100644 --- a/src/runtime/denoiser.hpp +++ b/src/runtime/denoiser.hpp @@ -1510,12 +1510,18 @@ static std::tuple get_ancestral_step(float sigma_from, } } +class NoiseSampler { +public: + virtual sd::Tensor operator()(double sigma_from, double sigma_to) = 0; + virtual ~NoiseSampler() = default; +}; + static sd::Tensor sample_euler_ancestral(denoise_cb_t model, sd::Tensor x, const std::vector& sigmas, - std::shared_ptr rng = nullptr, - bool is_flow_denoiser = false, - float eta = 0.f) { + NoiseSampler& noise_sampler, + bool is_flow_denoiser = false, + float eta = 0.f) { int steps = static_cast(sigmas.size()) - 1; for (int i = 0; i < steps; i++) { float sigma = sigmas[i]; @@ -1538,7 +1544,7 @@ static sd::Tensor sample_euler_ancestral(denoise_cb_t model, if (is_flow_denoiser) { x *= alpha_scale; } - x += sd::Tensor::randn_like(x, rng) * sigma_up; + x += noise_sampler(sigma, sigma_to) * sigma_up; } } } @@ -1656,7 +1662,7 @@ static sd::Tensor sample_dpm2(denoise_cb_t model, static sd::Tensor sample_dpmpp_2s_ancestral(denoise_cb_t model, sd::Tensor x, const std::vector& sigmas, - std::shared_ptr rng, + NoiseSampler& noise_sampler, float eta) { auto t_fn = [](float sigma) -> float { return -log(sigma); }; auto sigma_fn = [](float t) -> float { return exp(-t); }; @@ -1688,7 +1694,7 @@ static sd::Tensor sample_dpmpp_2s_ancestral(denoise_cb_t model, } if (sigmas[i + 1] > 0) { - x += sd::Tensor::randn_like(x, rng) * sigma_up; + x += noise_sampler(sigmas[i], sigmas[i + 1]) * sigma_up; } } return x; @@ -1697,7 +1703,7 @@ static sd::Tensor sample_dpmpp_2s_ancestral(denoise_cb_t model, static sd::Tensor sample_dpmpp_2s_ancestral_flow(denoise_cb_t model, sd::Tensor x, const std::vector& sigmas, - std::shared_ptr rng, + NoiseSampler& noise_sampler, float eta = 1.0f) { int steps = static_cast(sigmas.size()) - 1; for (int i = 0; i < steps; i++) { @@ -1780,7 +1786,7 @@ static sd::Tensor sample_dpmpp_2s_ancestral_flow(denoise_cb_t model, x = (x * sigma_down_i_ratio) + (D_i * (1.0f - sigma_down_i_ratio)); if (sigma_to > 0.0f && eta > 0.0f) { - x = alpha_scale * x + sd::Tensor::randn_like(x, rng) * sigma_up; + x = alpha_scale * x + noise_sampler(sigma, sigma_to) * sigma_up; } } } @@ -1856,169 +1862,20 @@ static sd::Tensor sample_dpmpp_2m_v2(denoise_cb_t model, return x; } -// DPM-Solver++(2M) SDE, midpoint variant. Ref: Lu et al. arXiv:2211.01095; -// k-diffusion sample_dpmpp_2m_sde. +// DPM-Solver++(2M) SDE, midpoint variant. +// Ref: Lu et al. arXiv:2211.01095; k-diffusion sample_dpmpp_2m_sde +// step-count-stable Brownian-tree noise variant: same trajectory shape at any step count for a given seed. +// Aliased in k-diffusion / ComfyUI as sample_dpmpp_2m_sde_gpu. +// Ref: Lu et al. arXiv:2211.01095; torchsde BrownianTree. static sd::Tensor sample_dpmpp_2m_sde(denoise_cb_t model, sd::Tensor x, const std::vector& sigmas, - std::shared_ptr rng, + NoiseSampler& noise_sampler, float eta) { sd::Tensor old_denoised; bool have_old_denoised = false; float h_last = 0.f; - int steps = static_cast(sigmas.size()) - 1; - for (int i = 0; i < steps; i++) { - auto denoised_opt = model(x, sigmas[i], i + 1); - if (denoised_opt.pred.empty()) { - return {}; - } - sd::Tensor denoised = std::move(denoised_opt.pred); - - if (sigmas[i + 1] == 0.f) { - x = denoised; - } else { - float t = -std::log(sigmas[i]); - float s = -std::log(sigmas[i + 1]); - float h = s - t; - float eta_h = eta * h; - float a = sigmas[i + 1] / sigmas[i] * std::exp(-eta_h); - float b = -std::expm1(-h - eta_h); - - x = a * x + b * denoised; - - if (have_old_denoised) { - float r = h_last / h; - x += (0.5f * b / r) * (denoised - old_denoised); - } - if (eta > 0.f) { - x += sd::Tensor::randn_like(x, rng) * (sigmas[i + 1] * std::sqrt(-std::expm1(-2.f * eta_h))); - } - h_last = h; - } - old_denoised = denoised; - have_old_denoised = true; - } - return x; -} - -// Seeded Brownian tree providing deterministic, step-count-stable Gaussian -// increments for stochastic samplers. Constructed once per generation; each -// call returns unit-variance noise for interval [sigma_a, sigma_b]. -// Reference: torchsde BrownianTree; k-diffusion BatchedBrownianTree. -class BrownianTreeNoiseSampler { -public: - BrownianTreeNoiseSampler(const sd::Tensor& x_template, - double sigma_min, - double sigma_max, - uint64_t seed) - : t_min_(sigma_min), - t_max_(sigma_max), - shape_(x_template.shape()), - root_seed_(mix64(seed, 0x9E3779B97F4A7C15ULL)) { - auto rng = std::make_shared(); - rng->manual_seed(mix64(seed, 0xBF58476D1CE4E5B9ULL)); - w_at_tmax_ = sd::Tensor::randn(shape_, rng) * std::sqrt(static_cast(t_max_ - t_min_)); - } - - sd::Tensor operator()(double sigma_a, double sigma_b) { - double a = clamp(std::min(sigma_a, sigma_b)); - double b = clamp(std::max(sigma_a, sigma_b)); - auto dW = w(b) - w(a); - float span = static_cast(std::max(std::abs(sigma_b - sigma_a), 1e-12)); - return dW * (1.0f / std::sqrt(span)); - } - -private: - static constexpr int kMaxDepth = 24; - - static uint64_t mix64(uint64_t v, uint64_t salt) { - uint64_t z = v + salt; - z = (z ^ (z >> 30)) * 0xBF58476D1CE4E5B9ULL; - z = (z ^ (z >> 27)) * 0x94D049BB133111EBULL; - return z ^ (z >> 31); - } - - double clamp(double t) const { - return std::min(std::max(t, t_min_), t_max_); - } - - sd::Tensor w(double t) { - auto it = cache_.find(t); - if (it != cache_.end()) { - return it->second; - } - sd::Tensor zero = sd::Tensor::zeros(shape_); - sd::Tensor out = bridge(t_min_, t_max_, zero, w_at_tmax_, t, root_seed_, kMaxDepth); - cache_.emplace(t, out); - return out; - } - - sd::Tensor bridge(double a, - double c, - const sd::Tensor& w_a, - const sd::Tensor& w_c, - double t, - uint64_t node_seed, - int depth) { - if (depth <= 0 || c - a < 1e-9) { - float alpha = (c > a) ? static_cast((t - a) / (c - a)) : 0.5f; - return (1.0f - alpha) * w_a + alpha * w_c; - } - double m = 0.5 * (a + c); - double std_dev = std::sqrt((c - m) * (m - a) / (c - a)); - auto rng = std::make_shared(); - rng->manual_seed(node_seed); - auto z = sd::Tensor::randn(shape_, rng); - auto w_m = 0.5f * (w_a + w_c) + static_cast(std_dev) * z; - if (t == m) { - return w_m; - } - if (t < m) { - return bridge(a, m, w_a, w_m, t, mix64(node_seed, 1), depth - 1); - } - return bridge(m, c, w_m, w_c, t, mix64(node_seed, 2), depth - 1); - } - - double t_min_; - double t_max_; - std::vector shape_; - uint64_t root_seed_; - sd::Tensor w_at_tmax_; - std::map> cache_; -}; - -// DPM-Solver++(2M) SDE, midpoint variant, with step-count-stable Brownian-tree -// noise. Same trajectory shape at any step count for a given seed. Aliased in -// k-diffusion / ComfyUI as sample_dpmpp_2m_sde_gpu. -// Ref: Lu et al. arXiv:2211.01095; torchsde BrownianTree. -static sd::Tensor sample_dpmpp_2m_sde_bt(denoise_cb_t model, - sd::Tensor x, - const std::vector& sigmas, - std::shared_ptr rng, - float eta) { - double sigma_max = 0.0; - double sigma_min = std::numeric_limits::infinity(); - for (float s : sigmas) { - if (s > 0.0f) { - sigma_max = std::max(sigma_max, static_cast(s)); - sigma_min = std::min(sigma_min, static_cast(s)); - } - } - if (sigma_max <= sigma_min) { - return x; - } - uint64_t tree_seed = 0; - { - auto draw = rng->randn(2); - std::memcpy(&tree_seed, draw.data(), sizeof(tree_seed)); - } - BrownianTreeNoiseSampler noise_sampler(x, sigma_min, sigma_max, tree_seed); - - sd::Tensor old_denoised; - bool have_old_denoised = false; - float h_last = 0.f; - int steps = static_cast(sigmas.size()) - 1; for (int i = 0; i < steps; i++) { auto denoised_opt = model(x, sigmas[i], i + 1); @@ -2059,7 +1916,7 @@ using SamplerExtraArgs = KeyValueArgs; static sd::Tensor sample_lcm(denoise_cb_t model, sd::Tensor x, const std::vector& sigmas, - std::shared_ptr rng, + NoiseSampler& noise_sampler, bool is_flow_denoiser, const SamplerExtraArgs& extra_sample_args) { struct LCMSampleArgs { @@ -2112,7 +1969,7 @@ static sd::Tensor sample_lcm(denoise_cb_t model, if (is_flow_denoiser) { x *= (1 - sigmas[i + 1]); } - auto noise = sd::Tensor::randn_like(x, rng); + auto noise = noise_sampler(sigmas[i], sigmas[i + 1]); if (args.noise_clip_std > 0.0f && noise.numel() > 0) { double mean = 0.0; for (int64_t j = 0; j < noise.numel(); ++j) { @@ -2230,7 +2087,7 @@ static sd::Tensor sample_ipndm_v(denoise_cb_t model, static sd::Tensor sample_res_multistep(denoise_cb_t model, sd::Tensor x, const std::vector& sigmas, - std::shared_ptr rng, + NoiseSampler& noise_sampler, bool is_flow_denoiser, float eta) { sd::Tensor old_denoised = x; @@ -2295,7 +2152,7 @@ static sd::Tensor sample_res_multistep(denoise_cb_t model, if (is_flow_denoiser) { x *= alpha_scale; } - x += sd::Tensor::randn_like(x, rng) * sigma_up; + x += noise_sampler(sigma_from, sigma_to) * sigma_up; } old_denoised = denoised; @@ -2308,7 +2165,7 @@ static sd::Tensor sample_res_multistep(denoise_cb_t model, static sd::Tensor sample_res_2s(denoise_cb_t model, sd::Tensor x, const std::vector& sigmas, - std::shared_ptr rng, + NoiseSampler& noise_sampler, bool is_flow_denoiser, float eta) { const float c2 = 0.5f; @@ -2371,7 +2228,7 @@ static sd::Tensor sample_res_2s(denoise_cb_t model, if (is_flow_denoiser) { x *= alpha_scale; } - x += sd::Tensor::randn_like(x, rng) * sigma_up; + x += noise_sampler(sigma_from, sigma_to) * sigma_up; } } return x; @@ -2380,7 +2237,7 @@ static sd::Tensor sample_res_2s(denoise_cb_t model, static sd::Tensor sample_er_sde(denoise_cb_t model, sd::Tensor x, std::vector sigmas, - std::shared_ptr rng, + NoiseSampler& noise_sampler, bool is_flow_denoiser, float eta) { constexpr int max_stage = 3; @@ -2502,7 +2359,7 @@ static sd::Tensor sample_er_sde(denoise_cb_t model, float noise_scale_sq = er_lambda_t * er_lambda_t - er_lambda_s * er_lambda_s * r * r; if (s_noise > 0.0f && noise_scale_sq > 0.0f) { float noise_scale = alpha_t * std::sqrt(std::max(noise_scale_sq, 0.0f)); - x += sd::Tensor::randn_like(x, rng) * noise_scale; + x += noise_sampler(sigmas[i], sigmas[i + 1]) * noise_scale; } } @@ -2515,7 +2372,7 @@ static sd::Tensor sample_er_sde(denoise_cb_t model, static sd::Tensor sample_tcd(denoise_cb_t model, sd::Tensor x, const std::vector& sigmas, - std::shared_ptr rng, + NoiseSampler& noise_sampler, float eta) { float beta_start = 0.00085f; float beta_end = 0.0120f; @@ -2572,7 +2429,7 @@ static sd::Tensor sample_tcd(denoise_cb_t model, if (eta > 0 && sigma_to > 0.0f) { x = std::sqrt(alpha_prod_t_prev / alpha_prod_s) * x + - std::sqrt(1.0f / alpha_prod_t_prev - 1.0f / alpha_prod_s) * sd::Tensor::randn_like(x, rng); + std::sqrt(1.0f / alpha_prod_t_prev - 1.0f / alpha_prod_s) * noise_sampler(sigma, sigma_to); } } return x; @@ -2707,7 +2564,7 @@ static sd::Tensor sample_euler_cfg_pp(denoise_cb_t model, static sd::Tensor sample_euler_ancestral_cfg_pp(denoise_cb_t model, sd::Tensor x, const std::vector& sigmas, - std::shared_ptr rng, + NoiseSampler& noise_sampler, float eta) { int steps = static_cast(sigmas.size()) - 1; for (int i = 0; i < steps; i++) { @@ -2726,7 +2583,7 @@ static sd::Tensor sample_euler_ancestral_cfg_pp(denoise_cb_t model, x = denoised + d * sigma_down; if (sigmas[i + 1] > 0) { - x += sd::Tensor::randn_like(x, rng) * sigma_up; + x += noise_sampler(sigmas[i], sigmas[i + 1]) * sigma_up; } } return x; @@ -2736,7 +2593,7 @@ static sd::Tensor sample_euler_ancestral_cfg_pp(denoise_cb_t model, static sd::Tensor sample_gradient_estimation(denoise_cb_t model, sd::Tensor x, const std::vector& sigmas, - std::shared_ptr rng, + NoiseSampler& noise_sampler, bool is_flow_denoiser, float eta, const SamplerExtraArgs& extra_sample_args) { @@ -2783,13 +2640,158 @@ static sd::Tensor sample_gradient_estimation(denoise_cb_t model, if (is_flow_denoiser) { x *= alpha_scale; } - x += sd::Tensor::randn_like(x, rng) * sigma_up; + x += noise_sampler(sigma, sigma_to) * sigma_up; } } } return x; } +// independent and identically distributed Gaussian noise (default for most samplers) +class IIDGaussianNoiseSampler : public NoiseSampler { +public: + IIDGaussianNoiseSampler(const sd::Tensor& x_template, std::shared_ptr r) + : rng(std::move(r)), shape(x_template.shape()) {} + sd::Tensor operator()(double sigma_from, double sigma_to) override { + (void)sigma_from; + (void)sigma_to; + return sd::Tensor::randn(shape, rng); + } + +private: + std::shared_ptr rng; + std::vector shape; +}; + +// Seeded Brownian tree providing deterministic, step-count-stable Gaussian +// increments for stochastic samplers. Constructed once per generation; each +// call returns unit-variance noise for interval [sigma_a, sigma_b]. +// Reference: torchsde BrownianTree; k-diffusion BatchedBrownianTree. +class BrownianTreeNoiseSampler : public NoiseSampler { +public: + BrownianTreeNoiseSampler(const sd::Tensor& x_template, + double sigma_min, + double sigma_max, + uint64_t seed) + : t_min_(sigma_min), + t_max_(sigma_max), + shape_(x_template.shape()), + root_seed_(mix64(seed, 0x9E3779B97F4A7C15ULL)) { + auto rng = std::make_shared(); + rng->manual_seed(mix64(seed, 0xBF58476D1CE4E5B9ULL)); + w_at_tmax_ = sd::Tensor::randn(shape_, rng) * std::sqrt(static_cast(t_max_ - t_min_)); + } + + sd::Tensor operator()(double sigma_a, double sigma_b) override { + double a = clamp(std::min(sigma_a, sigma_b)); + double b = clamp(std::max(sigma_a, sigma_b)); + auto dW = w(b) - w(a); + float span = static_cast(std::max(std::abs(sigma_b - sigma_a), 1e-12)); + return dW * (1.0f / std::sqrt(span)); + } + +private: + static constexpr int kMaxDepth = 24; + + static uint64_t mix64(uint64_t v, uint64_t salt) { + uint64_t z = v + salt; + z = (z ^ (z >> 30)) * 0xBF58476D1CE4E5B9ULL; + z = (z ^ (z >> 27)) * 0x94D049BB133111EBULL; + return z ^ (z >> 31); + } + + double clamp(double t) const { + return std::min(std::max(t, t_min_), t_max_); + } + + sd::Tensor w(double t) { + auto it = cache_.find(t); + if (it != cache_.end()) { + return it->second; + } + sd::Tensor zero = sd::Tensor::zeros(shape_); + sd::Tensor out = bridge(t_min_, t_max_, zero, w_at_tmax_, t, root_seed_, kMaxDepth); + cache_.emplace(t, out); + return out; + } + + sd::Tensor bridge(double a, + double c, + const sd::Tensor& w_a, + const sd::Tensor& w_c, + double t, + uint64_t node_seed, + int depth) { + if (depth <= 0 || c - a < 1e-9) { + float alpha = (c > a) ? static_cast((t - a) / (c - a)) : 0.5f; + return (1.0f - alpha) * w_a + alpha * w_c; + } + double m = 0.5 * (a + c); + double std_dev = std::sqrt((c - m) * (m - a) / (c - a)); + auto rng = std::make_shared(); + rng->manual_seed(node_seed); + auto z = sd::Tensor::randn(shape_, rng); + auto w_m = 0.5f * (w_a + w_c) + static_cast(std_dev) * z; + if (t == m) { + return w_m; + } + if (t < m) { + return bridge(a, m, w_a, w_m, t, mix64(node_seed, 1), depth - 1); + } + return bridge(m, c, w_m, w_c, t, mix64(node_seed, 2), depth - 1); + } + + double t_min_; + double t_max_; + std::vector shape_; + uint64_t root_seed_; + sd::Tensor w_at_tmax_; + std::map> cache_; +}; + +static std::unique_ptr make_noise_sampler(const sd::Tensor& x, std::shared_ptr rng, sample_method_t method, const std::vector& sigmas, const SamplerExtraArgs& extra_args) { + bool brownian_tree = (method == DPMPP2M_SDE_BT_SAMPLE_METHOD); + bool def_brownian_tree = brownian_tree; + + for (const auto& [key, value] : extra_args) { + if (key == "noise_sampler") { + if (value == "iid") { + brownian_tree = false; + } else if (value == "brownian_tree") { + brownian_tree = true; + } else { + LOG_WARN("unknown noise_sampler value '%s'; using default", value.c_str()); + } + } + } + + if (brownian_tree) { + double sigma_max = 0.0; + double sigma_min = std::numeric_limits::infinity(); + for (float s : sigmas) { + if (s > 0.0f) { + sigma_max = std::max(sigma_max, static_cast(s)); + sigma_min = std::min(sigma_min, static_cast(s)); + } + } + + if (sigma_max > sigma_min) { + uint64_t tree_seed = 0; + auto draw = rng->randn(2); + std::memcpy(&tree_seed, draw.data(), sizeof(tree_seed)); + if (!def_brownian_tree) { + LOG_INFO("setting noise sampler to Brownian tree"); + } + return std::make_unique(x, sigma_min, sigma_max, tree_seed); + } + } + + if (def_brownian_tree) { + LOG_INFO("setting noise sampler to independent and identically distributed (iid)"); + } + return std::make_unique(x, rng); +} + // 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, @@ -2806,9 +2808,12 @@ static sd::Tensor sample_k_diffusion(sample_method_t method, } } SamplerExtraArgs extra_args = parse_key_value_args(extra_sample_args, "extra sample arg"); + + std::unique_ptr noise_sampler = make_noise_sampler(x, rng, method, sigmas, extra_args); + switch (method) { case EULER_A_SAMPLE_METHOD: - return sample_euler_ancestral(model, std::move(x), sigmas, rng, is_flow_denoiser, eta); + return sample_euler_ancestral(model, std::move(x), sigmas, *noise_sampler, is_flow_denoiser, eta); case EULER_SAMPLE_METHOD: return sample_euler(model, std::move(x), sigmas); case HEUN_SAMPLE_METHOD: @@ -2817,42 +2822,41 @@ static sd::Tensor sample_k_diffusion(sample_method_t method, return sample_dpm2(model, std::move(x), sigmas); case DPMPP2S_A_SAMPLE_METHOD: if (is_flow_denoiser) - return sample_dpmpp_2s_ancestral_flow(model, std::move(x), sigmas, rng, eta); + return sample_dpmpp_2s_ancestral_flow(model, std::move(x), sigmas, *noise_sampler, eta); else - return sample_dpmpp_2s_ancestral(model, std::move(x), sigmas, rng, eta); + return sample_dpmpp_2s_ancestral(model, std::move(x), sigmas, *noise_sampler, eta); case DPMPP2M_SAMPLE_METHOD: return sample_dpmpp_2m(model, std::move(x), sigmas); 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, extra_args); + return sample_lcm(model, std::move(x), sigmas, *noise_sampler, is_flow_denoiser, extra_args); case IPNDM_SAMPLE_METHOD: return sample_ipndm(model, std::move(x), sigmas); case IPNDM_V_SAMPLE_METHOD: return sample_ipndm_v(model, std::move(x), sigmas); case RES_MULTISTEP_SAMPLE_METHOD: - return sample_res_multistep(model, std::move(x), sigmas, rng, is_flow_denoiser, eta); + return sample_res_multistep(model, std::move(x), sigmas, *noise_sampler, is_flow_denoiser, eta); case RES_2S_SAMPLE_METHOD: - return sample_res_2s(model, std::move(x), sigmas, rng, is_flow_denoiser, eta); + return sample_res_2s(model, std::move(x), sigmas, *noise_sampler, is_flow_denoiser, eta); case ER_SDE_SAMPLE_METHOD: - return sample_er_sde(model, std::move(x), sigmas, rng, is_flow_denoiser, eta); + return sample_er_sde(model, std::move(x), sigmas, *noise_sampler, is_flow_denoiser, eta); case DPMPP2M_SDE_SAMPLE_METHOD: - return sample_dpmpp_2m_sde(model, std::move(x), sigmas, rng, eta); case DPMPP2M_SDE_BT_SAMPLE_METHOD: - return sample_dpmpp_2m_sde_bt(model, std::move(x), sigmas, rng, eta); + return sample_dpmpp_2m_sde(model, std::move(x), sigmas, *noise_sampler, eta); case DDIM_TRAILING_SAMPLE_METHOD: // DDIM is equivalent to Euler Ancestral with the Simple scheduler - return sample_euler_ancestral(model, std::move(x), sigmas, rng, is_flow_denoiser, eta); + return sample_euler_ancestral(model, std::move(x), sigmas, *noise_sampler, is_flow_denoiser, eta); case TCD_SAMPLE_METHOD: - return sample_tcd(model, std::move(x), sigmas, rng, eta); + return sample_tcd(model, std::move(x), sigmas, *noise_sampler, eta); case LMS_SAMPLE_METHOD: return sample_lms(model, std::move(x), sigmas, extra_args); case EULER_CFG_PP_SAMPLE_METHOD: return sample_euler_cfg_pp(model, std::move(x), sigmas); case EULER_A_CFG_PP_SAMPLE_METHOD: - return sample_euler_ancestral_cfg_pp(model, std::move(x), sigmas, rng, eta); + return sample_euler_ancestral_cfg_pp(model, std::move(x), sigmas, *noise_sampler, eta); case EULER_GE_SAMPLE_METHOD: - return sample_gradient_estimation(model, std::move(x), sigmas, rng, is_flow_denoiser, eta, extra_args); + return sample_gradient_estimation(model, std::move(x), sigmas, *noise_sampler, is_flow_denoiser, eta, extra_args); default: return {}; }