From 9b9271be2113fac78c58eb0f2492895ad3dc8b0a Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Wed, 20 May 2026 18:39:39 +0300 Subject: [PATCH] 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 --- __init__.py | 7 +- context_transition.py | 196 +++--------------------------------------- 2 files changed, 14 insertions(+), 189 deletions(-) diff --git a/__init__.py b/__init__.py index feb20388..9d10a42a 100644 --- a/__init__.py +++ b/__init__.py @@ -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", diff --git a/context_transition.py b/context_transition.py index 3e636b77..48177a18 100644 --- a/context_transition.py +++ b/context_transition.py @@ -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.",