fix: H3 shadow negative never ran; now it does, i2v scenes too
It read the prompt length from transformer_options['c_crossattn'], which ComfyUI never sets, so every block ran the stock forward. A model wrapper now records it per call. Per-token mod rows (any pinned frame) no longer disable it, and the step window reads the real schedule. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 5.5
parent
94f755bb4b
commit
d4d4e29ee2
+36
-11
@@ -83,7 +83,7 @@ class ShadowState:
|
|||||||
same fact), rebuilt fresh from `negative_context` on first use afterwards.
|
same fact), rebuilt fresh from `negative_context` on first use afterwards.
|
||||||
"""
|
"""
|
||||||
__slots__ = ("dm", "negative_context", "video_scale", "audio_scale", "tau", "alpha",
|
__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")
|
"warned_layout", "warned_fallback", "hard_disabled")
|
||||||
|
|
||||||
def __init__(self, dm, negative_context, video_scale, audio_scale, tau, alpha,
|
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.end_percent = float(end_percent)
|
||||||
self.neg_h = None
|
self.neg_h = None
|
||||||
self.pos_text_len = 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_layout = False
|
||||||
self.warned_fallback = False
|
self.warned_fallback = False
|
||||||
self.hard_disabled = False
|
self.hard_disabled = False
|
||||||
@@ -118,6 +124,25 @@ def _prepare_negative(state, x, transformer_options):
|
|||||||
return neg
|
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):
|
def _mod_one(h, shift, scale, row):
|
||||||
return h.mul_(1.0 + scale[row].to(h.dtype)).add_(shift[row].to(h.dtype))
|
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:
|
if b > pos_text_len:
|
||||||
return None, None
|
return None, None
|
||||||
runs.append(seg)
|
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:
|
if not runs or not text_runs:
|
||||||
return None, None
|
return None, None
|
||||||
return runs, text_runs[-1]
|
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
|
h_parts, r_parts, cursor = [], [], 0
|
||||||
neg_a = neg_b = None
|
neg_a = neg_b = None
|
||||||
for a, b, row in presentation_runs:
|
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]
|
neg_a, neg_b = cursor, cursor + neg_norm.shape[0]
|
||||||
h_parts.append(neg_norm)
|
h_parts.append(neg_norm)
|
||||||
r_parts.append(neg_rope)
|
r_parts.append(neg_rope)
|
||||||
@@ -233,26 +258,26 @@ def _shadow_forward(state, block, args, extra_options):
|
|||||||
rope = args["rope_freqs"]
|
rope = args["rope_freqs"]
|
||||||
to = args["transformer_options"]
|
to = args["transformer_options"]
|
||||||
|
|
||||||
sigmas = to.get("sample_sigmas", to.get("sigmas"))
|
progress = step_progress(to)
|
||||||
progress = sigma_progress(args.get("timestep", to.get("timestep")), sigmas)
|
|
||||||
active = state.alpha > 0 and (state.video_scale != 1.0 or state.audio_scale != 1.0)
|
active = state.alpha > 0 and (state.video_scale != 1.0 or state.audio_scale != 1.0)
|
||||||
if progress is not None:
|
if progress is not None:
|
||||||
active = active and state.start_percent <= progress <= state.end_percent
|
active = active and state.start_percent <= progress <= state.end_percent
|
||||||
if not active or len(mod_segments) < 3:
|
if not active or len(mod_segments) < 3:
|
||||||
return extra_options["original_block"](args)
|
return extra_options["original_block"](args)
|
||||||
|
|
||||||
state.pos_text_len = int(x.shape[0]) # placeholder, corrected below once c_crossattn known
|
pos_text_len = state.text_len
|
||||||
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)
|
|
||||||
if pos_text_len is None:
|
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)
|
return extra_options["original_block"](args)
|
||||||
state.pos_text_len = pos_text_len
|
state.pos_text_len = pos_text_len
|
||||||
|
|
||||||
audio_a, audio_b, audio_row = mod_segments[-2]
|
audio_a, audio_b, audio_row = mod_segments[-2]
|
||||||
video_a, video_b, video_row = mod_segments[-1]
|
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]:
|
or video_b != x.shape[0]:
|
||||||
if not state.warned_layout:
|
if not state.warned_layout:
|
||||||
state.warned_layout = True
|
state.warned_layout = True
|
||||||
|
|||||||
+18
@@ -7991,6 +7991,24 @@ class FunPackLTXAVSceneChainSampler:
|
|||||||
dm=dm, negative_context=negative[0][0].detach(),
|
dm=dm, negative_context=negative[0][0].detach(),
|
||||||
video_scale=video_scale, audio_scale=audio_scale,
|
video_scale=video_scale, audio_scale=audio_scale,
|
||||||
tau=tau, alpha=alpha, start_percent=_lo, end_percent=_hi)
|
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 = []
|
_skipped = []
|
||||||
for i, block in enumerate(dm.blocks):
|
for i, block in enumerate(dm.blocks):
|
||||||
if compose and i in _claimed:
|
if compose and i in _claimed:
|
||||||
|
|||||||
@@ -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,
|
audio_scale=1.0, tau=2.5, alpha=0.35, start_percent=0.0,
|
||||||
end_percent=0.6)
|
end_percent=0.6)
|
||||||
assert sn._prepare_negative(state, torch.zeros(1, 4), {}) is None
|
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
|
||||||
|
|||||||
@@ -278,3 +278,31 @@ if __name__ == "__main__":
|
|||||||
test_both_scales_at_one_is_a_true_noop()
|
test_both_scales_at_one_is_a_true_noop()
|
||||||
test_enabled_clones_and_hooks_every_block()
|
test_enabled_clones_and_hooks_every_block()
|
||||||
print("ok (run via pytest for the rest)")
|
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
|
||||||
|
|||||||
Reference in New Issue
Block a user