feat: support Brownian tree noise in all noise injection samplers

[wbruna]

Summary

Adds a new extra sampling parameter noise_sampler=brownian_tree to optionally use Brownian tree noise in LCM and all ancestral samplers.

Related Issue / Discussion

#1743

Checklist

• I have read and confirmed this PR follows the contribution guidelines.

[vmobilis]

@wbruna, maybe it's possible to make it more generic – as an option for --sampler-rng or even --rng?

Because it looks like noise_sampler is duplicating their functionality.

Probably is not so good idea, because it does not substitute, but rather extends the rng.

[vmobilis]

@wbruna, hello, I digged a little more, currently the BrownianTreeNoiseSampler() explicitly uses std::make_shared<STDDefaultRNG>():

https://github.com/wbruna/stable-diffusion.cpp/blob/sd_brownian_tree/src/runtime/denoiser.hpp#L2724

auto rng = std::make_shared<STDDefaultRNG>();

(Which is the std option of --sampler-rng):
https://github.com/wbruna/stable-diffusion.cpp/blob/sd_brownian_tree/src/stable-diffusion.cpp#L659

if (rng_type == STD_DEFAULT_RNG) {
    return std::make_shared<STDDefaultRNG>();
}

Can you pass the application-wide rng, as in IIDGaussianNoiseSampler()?

https://github.com/wbruna/stable-diffusion.cpp/blob/sd_brownian_tree/src/runtime/denoiser.hpp#L2702

return sd::Tensor<float>::randn(shape, rng);

Then it should allow to use --sampler-rng choice of std, cpu and cuda in Brownian extension.

[wbruna]

Sounds sensible to me. But maybe @fszontagh adopted it to match results from ComfyUI? If so, we may want to keep it as the default for dpm++2m_sde_bt? (but then the rng choice would need to be moved from model loading time to an inference parameter - which also sounds sensible to me, but maybe a bit too much for this PR 🙂)

[vmobilis]

@wbruna, the std option should pass exactly the same <STDDefaultRNG> as currently is in BrownianTreeNoiseSampler(), so with --sampler-rng std the noise_sampler=brownian_tree should behave exactly as before, if I'm not mistaken.

And probably there is no need to move anything, because you are already using the app-wide rng in IIDGaussianNoiseSampler() and only need to pass it to BrownianTreeNoiseSampler() as well.

[fszontagh]

@wbruna Not a ComfyUI parity thing - mine is a hand-written bridge, ComfyUI uses torchsde.BrownianTree, so the noise never matched anyway. No need to keep STDDefaultRNG as the default for dpm++2m_sde_bt.

One caveat: bridge() re-seeds per node, so passing the shared app RNG would re-seed it mid-sampling. Better to clone it. The hardcoded std also keeps the tree device-independent, which is worth preserving.

[vmobilis]

@wbruna, hello, maybe something like wbruna#3?

I didn't test all variants yet, but so far iid and brownian_tree look same and others produce different images.

The changes are minimal, the rng in BrownianTreeNoiseSampler() is copied, as @fszontagh suggested.

P.s. If you don't need it, just ignore, I've simply shown what I meant. 😉

[leejet]

Thanks for the discussion. I’ve added brownian_tree_rng with these options:

• cpu (default)
• cuda
• std_default
• sampler_rng

Each tree owns an independent RNG. sampler_rng clones the selected sampler RNG once, so reseeding tree nodes cannot reset the shared RNG.

The default is now cpu, using the same host-side implementation regardless of the inference backend. This changes existing Brownian outputs; std_default preserves the previous node-noise generation. The Brownian bridge and tree-seed derivation remain unchanged.

I also made tree initialization lazy, so samplers that never request noise no longer consume RNG state unnecessarily.