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 <noreply@anthropic.com>

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

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)]