From 4faea24bcf481d474c59686bca407da4e64b7141 Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Tue, 10 Nov 2026 00:10:01 +0200 Subject: [PATCH] SDE samplers: torchsde's Brownian start warnings stay silent even when the widened noise sampler cannot be used Co-Authored-By: Claude Sonnet 5.5 --- core/brownian.py | 5 +++++ core/tests/test_brownian.py | 15 +++++++++++++++ 2 files changed, 20 insertions(+) diff --git a/core/brownian.py b/core/brownian.py index 407356b8..92439883 100644 --- a/core/brownian.py +++ b/core/brownian.py @@ -10,9 +10,13 @@ The class below builds the tree exactly as ComfyUI does (so the noise of a seed query to the tree's own span BEFORE torchsde sees it, which is the very same limit torchsde would apply, silently. """ +import warnings + import torch _installed = False +# torchsde clamps such a query itself, so the line is noise whenever the widened sampler below cannot be used or is bypassed. +_START_WARNING = r"Should have t[ab]" def _internals_present(ksampling) -> bool: @@ -41,6 +45,7 @@ def _log_skipped(): def install(): """Swap ComfyUI's BrownianTreeNoiseSampler for the widened one. Safe to call twice.""" global _installed + warnings.filterwarnings("ignore", message=_START_WARNING) import comfy.k_diffusion.sampling as ksampling if _installed or getattr(ksampling.BrownianTreeNoiseSampler, "_funui_widened", False): _installed = True diff --git a/core/tests/test_brownian.py b/core/tests/test_brownian.py index 1ee5aadf..92a1f7fa 100644 --- a/core/tests/test_brownian.py +++ b/core/tests/test_brownian.py @@ -123,3 +123,18 @@ def test_a_torchsde_with_other_internals_leaves_the_stock_sampler_in_place(monke def test_the_real_internals_are_recognised(): assert _brownian()._internals_present(ksampling) is True + + +def test_the_start_warning_is_silent_even_when_the_stock_sampler_stays(monkeypatch): + original = ksampling.BrownianTreeNoiseSampler + while getattr(original, "_funui_widened", False): + original = original.__mro__[1] + monkeypatch.setattr(ksampling, "BrownianTreeNoiseSampler", original) + mod = _brownian() + monkeypatch.setattr(mod, "_installed", False) + monkeypatch.setattr(mod, "_internals_present", lambda k: False) + sig = torch.tensor([0.9876543, 0.6, 0.4465445876]) + with warnings.catch_warnings(record=True) as seen: + mod.install() + original(torch.zeros(1, 2, 3), sig.min(), sig.max(), seed=7, cpu=True)(sig[0], sig[2]) + assert not [w for w in seen if "Should have" in str(w.message)]