Round 3: Brownian noise equals stock noise at any sigma height; every LoRA wrapper reaches every Diffusers module

Queries are limited as plain numbers inside the sampler call. Tests for 11 wrappers and sigma 14.6.

core/brownian.py

@@ -6,8 +6,8 @@ against the ROUNDED start. A schedule whose smallest sigma has a 7th decimal of
 `Should have ta>=t0 but got ta=... and t0=...`, then clamps the query to the start itself. The picture is not hurt,
 but the log fills with warnings on every SDE-type sampler, and Two stages / Turbo hand such sigmas to the sampler.
 
-The class below builds the tree exactly as ComfyUI does (so the noise of a seed is the stock noise) and clamps a
-query to the tree's own span BEFORE torchsde sees it, which is the very same clamp torchsde would apply, silently.
+The class below builds the tree exactly as ComfyUI does (so the noise of a seed is the stock noise) and limits a
+query to the tree's own span BEFORE torchsde sees it, which is the very same limit torchsde would apply, silently.
 """
 
 import torch
@@ -15,29 +15,6 @@ import torch
 _installed = False
 
 
-def _inside32(lo, hi):
-    """(lo, hi) moved inwards to the nearest float32 values: ComfyUI's CPU trees cast a query to float32, which
-    must not land outside the span again."""
-    import numpy as np
-    a, b = np.float32(lo), np.float32(hi)
-    if float(a) < lo:
-        a = np.nextafter(a, np.float32(np.inf))
-    if float(b) > hi:
-        b = np.nextafter(b, np.float32(-np.inf))
-    return float(a), float(b)
-
-
-def clamped(t, lo, hi):
-    """`t` limited to [lo, hi] (python floats) on t's own device and in t's dtype (float32 when t is narrower, so the
-    bound is not rounded back out of the span); a value already inside comes back untouched."""
-    v = float(t)
-    if lo <= v <= hi:
-        return t
-    lo32, hi32 = _inside32(lo, hi)
-    dtype = t.dtype if t.dtype in (torch.float32, torch.float64) else torch.float32
-    return torch.tensor(min(max(v, lo32), hi32), dtype=dtype, device=t.device)
-
-
 def install():
     """Swap ComfyUI's BrownianTreeNoiseSampler for the widened one. Safe to call twice."""
     global _installed
@@ -51,11 +28,20 @@ def install():
         _funui_widened = True
 
         def __call__(self, sigma, sigma_next):
+            """ComfyUI's own call, with one difference: a query outside the tree's span is limited to it HERE, as a plain
+            number, instead of inside torchsde (which does the same limiting, after a warning)."""
             t0, t1 = self.transform(torch.as_tensor(sigma)), self.transform(torch.as_tensor(sigma_next))
-            tree = self.tree.trees[0]._interval
-            lo, hi = float(tree._start), float(tree._end)
-            noise = self.tree(clamped(t0, lo, hi), clamped(t1, lo, hi)).to(dtype=t0.dtype)      # the dtype the stock sampler hands back
-            return noise / (t1 - t0).abs().sqrt()
+            batch = self.tree
+            a, b, sign = batch.sort(t0, t1)
+            device, dtype = a.device, a.dtype
+            span = batch.trees[0]._interval
+            lo, hi = float(span._start), float(span._end)
+
+            def query(t):                                  # the same cast ComfyUI applies to the query of a CPU tree
+                v = float(t.detach().cpu().float()) if batch.cpu_tree else float(t)
+                return min(max(v, lo), hi)
+            w = torch.stack([tree(query(a), query(b)) for tree in batch.trees]).to(device=device, dtype=dtype) * (batch.sign * sign)
+            return (w if batch.batched else w[0]) / (t1 - t0).abs().sqrt()
 
     ksampling.BrownianTreeNoiseSampler = BrownianTreeNoiseSampler
     _installed = True

core/tests/test_brownian.py

@@ -67,19 +67,25 @@ def test_the_noise_of_a_seed_is_the_stock_noise(patched):
         assert all(torch.equal(a, b) for a, b in zip(want, got))
 
 
