nodes/olivv-cs-ComfyUI-FunPack
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: step wrappers find the picture/sound split on a ModelPatcher

Browse Source 
_get_latent_shapes only looked under a CFGGuider's conds; wrappers hold a
ModelPatcher, so on packed AV latents every video-only wrapper passed through
unchanged. Also: late-branch and STAS say when they applied nothing, and a
learned value is only rated when the feature actually applied.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
digital-garbage
2026-09-28 19:45:08 +03:00
co-authored by Claude Opus 5.5
parent ce2de22b2e
commit aaa1c75240
4 changed files with 83 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
samplers.py
+46 -21
View File
@@ -763,9 +763,7 @@ def _get_latent_shapes(model):
try:
guider = getattr(model, "inner_model", None) # CFGGuider
conds = getattr(guider, "conds", None)
if not isinstance(conds, dict):
return None
for key in ("positive", "negative"):
for key in (("positive", "negative") if isinstance(conds, dict) else ()):
lst = conds.get(key)
if not lst:
continue
@@ -775,9 +773,15 @@ def _get_latent_shapes(model):
val = getattr(ls, "cond", None)
if val and len(val) > 1:
return val
except Exception:
pass
# A ModelPatcher (what every model_function_wrapper builder holds) has no inner_model.
# CFGGuider.inner_sample publishes the shapes on its BaseModel, which clones share.
try:
val = getattr(getattr(model, "model", None), "latent_shapes", None)
return list(val) if val and len(val) > 1 else None
except Exception:
return None
return None
def _packed_video_mask(model, x):
@@ -4167,10 +4171,10 @@ class FunPackLTXAVSceneChainSampler:
# Outermost of the block hooks: its weak pass must be able to skip every hook on the
# shared blocks, and run the tail exactly as the normal pass does.
_late_installed = False
_late_installed, _late_stats = False, None
if self._is_h3 and getattr(self, "_late_w", 0.0) > 0.0:
_model_before_late = model
model = self._install_late_branch(model, self._late_block, self._late_w)
model, _late_stats = self._install_late_branch(model, self._late_block, self._late_w)
_late_installed = model is not _model_before_late
_phrase_probe = self._install_phrase_probe(model, positive, latent)
@@ -4194,7 +4198,14 @@ class FunPackLTXAVSceneChainSampler:
callback=_progress_cb if pbar is not None else None,
)
if _stas_report is not None:
print(f"[FunPackSceneChain] STAS: {_stas_report()}")
_line, _n = _stas_report()
self._stas_ran = getattr(self, "_stas_ran", 0) + _n
print(f"[FunPackSceneChain] STAS: {_line}")
if _late_stats is not None:
self._late_ran = getattr(self, "_late_ran", 0) + _late_stats["guided"]
print("[FunPackSceneChain] late-branch guidance: " + (
f"guided {_late_stats['guided']} step call(s)" if _late_stats["guided"]
else (_late_stats["why"] or "no step call reached it, nothing guided")))
if _phrase_probe is not None:
self._report_phrase_probe(_phrase_probe[1])
# The masked probe passes ran through every capture hook too, so this
@@ -8546,14 +8557,15 @@ class FunPackLTXAVSceneChainSampler:
patched.model_options["transformer_options"] = to
def _report():
"""-> (log line, steered call count)."""
if stats["steered"]:
line = (f"alpha {_alpha:.2f} at block {block}: steered {stats['steered']} "
f"call(s), channels {stats['dims']}")
if stats["empty"]:
line += f"; {stats['empty']} call(s) had no massive channel"
return line
return line, stats["steered"]
return (stats["why"] or f"no channel at block {block} is over "
f"{_stas.MA_RATIO:.0f}x the mean, nothing steered -- try another block")
f"{_stas.MA_RATIO:.0f}x the mean, nothing steered -- try another block"), 0
print(f"[FunPackSceneChain] STAS: alpha {_alpha:.2f} at block {block}, first "
f"{_stas.EARLY_FRACTION:.0%} of steps, frame 0 + frame edges, picture only.")
@@ -8581,7 +8593,7 @@ class FunPackLTXAVSceneChainSampler:
_log.failed("FunPackSceneChain", "late-branch guidance",
f"block {branch} is not a block of this model",
"the run goes on without it")
return model
return model, None
patched = model.clone()
to = patched.model_options.get("transformer_options", {}).copy()
patches_replace = dict(to.get("patches_replace", {}))
@@ -8618,6 +8630,7 @@ class FunPackLTXAVSceneChainSampler:
old_wrapper = patched.model_options.get("model_function_wrapper")
_w = float(w)
stats = {"guided": 0, "why": None}
def _guided(apply_fn, cond_only):
def _apply(x, t, **c):
@@ -8639,7 +8652,13 @@ class FunPackLTXAVSceneChainSampler:
if span is None:
# A plain 5-D latent is picture only; an unreadable packed one is
# left alone rather than guessed at.
return _late.mix(normal, weak, _w) if normal.dim() == 5 else normal
if normal.dim() != 5:
stats["why"] = ("picture and sound could not be told apart in "
"the latent, nothing guided")
return normal
stats["guided"] += 1
return _late.mix(normal, weak, _w)
stats["guided"] += 1
off, sz, _shape = span
video = _late.mix(normal[..., off:off + sz], weak[..., off:off + sz], _w)
return torch.cat([normal[..., :off], video.to(normal.dtype),
@@ -8659,11 +8678,11 @@ class FunPackLTXAVSceneChainSampler:
print(f"[FunPackSceneChain] late-branch guidance: strength {_w:.2f}, weak copy "
f"skips block {branch} and shares 0-{branch - 1} "
f"(~{100.0 * (len(dit) - branch) / len(dit):.0f}% extra per step), picture only.")
return patched
return patched, stats
except Exception as _e: # noqa: BLE001
_log.failed("FunPackSceneChain", "late-branch guidance", _e,
"the run goes on without it")
return model
return model, None
def _install_block_influence(self, model, capture_holder, n_blocks=50, max_rows=512):
"""MEASUREMENT ONLY (H3). See block_influence.py.
@@ -9582,6 +9601,7 @@ class FunPackLTXAVSceneChainSampler:
self._tsr_k = 1.0
self._late_w, self._late_block = 0.0, int(late_guidance_block)
self._stas_alpha, self._stas_block = None, int(stas_block)
self._late_ran, self._stas_ran = 0, 0 # calls the features actually changed
if (shot_memory != "off" or tsr == "learned" or late_guidance == "learned"
or stas == "learned"):
if not refinement_key_input:
@@ -10866,14 +10886,19 @@ class FunPackLTXAVSceneChainSampler:
_tsr.save_pending(refinement_key_input, self._tsr_k)
else:
_tsr.clear_pending(refinement_key_input)
if late_guidance == "learned" and self._is_h3:
_late.DIAL.save_pending(refinement_key_input, self._late_w)
else:
_late.DIAL.clear_pending(refinement_key_input)
if stas == "learned" and self._stas_alpha is not None:
_stas.DIAL.save_pending(refinement_key_input, self._stas_alpha)
else:
_stas.DIAL.clear_pending(refinement_key_input)
# A learned value is only rated when it changed something: a run where the
# feature never applied must not teach the dial about the value it drew.
for _on, _ran, _dial, _v, _name in (
(late_guidance == "learned", self._late_ran, _late.DIAL, self._late_w,
"late-branch guidance"),
(stas == "learned", self._stas_ran, _stas.DIAL, self._stas_alpha, "STAS")):
if _on and _ran and _v is not None:
_dial.save_pending(refinement_key_input, _v)
continue
_dial.clear_pending(refinement_key_input)
if _on and self._is_h3:
print(f"[FunPackSceneChain] {_name}: did not apply this run, so this "
f"run's rating will not teach its learned strength")
# Trajectory probe: same run/rating pairing as the two above — the rating that scores
# THIS run appends these per-bucket descriptors to the measurement log.
tests/test_late_guidance.py
+7 -6
View File
@@ -71,7 +71,7 @@ def test_weak_pass_shares_the_head_skips_the_branch_and_pushes_away():
return extra["original_block"](args)
base = _Model(net)
base.model_options["transformer_options"] = {"patches_replace": {"dit": {("double_block", 5): spy}}}
patched = samplers.FunPackLTXAVSceneChainSampler()._install_late_branch(base, 3, 0.5)
patched = samplers.FunPackLTXAVSceneChainSampler()._install_late_branch(base, 3, 0.5)[0]
assert patched is not base
x = torch.zeros(1, 2, 1, 2, 2)
@@ -86,24 +86,25 @@ def test_refuses_a_block_the_model_does_not_have():
import samplers
base = _Model(_Blocks())
s = samplers.FunPackLTXAVSceneChainSampler()
assert s._install_late_branch(base, 0, 0.5) is base
assert s._install_late_branch(base, 6, 0.5) is base
assert s._install_late_branch(base, 0, 0.5)[0] is base
assert s._install_late_branch(base, 6, 0.5)[0] is base
def test_unreadable_packed_latent_is_left_alone(monkeypatch):
import samplers
monkeypatch.setattr(samplers, "_video_span", lambda model, x: None)
net = _Blocks()
patched = samplers.FunPackLTXAVSceneChainSampler()._install_late_branch(_Model(net), 3, 0.5)
patched, stats = samplers.FunPackLTXAVSceneChainSampler()._install_late_branch(_Model(net), 3, 0.5)
x = torch.zeros(1, 1, 8)
assert torch.allclose(_run(patched, net, x), torch.full_like(x, 21.0))
assert stats["guided"] == 0 and "nothing guided" in stats["why"] # said, not silent
def test_video_span_only_the_picture_is_guided(monkeypatch):
import samplers
monkeypatch.setattr(samplers, "_video_span", lambda model, x: (0, 5, (1, 1, 1, 1, 5)))
net = _Blocks()
patched = samplers.FunPackLTXAVSceneChainSampler()._install_late_branch(_Model(net), 3, 1.0)
patched = samplers.FunPackLTXAVSceneChainSampler()._install_late_branch(_Model(net), 3, 1.0)[0]
out = _run(patched, net, torch.zeros(1, 1, 8))
assert torch.allclose(out[..., :5], torch.full((1, 1, 5), 25.0)) # 21 + 1*(21-17)
assert torch.allclose(out[..., 5:], torch.full((1, 1, 3), 21.0)) # sound: normal pass
@@ -113,7 +114,7 @@ def test_negative_prompt_calls_are_not_guided():
"""CFG above 1: the negative call runs once, plain; the prompt call is still guided."""
import samplers
net = _Blocks()
patched = samplers.FunPackLTXAVSceneChainSampler()._install_late_branch(_Model(net), 3, 0.5)
patched = samplers.FunPackLTXAVSceneChainSampler()._install_late_branch(_Model(net), 3, 0.5)[0]
wrapper = patched.model_options["model_function_wrapper"]
x = torch.zeros(1, 2, 1, 2, 2)
c = {"transformer_options": patched.model_options["transformer_options"]}
tests/test_latent_shapes_on_patcher.py
+28
View File
@@ -0,0 +1,28 @@
"""The step wrappers get a ModelPatcher. ComfyUI publishes the packed picture+sound
shapes on its BaseModel (CFGGuider.inner_sample: inner_model.latent_shapes), not under
an inner_model, so the video span must be found from there too."""
import sys
from pathlib import Path
import torch
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
def test_video_span_found_from_a_model_patcher():
import samplers
class Base:
latent_shapes = [torch.Size([1, 24, 3, 4, 6]), torch.Size([1, 32, 5, 2])]
class Patcher: # what model_function_wrapper builders receive
model = Base()
model_options = {}
n = 24 * 3 * 4 * 6 + 32 * 5 * 2
assert samplers._video_span(Patcher(), torch.zeros(1, 1, n)) == (0, 24 * 3 * 4 * 6, (1, 24, 3, 4, 6))
# a stale shape list from another run does not match this tensor: no guess
assert samplers._video_span(Patcher(), torch.zeros(1, 1, n + 1)) is None
# single stream: nothing to split
Base.latent_shapes = [torch.Size([1, 24, 3, 4, 6])]
assert samplers._video_span(Patcher(), torch.zeros(1, 1, 24 * 3 * 4 * 6)) is None
tests/test_stas.py
+2 -2
View File
@@ -83,7 +83,7 @@ def test_hook_steers_picture_rows_on_early_steps_only(monkeypatch):
assert torch.equal(out[17:21], h[17:21]) # interior
late, _ = _call(patched, h.clone(), step=5, total=7, video_rows=None)
assert torch.equal(late, h) # past 40%
assert "channels 3" in report()
assert "channels 3" in report()[0] and report()[1] == 1
def test_unreadable_layout_is_declared(monkeypatch):
@@ -94,4 +94,4 @@ def test_unreadable_layout_is_declared(monkeypatch):
_Model(), 2, 2.0, torch.zeros(1, 24, 2, 4, 6))
h = torch.randn(30, 8)
out, _ = _call(patched, h.clone(), step=0, total=7, video_rows=None)
assert torch.equal(out, h) and "nothing steered" in report()
assert torch.equal(out, h) and "nothing steered" in report()[0] and report()[1] == 0