Remove FunPackContextTransitionWindows; keep Scene Cut Windows + Scene Noise

ComfyUI already has a built-in manual context windows node for long-video
windowed denoising. FunPackContextTransitionWindows duplicated that
functionality without adding anything unique, and never worked correctly
for the scene-cut use case (LTXAV's packed latent format was incompatible).

Remaining nodes:
- FunPackSceneCutWindows: attention-level scene masking (MODEL in/out)
- FunPackSceneNoise: per-scene seeded noise (LATENT out, pairs with KSamplerAdvanced)

context_transition.py is now ~180 lines vs ~320 before.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
digital-garbage
2026-05-20 18:39:39 +03:00
co-authored by Claude Sonnet 4.6
parent d255b13cc8
commit 9b9271be21
2 changed files with 14 additions and 189 deletions
+2 -5
View File
@@ -21,7 +21,7 @@ if __package__:
)
from .model_management import FunPackApplyLoraWeights, FunPackLoraLoader
from .samplers import FunPackHybridEuler2SSampler, FunPackDistilledFlowSampler
from .context_transition import FunPackContextTransitionWindows, FunPackSceneCutWindows, FunPackSceneNoise
from .context_transition import FunPackSceneCutWindows, FunPackSceneNoise
from .templates import FunPackRefinementKeyLoader, FunPackSceneBuilder
else:
# Standalone tests may not have the full ComfyUI/CUDA runtime loaded.
@@ -61,9 +61,8 @@ else:
FunPackHybridEuler2SSampler = None
FunPackDistilledFlowSampler = None
try:
from context_transition import FunPackContextTransitionWindows, FunPackSceneCutWindows, FunPackSceneNoise
from context_transition import FunPackSceneCutWindows, FunPackSceneNoise
except Exception:
FunPackContextTransitionWindows = None
FunPackSceneCutWindows = None
FunPackSceneNoise = None
try:
@@ -91,7 +90,6 @@ NODE_CLASS_MAPPINGS = {
"FunPackSaveRefinementLatent": FunPackSaveRefinementLatent,
"FunPackHybridEuler2SSampler": FunPackHybridEuler2SSampler,
"FunPackDistilledFlowSampler": FunPackDistilledFlowSampler,
"FunPackContextTransitionWindows": FunPackContextTransitionWindows,
"FunPackSceneCutWindows": FunPackSceneCutWindows,
"FunPackSceneNoise": FunPackSceneNoise,
"FunPackApplyLoraWeights": FunPackApplyLoraWeights,
@@ -118,7 +116,6 @@ NODE_DISPLAY_NAME_MAPPINGS = {
"FunPackSaveRefinementLatent": "FunPack Save Refinement Latent",
"FunPackHybridEuler2SSampler": "FunPack Hybrid Euler 2S Sampler",
"FunPackDistilledFlowSampler": "FunPack Distilled Flow Sampler",
"FunPackContextTransitionWindows": "FunPack Context Transition Windows",
"FunPackSceneCutWindows": "FunPack Scene Cut Windows",
"FunPackSceneNoise": "FunPack Scene Noise",
"FunPackApplyLoraWeights": "FunPack Apply LoRA Weights",
+12 -184
View File
@@ -1,164 +1,11 @@
import logging
import torch
import comfy.context_windows
import comfy.ldm.modules.attention as _attn_mod
# ---------------------------------------------------------------------------
# Original context-window node (kept for long-video windowed denoising)
# ---------------------------------------------------------------------------
_DIM = 2
_FALLBACK_CONTEXT_LENGTH = 16
def _offset_windows(local_windows, offset, num_frames):
shifted = []
for window in local_windows:
w = [idx + offset for idx in window if 0 <= idx + offset < num_frames]
if w:
shifted.append(w)
return shifted
def _make_auto_schedule():
base = comfy.context_windows.get_matching_context_schedule(
comfy.context_windows.ContextSchedules.STATIC_STANDARD
)
def auto_schedule(num_frames, handler, model_options):
num_windows = getattr(handler, "auto_num_windows", 1)
if num_windows <= 1 or num_frames <= 1:
return base.func(num_frames, handler, model_options)
boundaries = [
int(round(num_frames * i / num_windows))
for i in range(1, num_windows)
]
boundaries = [b for b in boundaries if 0 < b < num_frames]
if not boundaries:
return base.func(num_frames, handler, model_options)
cuts = [0] + boundaries + [num_frames]
windows = [
list(range(cuts[i], cuts[i + 1]))
for i in range(len(cuts) - 1)
if cuts[i + 1] > cuts[i]
]
logging.info(
"FunPack context windows: %s isolated scenes, frame boundaries at %s.",
num_windows, boundaries,
)
return windows or base.func(num_frames, handler, model_options)
return comfy.context_windows.ContextSchedule("funpack_isolated_scenes", auto_schedule)
_AUTO_SCHEDULE = _make_auto_schedule()
class FunPackTransitionContextHandler(comfy.context_windows.IndexListContextHandler):
def __init__(self, auto_total_frames=0):
super().__init__(
context_schedule=_AUTO_SCHEDULE,
fuse_method=comfy.context_windows.get_matching_fuse_method(
comfy.context_windows.ContextFuseMethods.PYRAMID
),
context_length=_FALLBACK_CONTEXT_LENGTH,
context_overlap=0,
context_stride=1,
closed_loop=False,
dim=_DIM,
freenoise=False,
cond_retain_index_list="",
split_conds_to_windows=True,
)
self.auto_total_frames = int(auto_total_frames or 0)
self.auto_num_windows = 1
def _detect_temporal_dim(self, x_in):
ndim = x_in.dim()
priority = [d for d in [2, 1, 0] if d < ndim]
rest = [d for d in range(ndim) if d not in priority]
for d in priority + rest:
size = int(x_in.size(d))
if 1 < size <= 4096:
return d
return None
def should_use_context(self, model, conds, x_in, timestep, model_options):
detected = self._detect_temporal_dim(x_in)
if detected is None:
logging.debug(
"FunPack context windows: x_in shape %s has no suitable temporal dim; skipping.",
list(x_in.shape),
)
return False
if detected != self.dim:
logging.info(
"FunPack context windows: auto-detected temporal dim=%s (shape %s).",
detected, list(x_in.shape),
)
self.dim = detected
try:
group_sizes = [len(c) for c in conds if c]
num_windows = max(group_sizes, default=1)
except Exception:
group_sizes = []
num_windows = 1
self.auto_num_windows = num_windows
logging.info(
"FunPack context windows: cond group sizes=%s, num_windows=%s, dim=%s, x_in=%s.",
group_sizes, num_windows, self.dim, list(x_in.shape),
)
if num_windows > 1:
total = self.auto_total_frames if self.auto_total_frames > 0 else int(x_in.size(self.dim))
self.context_length = max(1, total // num_windows)
return super().should_use_context(model, conds, x_in, timestep, model_options) or num_windows > 1
class FunPackContextTransitionWindows:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"model": ("MODEL", {"tooltip": "The model to wrap with per-scene context windows."}),
},
"optional": {
"total_frames": ("INT", {
"default": 0, "min": 0, "max": 16384, "step": 1,
"tooltip": (
"Latent frame count of your video. Used for VRAM estimation. "
"Leave 0 to auto-detect from the latent at sampling time."
),
}),
},
}
RETURN_TYPES = ("MODEL",)
RETURN_NAMES = ("model",)
FUNCTION = "apply"
CATEGORY = "FunPack/Sampling"
DESCRIPTION = (
"Wraps a model with per-scene context windows (denoising-loop level). "
"Works for standard LTX but NOT for LTXAV (packed latent format). "
"For LTXAV, use FunPack Scene Cut Windows instead."
)
def apply(self, model, total_frames=0):
model = model.clone()
handler = FunPackTransitionContextHandler(auto_total_frames=total_frames)
model.model_options["context_handler"] = handler
comfy.context_windows.create_prepare_sampling_wrapper(model)
return (model,)
# ---------------------------------------------------------------------------
# Scene-cut attention masking: correct approach for LTXAV and standard LTX
# Scene-cut attention masking
# ---------------------------------------------------------------------------
def _scene_split_self_attn(orig_fn, q, k, v, heads, num_scenes, cache, *args, **kwargs):
@@ -197,7 +44,7 @@ class FunPackSceneCutHandler:
except Exception:
self.num_scenes = 1
logging.info(
"FunPack scene cut: num_scenes=%s detected from conditioning.",
"FunPack scene cut: %s scene(s) detected from conditioning.",
self.num_scenes,
)
return False # full-pass denoising; masking applied inside each forward
@@ -209,7 +56,7 @@ class FunPackSceneCutWindows:
return {
"required": {
"model": ("MODEL", {
"tooltip": "The model to wrap with scene-boundary attention masking.",
"tooltip": "The model to apply scene-boundary attention masking to.",
}),
},
}
@@ -231,7 +78,6 @@ class FunPackSceneCutWindows:
attn_mask_cache = {}
handler = FunPackSceneCutHandler()
# Use context_handler slot so should_use_context runs each step
model.model_options["context_handler"] = handler
old_wrapper = model.model_options.get("model_function_wrapper")
@@ -248,7 +94,6 @@ class FunPackSceneCutWindows:
orig_attn_masked = _attn_mod.optimized_attention_masked
def _patched(q, k, v, heads, *a, **kw):
# Only mask self-attention (q and k same seq-len = same sequence = video tokens)
if q.shape[1] == k.shape[1]:
return _scene_split_self_attn(orig_attn, q, k, v, heads, num_scenes, attn_mask_cache, *a, **kw)
return orig_attn(q, k, v, heads, *a, **kw)
@@ -310,32 +155,24 @@ class FunPackSceneNoise:
def generate(self, latent, seed, positive=None):
samples = latent["samples"]
num_scenes = max(1, len(positive)) if positive else 1
if num_scenes <= 1:
# Single scene: standard noise, same as empty latent into KSampler
def _make_noise(shape, s):
rng = torch.Generator(device=samples.device)
rng.manual_seed(seed)
noise = torch.randn(samples.shape, generator=rng,
dtype=samples.dtype, device=samples.device)
return ({"samples": noise},)
rng.manual_seed(s)
return torch.randn(shape, generator=rng, dtype=samples.dtype, device=samples.device)
if num_scenes <= 1:
return ({"samples": _make_noise(samples.shape, seed)},)
# Find the temporal dimension: the axis with size > 1 and <= 4096
temporal_dim = None
for d in [2, 1, 0] + list(range(3, samples.dim())):
s = int(samples.size(d))
if 1 < s <= 4096:
if 1 < int(samples.size(d)) <= 4096:
temporal_dim = d
break
if temporal_dim is None:
# Packed format - fall back to single noise
rng = torch.Generator(device=samples.device)
rng.manual_seed(seed)
noise = torch.randn(samples.shape, generator=rng,
dtype=samples.dtype, device=samples.device)
return ({"samples": noise},)
return ({"samples": _make_noise(samples.shape, seed)},)
total_frames = int(samples.size(temporal_dim))
noise = torch.zeros_like(samples)
@@ -343,22 +180,13 @@ class FunPackSceneNoise:
for i in range(num_scenes):
scene_seed = (seed + i * 31337) & 0xffffffffffffffff
rng = torch.Generator(device=samples.device)
rng.manual_seed(scene_seed)
start = i * scene_size
end = (i + 1) * scene_size if i < num_scenes - 1 else total_frames
# Build a slice for the temporal dimension
slices = [slice(None)] * samples.dim()
slices[temporal_dim] = slice(start, end)
scene_shape = list(samples.shape)
scene_shape[temporal_dim] = end - start
noise[tuple(slices)] = torch.randn(
scene_shape, generator=rng,
dtype=samples.dtype, device=samples.device,
)
noise[tuple(slices)] = _make_noise(scene_shape, scene_seed)
logging.info(
"FunPack scene noise: %s scenes, temporal_dim=%s, total_frames=%s.",