-def test_a_query_is_clamped_in_the_dtype_it_meets_and_stays_inside_after_float32():
-    mod = _brownian()
-    for dtype in (torch.float32, torch.float64, torch.bfloat16):
-        t = torch.tensor(0.1, dtype=dtype)
-        out = mod.clamped(t, 0.4465455, 0.99)
-        assert out.device == t.device and out.dtype == (dtype if dtype != torch.bfloat16 else torch.float32)
-        assert float(out.float()) >= 0.4465455                  # the cast a CPU tree applies cannot put it back outside
-    inside = torch.tensor(0.6)
-    assert mod.clamped(inside, 0.446545, 0.99) is inside
-    meta = torch.empty((), device="meta")
-    assert meta.device.type == "meta"                          # the helper reads the value, so a meta tensor is refused loudly, not guessed
-    with pytest.raises(Exception):
-        mod.clamped(meta, 0.1, 0.9)
+@pytest.mark.parametrize("cpu", [True, False])
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float64])
+@pytest.mark.parametrize("top", [1.0, 14.6146421])
+def test_whole_runs_equal_the_stock_noise_without_warnings_for_any_schedule_height(patched, cpu, dtype, top):
+    stock = patched.__mro__[1]
+    schedule = [top, top * 0.6, top * 0.4465445876, 0.0]
+    sig = torch.tensor(schedule, dtype=dtype)
+    pos = sig[sig > 0]
+
+    def run(cls):
+        s = cls(torch.zeros(1, 2, 3, dtype=dtype), pos.min(), sig.max(), seed=7, cpu=cpu)
+        return [s(a, b) for a, b in zip(pos[:-1], pos[1:])]
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+        want = run(stock)
+    with warnings.catch_warnings():
+        warnings.simplefilter("error")
+        got = run(patched)
+    assert all(torch.equal(a, b) and a.dtype == b.dtype for a, b in zip(want, got))
 
 
 def test_random_schedules_on_the_cpu_tree_never_warn(patched):

modules/loaders/lora/keys.py

@@ -27,6 +27,7 @@ REKEY = (
 
 
 _DIFFUSERS = re.compile(r"(^|\.)(transformer_blocks|refiner_blocks)\.\d+\.")
+_WRAPPER = re.compile(r"^.*?(?=(?:token_refiner\.)?(?:transformer_blocks|refiner_blocks)\.\d+\.)")
 _PART = re.compile(r"^(?P<mod>.+?)(?P<tail>\.(?:lora_[AB]\.weight|lora_[AB]|alpha))$")
 _SUFFIX = (("attn.to_out.0", "attn.out_proj"), ("ff.net.0.proj", "mlp.fc1"), ("ff.net.2", "mlp.fc2"))
 _QKV = ("to_q", "to_k", "to_v")
@@ -45,7 +46,7 @@ def from_diffusers(lora):
         return lora, {}
     out, extra, sizes = {}, {}, {}
     for k, v in lora.items():
-        m = _PART.match(k[len("transformer."):] if k.startswith("transformer.") else k)
+        m = _PART.match(_WRAPPER.sub("", k))                      # whatever training wrapper sits in front of the block names
         if not m:
             out[k] = v
             continue
@@ -72,8 +73,7 @@ def from_diffusers(lora):
         if None in heights:
             continue
         start = sum(heights[: _QKV.index(part)])
-        plain = re.sub(r"^(?:base_model\.model\.|model\.diffusion_model\.|diffusion_model\.)+", "", base)       # the file's own wrapper is not part of the model's name
-        extra[mod] = (f"diffusion_model.{plain}.weight", (0, start, sizes[mod]))
+        extra[mod] = (f"diffusion_model.{base}.weight", (0, start, sizes[mod]))
     return out, extra
 
 

modules/loaders/tests/test_diffusers_h3_lora.py

@@ -110,16 +110,21 @@ def test_an_ltx_lora_with_transformer_blocks_names_is_left_as_it_is(tiny_ltx, mo
     assert len(patches) == 1 and dropped == 0
 
 
-@pytest.mark.parametrize("wrapper", ["diffusion_model.", "base_model.model.", "model.diffusion_model."])
-def test_a_wrapped_diffusers_file_still_reaches_the_fused_qkv_rows(patcher, wrapper):
+@pytest.mark.parametrize("wrapper", ["", "transformer.", "diffusion_model.", "base_model.model.", "model.diffusion_model.",
+                                     "base_model.model.diffusion_model.", "base_model.model.dit.", "dit.", "model.dit.", "lora_model.", "net."])
+def test_a_wrapped_diffusers_file_gets_every_change_on_a_weight_the_model_has(patcher, wrapper):
     from modules.loaders.lora.keys import match
     lora, want = _file("transformer_blocks.1", torch.Generator().manual_seed(0))
     lora = {wrapper + k[len("transformer."):]: v for k, v in lora.items()}
     patches, how, dropped = match(patcher, None, lora)
+    assert dropped == 0 and len(patches) == 6, (wrapper, how, len(patches))
+    state = patcher.model.state_dict()
+    assert all((k[0] if isinstance(k, tuple) else k) in state for k in patches)       # every patch aims at a weight that exists
     clone = patcher.clone()
     clone.add_patches(patches, 1.0)
     got = _delta(clone, "diffusion_model.blocks.1.attn.qkv_proj.weight", (192, 64))
-    assert torch.allclose(got, torch.cat([want["attn.to_q"], want["attn.to_k"], want["attn.to_v"]], 0), atol=1e-5), (how, len(patches))
+    assert torch.allclose(got, torch.cat([want["attn.to_q"], want["attn.to_k"], want["attn.to_v"]], 0), atol=1e-5)
+    assert torch.allclose(_delta(clone, "diffusion_model.blocks.1.attn.out_proj.weight", (64, 64)), want["attn.to_out.0"], atol=1e-5)
 
 
 def test_a_qkv_pair_of_the_wrong_width_is_dropped_not_merged_scrambled(patcher):