From 4b306ae65fd351fbb0ac7594ec0d6af88ef70544 Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Thu, 5 Nov 2026 15:44:24 +0200 Subject: [PATCH] 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 | 44 +++++++------------ core/tests/test_brownian.py | 32 ++++++++------ modules/loaders/lora/keys.py | 6 +-- .../loaders/tests/test_diffusers_h3_lora.py | 11 +++-- 4 files changed, 45 insertions(+), 48 deletions(-) diff --git a/core/brownian.py b/core/brownian.py index 76d779b2..be65036a 100644 --- a/core/brownian.py +++ b/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 diff --git a/core/tests/test_brownian.py b/core/tests/test_brownian.py index 6988a6fe..5a8b1e52 100644 --- a/core/tests/test_brownian.py +++ b/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): diff --git a/modules/loaders/lora/keys.py b/modules/loaders/lora/keys.py index 1a36d3cf..e3a80018 100644 --- a/modules/loaders/lora/keys.py +++ b/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.+?)(?P\.(?: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 diff --git a/modules/loaders/tests/test_diffusers_h3_lora.py b/modules/loaders/tests/test_diffusers_h3_lora.py index 076ab6b1..37736da6 100644 --- a/modules/loaders/tests/test_diffusers_h3_lora.py +++ b/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):