diff --git a/src/core/rng.hpp b/src/core/rng.hpp index 20ffec7de..e2eb694de 100644 --- a/src/core/rng.hpp +++ b/src/core/rng.hpp @@ -4,17 +4,29 @@ #include #include +#include "stable-diffusion.h" // for SD_API, str_to_rng_type() in denoiser + class RNG { public: virtual void manual_seed(uint64_t seed) = 0; virtual std::vector randn(uint32_t n) = 0; + + virtual const char* const rn() const { return "rng"; } + virtual const std::shared_ptr clone() const = 0; }; +extern SD_API std::shared_ptr get_rng(rng_type_t rng_type); + class STDDefaultRNG : public RNG { private: std::default_random_engine generator; public: + virtual const char* const rn() const override { return "std"; } + virtual const std::shared_ptr clone() const override { + return std::make_shared(*this); + } + void manual_seed(uint64_t seed) override { generator.seed((unsigned int)seed); } diff --git a/src/core/rng_mt19937.hpp b/src/core/rng_mt19937.hpp index 85da05cfe..b1fd282c3 100644 --- a/src/core/rng_mt19937.hpp +++ b/src/core/rng_mt19937.hpp @@ -123,6 +123,11 @@ class MT19937RNG : public RNG { public: MT19937RNG(uint64_t seed = 0) { manual_seed(seed); } + virtual const char* const rn() const override { return "cpu"; } + virtual const std::shared_ptr clone() const override { + return std::make_shared(*this); + } + void manual_seed(uint64_t seed) override { s.seed_ = seed; s.seeded_ = true; diff --git a/src/core/rng_philox.hpp b/src/core/rng_philox.hpp index 8258ed950..8f8c7aee9 100644 --- a/src/core/rng_philox.hpp +++ b/src/core/rng_philox.hpp @@ -93,6 +93,11 @@ class PhiloxRNG : public RNG { this->offset = 0; } + virtual const char* const rn() const override { return "cuda"; } + virtual const std::shared_ptr clone() const override { + return std::make_shared(*this); + } + void manual_seed(uint64_t seed) override { this->seed = seed; this->offset = 0; diff --git a/src/runtime/denoiser.hpp b/src/runtime/denoiser.hpp index a9bf91699..be7afa023 100644 --- a/src/runtime/denoiser.hpp +++ b/src/runtime/denoiser.hpp @@ -16,6 +16,10 @@ #include "runtime/gits_noise.h" #include "runtime/guidance.h" +// #include "core/rng.hpp" // included in core/tensor.hpp +#include "core/rng_mt19937.hpp" +#include "core/rng_philox.hpp" + /*================================================= CompVisDenoiser ==================================================*/ // Ref: https://github.com/crowsonkb/k-diffusion/blob/master/k_diffusion/external.py @@ -2714,14 +2718,16 @@ class IIDGaussianNoiseSampler : public NoiseSampler { class BrownianTreeNoiseSampler : public NoiseSampler { public: BrownianTreeNoiseSampler(const sd::Tensor& x_template, + std::shared_ptr r, double sigma_min, double sigma_max, uint64_t seed) : t_min_(sigma_min), t_max_(sigma_max), + app_wide_rng(std::move(r)), shape_(x_template.shape()), root_seed_(mix64(seed, 0x9E3779B97F4A7C15ULL)) { - auto rng = std::make_shared(); + auto rng = app_wide_rng->clone(); rng->manual_seed(mix64(seed, 0xBF58476D1CE4E5B9ULL)); w_at_tmax_ = sd::Tensor::randn(shape_, rng) * std::sqrt(static_cast(t_max_ - t_min_)); } @@ -2754,12 +2760,13 @@ class BrownianTreeNoiseSampler : public NoiseSampler { 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); + sd::Tensor out = bridge(app_wide_rng, t_min_, t_max_, zero, w_at_tmax_, t, root_seed_, kMaxDepth); cache_.emplace(t, out); return out; } - sd::Tensor bridge(double a, + sd::Tensor bridge(std::shared_ptr r, + double a, double c, const sd::Tensor& w_a, const sd::Tensor& w_c, @@ -2772,7 +2779,7 @@ class BrownianTreeNoiseSampler : public NoiseSampler { } double m = 0.5 * (a + c); double std_dev = std::sqrt((c - m) * (m - a) / (c - a)); - auto rng = std::make_shared(); + auto rng = r->clone(); 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; @@ -2780,13 +2787,14 @@ class BrownianTreeNoiseSampler : public NoiseSampler { return w_m; } if (t < m) { - return bridge(a, m, w_a, w_m, t, mix64(node_seed, 1), depth - 1); + return bridge(r, 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); + return bridge(r, m, c, w_m, w_c, t, mix64(node_seed, 2), depth - 1); } double t_min_; double t_max_; + const std::shared_ptr app_wide_rng; std::vector shape_; uint64_t root_seed_; sd::Tensor w_at_tmax_; @@ -2796,6 +2804,7 @@ class BrownianTreeNoiseSampler : public NoiseSampler { 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; + int bt_rng_type = STD_DEFAULT_RNG; // int instead of enum because of RNG_TYPE_COUNT + 1 for (const auto& [key, value] : extra_args) { if (key == "noise_sampler") { @@ -2807,6 +2816,16 @@ static std::unique_ptr make_noise_sampler(const sd::Tensor& LOG_WARN("unknown noise_sampler value '%s'; using default", value.c_str()); } } + if (key == "brownian_tree_rng") { + if (value == "sampler_rng") { + bt_rng_type = RNG_TYPE_COUNT + 1; + } else { + bt_rng_type = str_to_rng_type(value.c_str()); + if (bt_rng_type == RNG_TYPE_COUNT) { + LOG_WARN("invalid '%s' type '%s'; using std_default", key.c_str(), value.c_str()); + } + } + } } if (brownian_tree) { @@ -2823,10 +2842,22 @@ static std::unique_ptr make_noise_sampler(const sd::Tensor& uint64_t tree_seed = 0; auto draw = rng->randn(2); std::memcpy(&tree_seed, draw.data(), sizeof(tree_seed)); + + std::shared_ptr r; + switch(bt_rng_type) { + case RNG_TYPE_COUNT + 1: + r = rng; + break; + case RNG_TYPE_COUNT: + r = get_rng(STD_DEFAULT_RNG); + break; + default: + r = get_rng((rng_type_t) bt_rng_type); + } if (!def_brownian_tree) { - LOG_INFO("setting noise sampler to Brownian tree"); + LOG_INFO("setting noise sampler to Brownian tree (%s)", r->rn()); } - return std::make_unique(x, sigma_min, sigma_max, tree_seed); + return std::make_unique(x, std::move(r), sigma_min, sigma_max, tree_seed); } } diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index d2193e14c..b635eb3bf 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -654,16 +654,6 @@ class StableDiffusionGGML { return false; } - std::shared_ptr get_rng(rng_type_t rng_type) { - if (rng_type == STD_DEFAULT_RNG) { - return std::make_shared(); - } else if (rng_type == CPU_RNG) { - return std::make_shared(); - } else { // default: CUDA_RNG - return std::make_shared(); - } - } - void refresh_compvis_denoiser_sigmas() { auto comp_vis_denoiser = std::dynamic_pointer_cast(denoiser); if (!comp_vis_denoiser) { @@ -3302,6 +3292,16 @@ enum rng_type_t str_to_rng_type(const char* str) { return RNG_TYPE_COUNT; } +std::shared_ptr get_rng(rng_type_t rng_type) { + if (rng_type == STD_DEFAULT_RNG) { + return std::make_shared(); + } else if (rng_type == CPU_RNG) { + return std::make_shared(); + } else { // default: CUDA_RNG + return std::make_shared(); + } +} + const char* sample_method_to_str[] = { "euler", "euler_a",