diff --git a/h3_shadow_negative.py b/h3_shadow_negative.py index 3e0eef22..9ec04423 100644 --- a/h3_shadow_negative.py +++ b/h3_shadow_negative.py @@ -83,7 +83,7 @@ class ShadowState: same fact), rebuilt fresh from `negative_context` on first use afterwards. """ __slots__ = ("dm", "negative_context", "video_scale", "audio_scale", "tau", "alpha", - "start_percent", "end_percent", "neg_h", "pos_text_len", + "start_percent", "end_percent", "neg_h", "pos_text_len", "text_len", "warned_layout", "warned_fallback", "hard_disabled") def __init__(self, dm, negative_context, video_scale, audio_scale, tau, alpha, @@ -98,6 +98,12 @@ class ShadowState: self.end_percent = float(end_percent) self.neg_h = None self.pos_text_len = None + # The POSITIVE prompt's token count for the model call in progress. Set by the + # model_function_wrapper _install_h3_shadow_negative puts on the model, which is the + # one place the conditioning tensor is visible -- ComfyUI never copies it into + # transformer_options, so reading it from there (as this first shipped) found None on + # every block and the whole feature silently ran the stock forward. + self.text_len = None self.warned_layout = False self.warned_fallback = False self.hard_disabled = False @@ -118,6 +124,25 @@ def _prepare_negative(state, x, transformer_options): return neg +def _tag_is(row, tag): + """A mod row is an int, or -- once a pinned frame or mask gives rows their own strength + (any i2v scene) -- a per-token LongTensor. int() on that tensor raised, which disabled + shadow negative for the rest of the run on every i2v scene.""" + if isinstance(row, torch.Tensor): + return bool(((row % 3) == tag).all()) + return int(row) % 3 == tag + + +def step_progress(transformer_options): + """0 at the first denoise step, 1 at the last, from the schedule core puts in + transformer_options (sample_sigmas + this call's sigmas). None if unresolvable.""" + to = transformer_options or {} + sched, cur = to.get("sample_sigmas"), to.get("sigmas") + if not isinstance(sched, torch.Tensor) or cur is None: + return None + return sigma_progress(cur, sched) + + def _mod_one(h, shift, scale, row): return h.mul_(1.0 + scale[row].to(h.dtype)).add_(shift[row].to(h.dtype)) @@ -142,7 +167,7 @@ def _presentation_plan(mod_segments, pos_text_len): if b > pos_text_len: return None, None runs.append(seg) - text_runs = [r for r in runs if (r[2] % 3) == 1] + text_runs = [r for r in runs if _tag_is(r[2], 1)] if not runs or not text_runs: return None, None return runs, text_runs[-1] @@ -183,7 +208,7 @@ def _build_shadow_attention_input(state, block, neg_h, h_pos, rope, presentation h_parts, r_parts, cursor = [], [], 0 neg_a = neg_b = None for a, b, row in presentation_runs: - if (a, b, row) == user_run: + if (a, b) == tuple(user_run[:2]): neg_a, neg_b = cursor, cursor + neg_norm.shape[0] h_parts.append(neg_norm) r_parts.append(neg_rope) @@ -233,26 +258,26 @@ def _shadow_forward(state, block, args, extra_options): rope = args["rope_freqs"] to = args["transformer_options"] - sigmas = to.get("sample_sigmas", to.get("sigmas")) - progress = sigma_progress(args.get("timestep", to.get("timestep")), sigmas) + progress = step_progress(to) active = state.alpha > 0 and (state.video_scale != 1.0 or state.audio_scale != 1.0) if progress is not None: active = active and state.start_percent <= progress <= state.end_percent if not active or len(mod_segments) < 3: return extra_options["original_block"](args) - state.pos_text_len = int(x.shape[0]) # placeholder, corrected below once c_crossattn known - cross_attn = to.get("c_crossattn") - pos_text_len = (int(cross_attn.shape[1]) - if isinstance(cross_attn, torch.Tensor) and cross_attn.ndim >= 3 - else None) + pos_text_len = state.text_len if pos_text_len is None: + if not state.warned_layout: + state.warned_layout = True + print("[FunPackSceneChain] H3 shadow negative: the prompt length was not seen for " + "this model call (the wrapper that records it is missing) -- bypassing " + "shadow guidance.") return extra_options["original_block"](args) state.pos_text_len = pos_text_len audio_a, audio_b, audio_row = mod_segments[-2] video_a, video_b, video_row = mod_segments[-1] - if (int(audio_row) % 3) != 2 or (int(video_row) % 3) != 0 or audio_b != video_a \ + if not _tag_is(audio_row, 2) or not _tag_is(video_row, 0) or audio_b != video_a \ or video_b != x.shape[0]: if not state.warned_layout: state.warned_layout = True diff --git a/samplers.py b/samplers.py index 8b226727..1e8a0109 100644 --- a/samplers.py +++ b/samplers.py @@ -7991,6 +7991,24 @@ class FunPackLTXAVSceneChainSampler: dm=dm, negative_context=negative[0][0].detach(), video_scale=video_scale, audio_scale=audio_scale, tau=tau, alpha=alpha, start_percent=_lo, end_percent=_hi) + + # The positive prompt's length is only visible here, on the model call's own + # conditioning -- core never copies it into transformer_options. Recorded per + # call, and the shadow stream reset per call, so each denoise evaluation starts + # the negative fresh at block 0 whatever blocks are hooked. + _old_wrapper = patched.model_options.get("model_function_wrapper") + + def _shadow_wrapper(apply_fn, args): + ctx = (args.get("c") or {}).get("c_crossattn") + state.text_len = (int(ctx.shape[1]) if isinstance(ctx, torch.Tensor) + and ctx.ndim >= 3 else None) + state.neg_h = None + if _old_wrapper is not None: + return _old_wrapper(apply_fn, args) + return apply_fn(args["input"], args["timestep"], **args.get("c", {})) + + patched.model_options["model_function_wrapper"] = _tag_scene_wrapper( + _shadow_wrapper, _old_wrapper) _skipped = [] for i, block in enumerate(dm.blocks): if compose and i in _claimed: diff --git a/tests/test_h3_shadow_negative.py b/tests/test_h3_shadow_negative.py index 018552a3..8498a684 100644 --- a/tests/test_h3_shadow_negative.py +++ b/tests/test_h3_shadow_negative.py @@ -141,3 +141,42 @@ def test_empty_negative_context_yields_nothing(): audio_scale=1.0, tau=2.5, alpha=0.35, start_percent=0.0, end_percent=0.6) assert sn._prepare_negative(state, torch.zeros(1, 4), {}) is None + + +# --- 2026-09-27: the fix for "never ran" -------------------------------------------- +# The prompt length used to be read from transformer_options["c_crossattn"], which ComfyUI +# never sets, so every block ran the stock forward. Proven against a real tiny H3 forward +# outside this (stubbed) suite; these pin the pieces the fix rests on. + +def test_a_per_token_mod_row_is_accepted_not_int_cast(): + rows = torch.tensor([3, 6, 9]) # all tag 0 (video), per-token strengths + assert sn._tag_is(rows, 0) and not sn._tag_is(rows, 2) + assert sn._tag_is(5, 2) and not sn._tag_is(5, 0) + + +def test_the_user_prompt_run_is_found_when_rows_are_tensors(): + segs = [(0, 3, 1), (3, 5, torch.tensor([0, 0])), (5, 9, 1), (9, 12, 2)] + runs, user = sn._presentation_plan(segs, 9) + assert user == (5, 9, 1) and len(runs) == 3 + + +def test_progress_comes_from_the_schedule_core_publishes(): + to = {"sample_sigmas": torch.tensor([1.0, 0.7, 0.3, 0.0]), "sigmas": torch.tensor([0.3])} + assert sn.step_progress(to) == 1.0 + to["sigmas"] = torch.tensor([1.0]) + assert sn.step_progress(to) == 0.0 + assert sn.step_progress({}) is None + + +def test_without_a_recorded_prompt_length_it_bypasses_and_says_so(capsys): + state = sn.ShadowState(dm=None, negative_context=torch.randn(1, 3, 16), video_scale=3.0, + audio_scale=1.0, tau=2.5, alpha=1.0, start_percent=0.0, + end_percent=1.0) + ran = [] + hook = sn.make_block_hook(object(), state, 1) + args = {"img": torch.zeros(4, 16), "t_emb": None, + "mod_segments": [(0, 1, 1), (1, 2, 2), (2, 4, 0)], "rope_freqs": None, + "transformer_options": {}} + hook(args, {"original_block": lambda a: ran.append(1) or {"img": a["img"]}}) + assert ran == [1] + assert "prompt length was not seen" in capsys.readouterr().out diff --git a/tests/test_h3_shadow_negative_install.py b/tests/test_h3_shadow_negative_install.py index 826585c3..4464b3c6 100644 --- a/tests/test_h3_shadow_negative_install.py +++ b/tests/test_h3_shadow_negative_install.py @@ -278,3 +278,31 @@ if __name__ == "__main__": test_both_scales_at_one_is_a_true_noop() test_enabled_clones_and_hooks_every_block() print("ok (run via pytest for the rest)") + + +def test_the_model_wrapper_records_the_prompt_length_and_chains_the_old_one(): + """The prompt length is only visible on the model call's conditioning, so the installer + puts a wrapper there; it must chain whatever wrapper was already installed.""" + import h3_shadow_negative as sn + model = _FakeModel() + seen = [] + model.model_options["model_function_wrapper"] = lambda fn, a: seen.append("old") or "out" + patched = S()._install_h3_shadow_negative(model, _negative(), 3.0, 1.0, 2.5, 0.35, 0.0, 0.6) + wrapper = patched.model_options["model_function_wrapper"] + captured = {} + orig_forward = sn._shadow_forward + + def spy(state, block, args, extra): + captured["len"] = state.text_len + return extra["original_block"](args) + + sn._shadow_forward = spy + try: + out = wrapper(None, {"input": None, "timestep": None, + "c": {"c_crossattn": torch.zeros(1, 7, 16)}}) + hook = patched.model_options["transformer_options"]["patches_replace"]["dit"][("double_block", 0)] + hook({"img": torch.zeros(2, 16)}, {"original_block": lambda a: {"img": a["img"]}}) + finally: + sn._shadow_forward = orig_forward + assert out == "out" and seen == ["old"] + assert captured["len"] == 7