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: