feat: add options for brownian_tree

src/core/rng.hpp

@@ -4,17 +4,29 @@
 #include <random>
 #include <vector>
 
+#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<float> randn(uint32_t n) = 0;
+
+    virtual const char* const rn() const { return "rng"; }
+    virtual const std::shared_ptr<RNG> clone() const = 0;
 };
 
+extern SD_API std::shared_ptr<RNG> 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<RNG> clone() const override {
+        return std::make_shared<STDDefaultRNG>(*this);
+    }
+
     void manual_seed(uint64_t seed) override {
         generator.seed((unsigned int)seed);
     }

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<RNG> clone() const override {
+        return std::make_shared<MT19937RNG>(*this);
+    }
+
     void manual_seed(uint64_t seed) override {
         s.seed_     = seed;
         s.seeded_   = true;

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<RNG> clone() const override {
+        return std::make_shared<PhiloxRNG>(*this);
+    }
+
     void manual_seed(uint64_t seed) override {
         this->seed   = seed;
         this->offset = 0;

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<float>& x_template,
+                             std::shared_ptr<RNG> 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<STDDefaultRNG>();
+        auto rng = app_wide_rng->clone();
         rng->manual_seed(mix64(seed, 0xBF58476D1CE4E5B9ULL));
         w_at_tmax_ = sd::Tensor<float>::randn(shape_, rng) * std::sqrt(static_cast<float>(t_max_ - t_min_));
     }
@@ -2754,12 +2760,13 @@ class BrownianTreeNoiseSampler : public NoiseSampler {
             return it->second;
         }
         sd::Tensor<float> zero = sd::Tensor<float>::zeros(shape_);
-        sd::Tensor<float> out  = bridge(t_min_, t_max_, zero, w_at_tmax_, t, root_seed_, kMaxDepth);
+        sd::Tensor<float> 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<float> bridge(double a,
+    sd::Tensor<float> bridge(std::shared_ptr<RNG> r,
+                             double a,
                              double c,
                              const sd::Tensor<float>& w_a,
                              const sd::Tensor<float>& w_c,
@@ -2772,21 +2779,22 @@ 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<STDDefaultRNG>();
+        auto rng = r->clone();
         rng->manual_seed(node_seed);
         auto z   = sd::Tensor<float>::randn(shape_, rng);
         auto w_m = 0.5f * (w_a + w_c) + static_cast<float>(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(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<RNG> app_wide_rng;
     std::vector<int64_t> shape_;
     uint64_t root_seed_;
     sd::Tensor<float> w_at_tmax_;
@@ -2796,6 +2804,7 @@ class BrownianTreeNoiseSampler : public NoiseSampler {
 static std::unique_ptr<NoiseSampler> make_noise_sampler(const sd::Tensor<float>& x, std::shared_ptr<RNG> rng, sample_method_t method, const std::vector<float>& 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<NoiseSampler> make_noise_sampler(const sd::Tensor<float>&
                 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<NoiseSampler> make_noise_sampler(const sd::Tensor<float>&
             uint64_t tree_seed = 0;
             auto draw          = rng->randn(2);
             std::memcpy(&tree_seed, draw.data(), sizeof(tree_seed));
+
+            std::shared_ptr<RNG> 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<BrownianTreeNoiseSampler>(x, sigma_min, sigma_max, tree_seed);
+            return std::make_unique<BrownianTreeNoiseSampler>(x, std::move(r), sigma_min, sigma_max, tree_seed);
         }
     }
 

src/stable-diffusion.cpp

@@ -654,16 +654,6 @@ class StableDiffusionGGML {
         return false;
     }
 
-    std::shared_ptr<RNG> get_rng(rng_type_t rng_type) {
-        if (rng_type == STD_DEFAULT_RNG) {
-            return std::make_shared<STDDefaultRNG>();
-        } else if (rng_type == CPU_RNG) {
-            return std::make_shared<MT19937RNG>();
-        } else {  // default: CUDA_RNG
-            return std::make_shared<PhiloxRNG>();
-        }
-    }
-
     void refresh_compvis_denoiser_sigmas() {
         auto comp_vis_denoiser = std::dynamic_pointer_cast<CompVisDenoiser>(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<RNG> get_rng(rng_type_t rng_type) {
+    if (rng_type == STD_DEFAULT_RNG) {
+        return std::make_shared<STDDefaultRNG>();
+    } else if (rng_type == CPU_RNG) {
+        return std::make_shared<MT19937RNG>();
+    } else {  // default: CUDA_RNG
+        return std::make_shared<PhiloxRNG>();
+    }
+}
+
 const char* sample_method_to_str[] = {
     "euler",
     "euler_a",