From 88cc86aef0c8475711997e5b2612af4c427df0cf Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Wed, 26 Aug 2026 23:25:09 +0300 Subject: [PATCH] =?UTF-8?q?fix(h3):=20the=20RAVEN=20LoRA=20needs=20its=20h?= =?UTF-8?q?ead=20RENAMED,=20not=20stripped=20=E2=80=94=20and=20the=20lane?= =?UTF-8?q?=20shows=20progress?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two things, both from the rental. THE LORA. Its keys are ComfyUI's module paths exactly, under a different head: base_model.model.dit.audio_patch_proj.lora_A.weight diffusion_model.audio_patch_proj <- what comfy's map holds Every candidate the loader tried was a STRIP, which left `audio_patch_proj...` — and comfy's key map has no bare form to land on, so 532 keys matched nothing and the file did nothing at all. LORA_KEY_RENAMES adds (prefix, replacement) candidates, scored against the real key map like the strips, so a rename that lands beats a strip that does not without either being ordered first. THE PROGRESS. The causal lane replaces comfy's sampler rather than wrapping it, so nothing ticked the progress bar and nothing answered an interrupt: a long generation looked exactly like a hung one, and Cancel did nothing until the whole clip finished. The rollout now reports every model call, and the sampler maps that onto the budget this scene was allotted — the bar was sized in scenes x steps and a chunked scene runs that many steps PER CHUNK, so ticking raw would have overrun it. The phase readout says "scene 1/3 · chunk 7/22", and the interrupt is checked before every step, not just every chunk. Tests: 1829 python (+5), 68 JS. --- h3_causal.py | 17 +++++++++-- model_management.py | 56 ++++++++++++++++++++++++++++------ samplers.py | 29 ++++++++++++++++-- tests/test_causal_rollout.py | 28 +++++++++++++++-- tests/test_h3_causal_loader.py | 35 ++++++++++++++++++++- 5 files changed, 145 insertions(+), 20 deletions(-) diff --git a/h3_causal.py b/h3_causal.py index 35994744..ba80372b 100644 --- a/h3_causal.py +++ b/h3_causal.py @@ -328,7 +328,7 @@ def audio_sigma_for(sigma, shift_video, shift_audio): def causal_rollout(*, chunks, sigmas, forward, commit, draw_noise, video_noise, audio_noise, step_rule="consistency", eta=1.0, known_chunks=0, known_video=None, known_audio=None, - on_chunk=None, cancel=None): + on_chunk=None, on_step=None, cancel=None): """Generate a clip chunk by chunk, each one attending to the cache of the ones before it. Everything the model does arrives as a callable, so the loop is testable without ComfyUI, @@ -343,6 +343,11 @@ def causal_rollout(*, chunks, sigmas, forward, commit, draw_noise, is a chunk the rest of the clip cannot see. draw_noise(shape) fresh noise. Called in a fixed order so a seeded run reproduces exactly. + on_step(chunk_index, done, total) + after every model call, with the count of chunk-steps finished out of the clip's + total. This lane does not go through comfy's sampler, so nothing else ticks a + progress bar — and a long generation showing no progress is indistinguishable from + a hung one. `known_chunks` leading chunks are taken from `known_video`/`known_audio` instead of being sampled. That is how an i2v anchor survives: chunk 0 comes from FunPack's ORDINARY dense @@ -382,6 +387,8 @@ def causal_rollout(*, chunks, sigmas, forward, commit, draw_noise, video_xt = video_noise[:, :, v_start:v_stop].clone() audio_xt = audio_noise[..., a_start:a_stop].clone() for position in range(len(sigmas) - 1): + if cancel is not None: + cancel(index) sigma, sigma_next = sigmas[position], sigmas[position + 1] video_v, audio_v = forward(video_xt, audio_xt, index, sigma) # H3's head is data-ward velocity; x0 = x_t + sigma * v, no negation. @@ -394,6 +401,9 @@ def causal_rollout(*, chunks, sigmas, forward, commit, draw_noise, audio_eps = draw_noise(audio_xt.shape) video_xt = step(step_rule, video_xt, video_x0, sigma, sigma_next, video_eps, eta) audio_xt = step(step_rule, audio_xt, audio_x0, sigma, sigma_next, audio_eps, eta) + if on_step is not None: + on_step(index, index * (len(sigmas) - 1) + position + 1, + len(chunks) * (len(sigmas) - 1)) video_out[:, :, v_start:v_stop] = video_xt audio_out[..., a_start:a_stop] = audio_xt @@ -924,7 +934,7 @@ def build_session(model, positive, latent, *, sink=2, window=2, device=None, off def run_session(session, *, sigmas, step_rule="consistency", eta=1.0, seed=0, known_chunks=0, known_video=None, known_audio=None, - on_chunk=None, cancel=None, transformer_options=None): + on_chunk=None, on_step=None, cancel=None, transformer_options=None): """Drive one causal rollout over an assembled session. Returns (video, audio).""" model, cache, plan = session["model"], session["cache"], session["plan"] device, payload = session["device"], session["payload"] @@ -981,4 +991,5 @@ def run_session(session, *, sigmas, step_rule="consistency", eta=1.0, seed=0, chunks=plan.bounds, sigmas=sigmas, forward=forward, commit=commit, draw_noise=draw_noise, video_noise=video_noise, audio_noise=audio_noise, step_rule=step_rule, eta=eta, known_chunks=known_chunks, - known_video=known_video, known_audio=known_audio, on_chunk=on_chunk, cancel=cancel) + known_video=known_video, known_audio=known_audio, on_chunk=on_chunk, + on_step=on_step, cancel=cancel) diff --git a/model_management.py b/model_management.py index cd10f1e4..62b90020 100644 --- a/model_management.py +++ b/model_management.py @@ -293,6 +293,20 @@ LORA_KEY_PREFIXES = ( "net.", ) +#: (prefix, replacement) for adapters that name the same modules under a different head. +#: Stripping is not enough for these: comfy's key map has no bare form, only +#: `diffusion_model.`, so the head has to be REPLACED rather than removed. +LORA_KEY_RENAMES = ( + # RAVEN's streaming LoRA for MiniMax H3: `base_model.model.dit.blocks.0.adaln_proj.linear` + # is comfy's `diffusion_model.blocks.0.adaln_proj.linear`, module for module. Stripping + # left `blocks.0...`, which matches nothing, and the file silently did nothing at all. + ("base_model.model.dit.", "diffusion_model."), + ("base_model.model.", "diffusion_model."), + ("dit.", "diffusion_model."), + ("model.dit.", "diffusion_model."), + ("transformer.", "diffusion_model."), +) + def _sample_keys(mapping, count=3): """A few representative names, shortened. Enough to compare two namings by eye.""" @@ -309,10 +323,19 @@ def _count_lora_matches(lora, key_map): return len(comfy.lora.load_lora(lora, key_map)) -def _strip_key_prefix(lora, prefix): +def _rekey(lora, prefix, replacement=""): + """`lora` with `prefix` replaced by `replacement` on every key that carries it, or None. + + Stripping is the `replacement=""` case. RENAMING matters as much: an adapter trained + against a wrapper that calls the DiT something else has ComfyUI's exact module path + underneath a different head — `base_model.model.dit.audio_patch_proj` against comfy's + `diffusion_model.audio_patch_proj` — and stripping alone leaves it matching nothing, + because comfy's key map has no bare form to land on. + """ if not any(k.startswith(prefix) for k in lora): return None - return {(k[len(prefix):] if k.startswith(prefix) else k): v for k, v in lora.items()} + return {(replacement + k[len(prefix):] if k.startswith(prefix) else k): v + for k, v in lora.items()} def _lora_pair_dims(adapter): @@ -410,6 +433,26 @@ def _dropped_all_of(line): return any(int(m.group(1)) == int(m.group(2)) for m in _DROPPED_FRAGMENT.finditer(line)) +def _best_lora_keying(converted, key_map): + """The keying of `converted` that matches the most of `key_map`. -> (lora, note). + + Every candidate is SCORED rather than ordered, so a rename that lands beats a strip that + does not without either having to be tried first, and a format nobody anticipated still + wins if its keys are recognisable underneath. + """ + best, best_count, best_note = converted, _count_lora_matches(converted, key_map), "as-is" + candidates = [(p, "", f"stripped {p}") for p in LORA_KEY_PREFIXES] + candidates += [(p, r, f"{p} -> {r}") for p, r in LORA_KEY_RENAMES] + for prefix, replacement, label in candidates: + variant = _rekey(converted, prefix, replacement) + if variant is None: + continue + count = _count_lora_matches(variant, key_map) + if count > best_count: + best, best_count, best_note = variant, count, label + return best, best_note + + def resolve_lora_patches(model, lora, clip=None, name=None): """Match a LoRA against a model, trying known wrapper prefixes. -> (patches, note). @@ -429,14 +472,7 @@ def resolve_lora_patches(model, lora, clip=None, name=None): key_map = comfy.lora.model_lora_keys_clip(clip.cond_stage_model, key_map) converted = comfy.lora_convert.convert_lora(lora) - best, best_count, best_note = converted, _count_lora_matches(converted, key_map), "as-is" - for prefix in LORA_KEY_PREFIXES: - variant = _strip_key_prefix(converted, prefix) - if variant is None: - continue - count = _count_lora_matches(variant, key_map) - if count > best_count: - best, best_count, best_note = variant, count, f"stripped {prefix}" + best, best_note = _best_lora_keying(converted, key_map) who = f"{name}: " if name else "" patches = comfy.lora.load_lora(best, key_map) diff --git a/samplers.py b/samplers.py index 4fd74553..ea136d12 100644 --- a/samplers.py +++ b/samplers.py @@ -8,6 +8,7 @@ import types import torch import comfy.k_diffusion.sampling as k_diffusion_sampling +import comfy.model_management import comfy.model_sampling import comfy.nested_tensor import comfy.sample @@ -4023,7 +4024,8 @@ class FunPackLTXAVSceneChainSampler: try: # The causal lane replaces the sampler entirely rather than wrapping it, so it is # tried here and falls through (loudly) when it cannot run. - sampled = self._run_h3_causal(model, positive, samples, sigmas, seed) + sampled = self._run_h3_causal(model, positive, samples, sigmas, seed, + pbar=pbar, step_offset=step_offset) if sampled is None: sampled = comfy.sample.sample_custom( model, noise, float(cfg), sampler, sigmas, positive, negative, samples, @@ -6871,7 +6873,8 @@ class FunPackLTXAVSceneChainSampler: return None return (projected.float() / norm) * float(strength) - def _run_h3_causal(self, model, positive, samples, sigmas, seed): + def _run_h3_causal(self, model, positive, samples, sigmas, seed, + pbar=None, step_offset=0): """Generate the clip in remembering chunks instead of denoising it all at once. Returns the finished latent, or None to fall through to ordinary sampling — every @@ -6916,16 +6919,36 @@ class FunPackLTXAVSceneChainSampler: f"never seen a key/value cache. The wired sampler and CFG are NOT used on this " f"path: the chunk loop is not a k-diffusion sampler, and H3 generates at CFG 1 " f"anyway. There is no per-chunk redo.") + # This lane never reaches comfy's sampler, so nothing else ticks the progress bar or + # checks for a cancel. A long generation showing no progress is indistinguishable + # from a hung one, and an interrupt that does nothing is worse. + base_label = self._current_phase_label() + scene_steps = max(1, len(schedule) - 1) + + def _on_step(index, done, total): + self._set_phase(f"{base_label}{' · ' if base_label else ''}" + f"chunk {int(index) + 1}/{session['plan'].n_chunks}") + if pbar is not None and total: + # Mapped onto the budget this scene was allotted, not counted raw: the bar was + # sized in scenes x steps, and a chunked scene runs that many steps PER CHUNK. + pbar.update_absolute(int(step_offset) + + int(round(scene_steps * float(done) / float(total)))) + + def _cancel(_index): + comfy.model_management.throw_exception_if_processing_interrupted() + try: video, audio = causal.run_session( session, sigmas=schedule, step_rule=str(getattr(self, "_h3_causal_step_rule", "consistency")), - seed=int(seed), + seed=int(seed), on_step=_on_step, cancel=_cancel, ) except Exception as error: # noqa: BLE001 _log.failed("FunPackSceneChain", "H3 causal rollout", error, "this scene is NOT generated — nothing is silently substituted") raise + finally: + self._set_phase(base_label) return self._pack_h3_av(samples, video, audio) @staticmethod diff --git a/tests/test_causal_rollout.py b/tests/test_causal_rollout.py index 2b9740d3..6af64d30 100644 --- a/tests/test_causal_rollout.py +++ b/tests/test_causal_rollout.py @@ -348,17 +348,39 @@ def test_a_schedule_with_no_steps_is_refused(): r.run(sigmas=(1.0,)) -def test_cancellation_is_checked_before_each_chunk(): +def test_cancellation_is_checked_before_each_chunk_and_each_step(): + """Per chunk is not often enough: a chunk is a full denoise, so an interrupt could sit + unanswered for every step of it while the user watches nothing happen.""" r = Recorder() seen = [] def cancel(index): seen.append(index) - if index == 1: + if len(seen) == 3: raise KeyboardInterrupt with pytest.raises(KeyboardInterrupt): r.run(cancel=cancel) - assert seen == [0, 1] + assert seen == [0, 0, 0] # chunk 0, then once per step inside it + + +def test_every_model_call_reports_progress(): + """This lane never reaches comfy's sampler, so nothing else ticks a progress bar — and a + long generation showing none is indistinguishable from a hung one.""" + r = Recorder() + ticks = [] + r.run(sigmas=(1.0, 0.5, 0.0), on_step=lambda i, done, total: ticks.append((i, done, total))) + steps = 2 + assert len(ticks) == len(r.chunks) * steps + assert [t[1] for t in ticks] == list(range(1, len(ticks) + 1)) + assert {t[2] for t in ticks} == {len(r.chunks) * steps} # a stable denominator + + +def test_progress_counts_the_whole_clip_not_the_chunk(): + """A bar that restarts at every chunk says nothing about how far along the clip is.""" + r = Recorder() + ticks = [] + r.run(sigmas=(1.0, 0.0), on_step=lambda i, done, total: ticks.append(done)) + assert ticks == list(range(1, len(r.chunks) + 1)) def test_the_chunk_callback_reports_whether_it_was_generated(): diff --git a/tests/test_h3_causal_loader.py b/tests/test_h3_causal_loader.py index 63b31e00..169d9dca 100644 --- a/tests/test_h3_causal_loader.py +++ b/tests/test_h3_causal_loader.py @@ -62,7 +62,7 @@ def test_the_longer_prefix_is_tried_before_the_shorter_one_can_win(): def test_stripping_that_prefix_yields_plain_module_paths(): - stripped = mm._strip_key_prefix( + stripped = mm._rekey( {"base_model.model.dit.blocks.0.attn.qkv_proj.lora_A.weight": 1, "base_model.model.dit.blocks.0.attn.qkv_proj.lora_B.weight": 2}, "base_model.model.dit.") @@ -70,6 +70,39 @@ def test_stripping_that_prefix_yields_plain_module_paths(): "blocks.0.attn.qkv_proj.lora_B.weight"} +# ── and why stripping was not enough ──────────────────────────────────────── +# +# Measured on the rental: the RAVEN LoRA matched NOTHING through every strip. Its keys are +# comfy's module paths exactly, under a different head — `base_model.model.dit.` for +# comfy's `diffusion_model.` — and comfy's key map has no bare form to land on, so +# stripping produced `blocks.0...` and matched none of it. The file did nothing at all. + +def test_the_raven_head_is_renamed_to_comfys_not_removed(): + assert ("base_model.model.dit.", "diffusion_model.") in mm.LORA_KEY_RENAMES + + +def test_renaming_that_head_yields_the_keys_comfys_map_holds(): + renamed = mm._rekey( + {"base_model.model.dit.audio_patch_proj.lora_A.weight": 1, + "base_model.model.dit.blocks.0.adaln_proj.linear.lora_B.weight": 2}, + "base_model.model.dit.", "diffusion_model.") + assert set(renamed) == {"diffusion_model.audio_patch_proj.lora_A.weight", + "diffusion_model.blocks.0.adaln_proj.linear.lora_B.weight"} + + +def test_a_rename_wins_over_a_strip_when_it_matches_more(monkeypatch): + """Both are scored against the real key map and the best is kept, so a rename that lands + beats a strip that does not without either being ordered ahead of the other.""" + key_map = {"diffusion_model.audio_patch_proj": "audio_patch_proj.weight"} + monkeypatch.setattr(mm, "_count_lora_matches", + lambda lora, km: sum(1 for k in lora + if k.rsplit(".lora_", 1)[0] in km)) + lora = {"base_model.model.dit.audio_patch_proj.lora_A.weight": 1} + best, note = mm._best_lora_keying(lora, key_map) + assert note == "base_model.model.dit. -> diffusion_model." + assert "diffusion_model.audio_patch_proj.lora_A.weight" in best + + # ── re-classing an already-loaded model ───────────────────────────────────── class _FakeBlock: