From fefe65bc9319bcb687c92144ff11d81920fc0be6 Mon Sep 17 00:00:00 2001 From: drozbay <17261091+drozbay@users.noreply.github.com> Date: Thu, 4 Jun 2026 23:34:54 -0600 Subject: [PATCH] Add per-window guide node --- README.md | 55 ++++++++----- context_windows.py | 36 ++++++++- nodes.py | 68 ++++++++++++++++ perwindow_guides.py | 186 ++++++++++++++++++++++++++++++++++++++++++++ pyproject.toml | 4 +- 5 files changed, 327 insertions(+), 22 deletions(-) create mode 100644 perwindow_guides.py diff --git a/README.md b/README.md index e15fc87..f73cf87 100644 --- a/README.md +++ b/README.md @@ -1,30 +1,47 @@ # LTExperiments -Context-window sampling for **LTXV / LTX-2 (LTXAV)** video models in native ComfyUI, including the -multimodal (video + audio) path and **IC-LoRA guide** splitting — guides placed at arbitrary positions -are spliced into each window they overlap. +Context-window sampling for LTXV and LTX-2 (LTXAV) video models in native ComfyUI, with support for the multimodal video and audio path and for splitting IC-LoRA guides across windows. -This is a self-contained extraction of the work in -[Comfy-Org/ComfyUI#13325](https://github.com/Comfy-Org/ComfyUI/pull/13325). It plugs into ComfyUI's -existing extension points (`model_options["context_handler"]` + sampling wrappers) and **does not edit -or monkey-patch core**. It works whether or not that PR is merged into your ComfyUI: when the base -model provides the guide hooks it defers to them, otherwise it uses its own bundled implementation. +This pack is a self-contained extraction of [Comfy-Org/ComfyUI#13325](https://github.com/Comfy-Org/ComfyUI/pull/13325). It uses ComfyUI's existing extension points, the `model_options["context_handler"]` slot and the sampling wrappers, and it does not edit or monkey-patch core. It works whether or not that pull request has been merged into your copy of ComfyUI. When the base model already provides the guide hooks the pack defers to them, otherwise it uses its own bundled implementation. -## Node +## Nodes -**LTEx LTXV Context Windows** — drop it between your model and the sampler. Inputs are in real (pixel) -frames; `context_length` snaps to `8n+1`, `context_overlap` to multiples of 8. `retain_first_frame` -keeps latent sub-pos 0 in both the conditioning and the noise latent (the right combo for -inplace-style LTX-2 first-frame / AnimateDiff-style I2V). Advanced settings (stride, closed_loop, -freenoise) are collapsed by default. +### LTEx LTXV Context Windows + +Turns on context-window sampling for a model. Place it between your model and the sampler. + +- Frame inputs are given in real pixel frames. `context_length` snaps to the nearest `8n+1` and `context_overlap` snaps to a multiple of 8. +- `retain_first_frame` keeps the first latent frame in both the conditioning and the noise latent for every window, which is the right setup for inplace-style LTX-2 first-frame conditioning and AnimateDiff-style image-to-video. +- Advanced settings (stride, closed loop, freenoise) are collapsed by default. + +### LTEx LTXV Add Per-Window Guide + +Inserts a guide image into every context window at the same position relative to that window. A normal LTXVAddGuide pins a guide to one absolute position in the video, so it only appears in the windows that cover it. This node repeats the guide in each window instead, which makes it a useful anchor for the start, or any chosen point, of every window. The guide rides on the conditioning and never adds frames to the working latent, so the latent stays clean between nodes. It needs the Context Windows node on the same model. + +`rel_index` sets where the guide sits inside each window, measured in latent frames: + +- `0` places it on the window's first frame. +- A positive value places it that many frames into the window. +- A negative value places it before the first frame, which is a real rope position rather than an index counted from the end. + +`strength` accepts a single value or a list of per-window strengths: + +- A single value applies to every window. +- A list applies its first value to the first window, its second to the second, and so on, reusing the last value for any further windows. The String to Float List node from ComfyUI-KJNodes is a convenient way to build one. +- A window value of `0` skips the guide for that window. +- The list maps to windows by the standard_static layout, so it stays consistent even when you sample with a uniform schedule. + +You can chain several of these nodes. Each one carries its own image, position, and strengths, and they are applied independently within each window. ## IC-LoRA guides -Guide support consumes the conds produced by the native **LTXVAddGuide** node — no extra setup beyond -adding your guide(s) as usual. Guides are stripped from the working latent, then re-injected only into -the windows they overlap, with `keyframe_idxs` / `guide_attention_entries` regenerated per window -(downscaled IC-LoRA guides included). +The pack also works with guides produced by the native LTXVAddGuide node, with no extra setup beyond adding your guides as usual. Those guides are stripped from the working latent and then reinserted only into the windows they overlap, with `keyframe_idxs` and `guide_attention_entries` regenerated for each window. Downscaled IC-LoRA guides are supported. ## Requirements -Native ComfyUI with LTXV / LTX-2 support. No extra Python dependencies. +- A native ComfyUI install with LTXV or LTX-2 support. +- No extra Python dependencies. + +## Debugging + +Set the environment variable `LTEX_DEBUG` to 1 before starting ComfyUI to log the window index, position, and strength for each per-window guide as it is injected. Leave the variable unset for normal use. diff --git a/context_windows.py b/context_windows.py index ff06fa5..ae517fb 100644 --- a/context_windows.py +++ b/context_windows.py @@ -12,6 +12,7 @@ otherwise (see ``guides``). """ from __future__ import annotations from typing import TYPE_CHECKING, Callable +import bisect import torch import logging from dataclasses import dataclass @@ -37,6 +38,7 @@ from comfy.context_windows import ( create_weights_pyramid, ) from . import guides +from . import perwindow_guides as pw if TYPE_CHECKING: from comfy.model_base import BaseModel @@ -143,6 +145,7 @@ class WindowingState: dim: int = 0 # primary modality temporal dim for context windowing is_multimodal: bool = False temporal_downscale_ratio: int = 1 # model's pixel-to-latent temporal compression ratio + perwindow_guides: list = None # prepared per-window guides injected into every window (primary modality) def prepare_window(self, window: IndexListContextWindow, model) -> IndexListContextWindow: """Reformat window for multimodal contexts by deriving per-modality index lists. @@ -191,6 +194,12 @@ class WindowingState: s, ng = self._inject_guide_frames(s, modality_window, modality_idx=idx) else: ng = 0 + # per-window guides ride on the primary modality, appended after any in-latent guides + if idx == 0 and self.perwindow_guides: + active = pw.active_guides(self.perwindow_guides, getattr(window, "canonical_index", 0)) + if active: + s = torch.cat([s, pw.window_guide_latent(active, self.dim, s.device)], dim=self.dim) + ng += pw.total_frame_count(active) sliced.append(s) guide_frame_counts.append(ng) return sliced, guide_frame_counts @@ -339,6 +348,9 @@ class IndexListContextHandler(ContextHandlerABC): extracted_guide_entries = self._get_guide_entries(conds) extracted_keyframe_idxs = self._get_keyframe_idxs(conds) + pw_raw = pw.get_perwindow_guides(conds) + perwindow_prepared = pw.prepare(model, pw_raw) if pw_raw else None + # Strip guide frames (only from first modality for now) if extracted_guide_entries is not None: guide_count = sum(e["latent_shape"][0] for e in extracted_guide_entries) @@ -358,7 +370,8 @@ class IndexListContextHandler(ContextHandlerABC): latent_shapes=latent_shapes, dim=self.dim, is_multimodal=is_multimodal, - temporal_downscale_ratio=model.latent_format.temporal_downscale_ratio) + temporal_downscale_ratio=model.latent_format.temporal_downscale_ratio, + perwindow_guides=perwindow_prepared) def should_use_context(self, model: BaseModel, conds: list[list[dict]], x_in: torch.Tensor, timestep: torch.Tensor, model_options: dict[str]) -> bool: window_state = self._build_window_state(x_in, conds, model) # build window_state to check frame counts, will be built again in execute @@ -378,6 +391,10 @@ class IndexListContextHandler(ContextHandlerABC): self.prepare_control_objects(control.previous_controlnet, device) return control + def _canonical_index(self, first_frame: int) -> int: + """Map a window's first frame to the standard_static window it starts in (its strength bucket).""" + return max(0, bisect.bisect_right(self._canonical_starts, first_frame) - 1) + def get_resized_cond(self, cond_in: list[dict], x_in: torch.Tensor, window: IndexListContextWindow, device=None) -> list: if cond_in is None: return None @@ -391,8 +408,13 @@ class IndexListContextHandler(ContextHandlerABC): # cond object is a list containing a dict - outer list is irrelevant, so just loop through it for actual_cond in cond_in: resized_actual_cond = actual_cond.copy() + # per-window guides ride as a top-level cond key; consume it here and inject into model_conds below + has_pw = pw.PERWINDOW_KEY in actual_cond and bool(getattr(self, "_perwindow_prepared", None)) + resized_actual_cond.pop(pw.PERWINDOW_KEY, None) # now we are in the inner dict - "pooled_output" is a tensor, "control" is a ControlBase object, "model_conds" is dictionary for key in actual_cond: + if key == pw.PERWINDOW_KEY: + continue try: cond_item = actual_cond[key] if isinstance(cond_item, torch.Tensor): @@ -461,6 +483,8 @@ class IndexListContextHandler(ContextHandlerABC): elif cond_key == "num_video_frames": # for SVD new_cond_item[cond_key] = cond_value._copy_with(cond_value.cond) new_cond_item[cond_key].cond = window.context_length + if has_pw: + pw.apply_to_window_cond(new_cond_item, self._perwindow_prepared, self._model, window, x_in) resized_actual_cond[key] = new_cond_item else: resized_actual_cond[key] = cond_item @@ -489,6 +513,15 @@ class IndexListContextHandler(ContextHandlerABC): self.set_step(timestep, model_options) window_state = self._build_window_state(x_in, conds, model) + self._perwindow_prepared = window_state.perwindow_guides + # Canonical (standard_static) window starts, used to map any schedule's windows to a stable + # per-window strength index via their first frame (see perwindow_guides.resolve_strength). + if window_state.perwindow_guides: + static_windows = get_matching_context_schedule(ContextSchedules.STATIC_STANDARD).func( + window_state.latents[0].size(self.dim), self, model_options) + self._canonical_starts = sorted({int(w[0]) for w in static_windows}) + else: + self._canonical_starts = [0] num_modalities = len(window_state.latents) context_windows = self.get_context_windows(model, window_state.latents[0], model_options) @@ -566,6 +599,7 @@ class IndexListContextHandler(ContextHandlerABC): # prepare the window accounting for multimodal windows window = window_state.prepare_window(window, model) + window.canonical_index = self._canonical_index(window.index_list[0]) # causal_window_fix: prepend a pre-window frame that will be stripped post-forward. # Set anchor before slice_for_window so the latent slice and downstream cond slices both pick it up. diff --git a/nodes.py b/nodes.py index 0ac8d52..0c26df5 100644 --- a/nodes.py +++ b/nodes.py @@ -5,12 +5,35 @@ LTXV / LTX-2 sampling. """ from comfy_api.latest import ComfyExtension, io import nodes +import node_helpers +import comfy.utils +import comfy.model_management from . import context_windows as cw +from . import perwindow_guides as pw CATEGORY = "LTExperiments" +def _encode_guide(vae, latent_width, latent_height, images): + """Encode guide image(s) to latent space sized to match the sampling latent (mirrors LTXVAddGuide.encode).""" + scale_factors = vae.downscale_index_formula # (time, height, width) + time_scale_factor = scale_factors[0] + images = images[:(images.shape[0] - 1) // time_scale_factor * time_scale_factor + 1] + target_w = int(latent_width * scale_factors[2]) + target_h = int(latent_height * scale_factors[1]) + pixels = comfy.utils.common_upscale(images.movedim(-1, 1), target_w, target_h, "bilinear", crop="center").movedim(1, -1) + return vae.encode(pixels[:, :, :, :3]) + + +def _existing_perwindow_guides(cond): + for t in cond: + v = t[1].get(pw.PERWINDOW_KEY) + if v: + return list(v) + return [] + + class LTEx_ContextWindowsManualNode(io.ComfyNode): @classmethod def define_schema(cls) -> io.Schema: @@ -106,10 +129,55 @@ class LTEx_LTXVContextWindowsNode(LTEx_ContextWindowsManualNode): cond_retain_index_list=retain_index_list, latent_retain_index_list=retain_index_list, split_conds_to_windows=split_conds_to_windows) +class LTEx_LTXVAddPerWindowGuideNode(io.ComfyNode): + @classmethod + def define_schema(cls) -> io.Schema: + return io.Schema( + node_id="LTEx_LTXVAddPerWindowGuide", + display_name="LTEx LTXV Add Per-Window Guide", + category=CATEGORY, + description="Inject a guide image at the same relative position in every context window. Requires the LTEx LTXV Context Windows node on the model.", + inputs=[ + io.Conditioning.Input("positive"), + io.Conditioning.Input("negative"), + io.Vae.Input("vae", tooltip="VAE used to encode the guide image."), + io.Latent.Input("latent", tooltip="The latent being sampled. Used to size the guide encode; not modified."), + io.Image.Input("image", tooltip="Guide image. Multiple frames fill consecutive positions starting at index."), + io.Int.Input("rel_index", default=-1, min=-nodes.MAX_RESOLUTION, max=nodes.MAX_RESOLUTION, + tooltip="Position of the guide within every window, counted in latent frames. 0 is the window's first frame; use negative values to place it before the first frame."), + io.Float.Input("strength", default=1.0, min=0.0, max=1.0, step=0.01, + tooltip="How strongly the guide influences each window. Accepts a single value, or a list of per-window strengths (e.g. from a String to Float List node) indexed by window order, clamping to the last value; set a window's value to 0 to skip the guide for that window."), + io.Mask.Input("attention_mask", optional=True, tooltip="Optional mask limiting where the guide applies."), + ], + outputs=[ + io.Conditioning.Output(display_name="positive"), + io.Conditioning.Output(display_name="negative"), + ], + is_experimental=True, + ) + + @classmethod + def execute(cls, positive, negative, vae, latent, image, rel_index: int, strength, attention_mask=None) -> io.NodeOutput: + _, _, _, latent_height, latent_width = latent["samples"].shape + t = _encode_guide(vae, latent_width, latent_height, image).to(comfy.model_management.intermediate_device()) + guide = { + "latent": t, + "rel_index": int(rel_index), + "strength": [float(x) for x in strength] if isinstance(strength, (list, tuple)) else float(strength), + "pixel_mask": attention_mask.unsqueeze(0).unsqueeze(0) if attention_mask is not None else None, + } + outs = [] + for cond in (positive, negative): + guides = [*_existing_perwindow_guides(cond), guide] + outs.append(node_helpers.conditioning_set_values(cond, {pw.PERWINDOW_KEY: guides})) + return io.NodeOutput(outs[0], outs[1]) + + class LTExperimentsExtension(ComfyExtension): async def get_node_list(self) -> list[type[io.ComfyNode]]: return [ LTEx_LTXVContextWindowsNode, + LTEx_LTXVAddPerWindowGuideNode, ] diff --git a/perwindow_guides.py b/perwindow_guides.py new file mode 100644 index 0000000..b9d8719 --- /dev/null +++ b/perwindow_guides.py @@ -0,0 +1,186 @@ +"""Per-window guide injection for LTXV context windows. + +A per-window guide is a clean guide latent injected into each context window at a fixed +window-relative position, rather than at one absolute position like a normal LTXVAddGuide guide. +Its per-window strength can vary, and a strength of 0 skips a given window. It rides on the +conditioning (key ``PERWINDOW_KEY``) and never touches the working latent. The engine appends it +to a window's token sequence, manufactures the matching ``keyframe_idxs``, +``guide_attention_entries``, and ``denoise_mask`` for that window, then strips it back off the +output. The guide is physically a tail token, and its position is purely its RoPE coordinate, so +``rel_index`` may be negative (a coordinate before the window's first frame). +""" +from __future__ import annotations +import os +import logging +import torch +import comfy.conds +from comfy.ldm.lightricks.symmetric_patchifier import latent_to_pixel_coords + +PERWINDOW_KEY = "ltex_perwindow_guides" + +LTEX_DEBUG = bool(os.environ.get("LTEX_DEBUG")) # off by default; set LTEX_DEBUG=1 to trace per-window guide injection + + +def _dbg(msg): + if LTEX_DEBUG: + logging.info("[LTEx perwindow] " + msg) + + +def get_perwindow_guides(conds) -> list | None: + """Return the raw per-window guide dicts carried on the conditioning, if any. + + The guides ride as a top-level cond-dict key (set via conditioning_set_values); unrecognized + keys are not wrapped into model_conds by extra_conds, so they live at the top level here. + """ + for cond_list in conds: + if cond_list is None: + continue + for cond_dict in cond_list: + entry = cond_dict.get(PERWINDOW_KEY) + if entry: + return entry + return None + + +def prepare(model, guides: list) -> list: + """Move guide latents into model latent space once, caching alongside the metadata. + + Returns a list of dicts: {latent (process_latent_in'd, on intermediate device), rel_index, + strength, pixel_mask, latent_shape [F,H,W]}. + """ + prepared = [] + for g in guides: + latent = model.process_latent_in(g["latent"]) + prepared.append({ + "latent": latent, + "rel_index": int(g["rel_index"]), + "strength": g.get("strength", 1.0), # float, or a per-window list of floats + "pixel_mask": g.get("pixel_mask"), + "latent_shape": list(latent.shape[2:]), # [F, H, W] + }) + return prepared + + +def resolve_strength(guide, canonical_index: int) -> float: + """Resolve a guide's strength for a given canonical (standard_static) window index. + + A scalar applies to every window; a list indexes by canonical window, clamping to the last entry. + """ + s = guide["strength"] + if isinstance(s, (list, tuple)): + if not s: + return 1.0 + return float(s[min(int(canonical_index), len(s) - 1)]) + return float(s) + + +def active_guides(prepared: list, canonical_index: int) -> list: + """Guides whose resolved strength for this window is > 0; the rest are skipped entirely.""" + return [p for p in prepared if resolve_strength(p, canonical_index) > 0.0] + + +def total_frame_count(prepared: list) -> int: + return sum(p["latent_shape"][0] for p in prepared) + + +def window_guide_latent(prepared: list, dim: int, device) -> torch.Tensor: + """Concatenated guide latents to append to a window's primary-modality slice.""" + frames = [p["latent"].to(device) for p in prepared] + return torch.cat(frames, dim=dim) + + +def _guide_pixel_coords(model, guide_latent, rel_index, device): + """Pixel (RoPE) coords for one guide latent placed at window-local position rel_index. + + Mirrors LTXVAddGuide.add_keyframe_index: patchify for the base grid, then offset the temporal + channel by rel_index in pixel space (one latent step == time scale factor), bypassing the causal + clamp so negative positions survive. rel_index is measured from the window's first slot, so a + causal_window_fix anchor (which occupies that slot) sits between rel_index 0 and -1: the guide at + -1 lands one step before the anchor, not on it. + """ + patchifier = model.diffusion_model.patchifier + scale_factors = model.diffusion_model.vae_scale_factors + time_scale_factor = scale_factors[0] + _, latent_coords = patchifier.patchify(guide_latent.to(device)) + # causal_fix only affects the base frame-0 token; the explicit offset below carries position. + pixel_coords = latent_to_pixel_coords(latent_coords, scale_factors, causal_fix=False) + pixel_coords[:, 0] += rel_index * time_scale_factor + return pixel_coords + + +def apply_to_window_cond(new_cond_item: dict, prepared: list, model, window, x_in: torch.Tensor): + """Extend (or create) keyframe_idxs / guide_attention_entries / denoise_mask for a window so + the appended per-window guide tokens are treated as clean, RoPE-positioned guide tokens. + + Order is [existing guide tokens][per-window guide tokens], matching the latent append order. + """ + device = x_in.device # build every per-window cond on the working-latent device + dim = window.dim + H, W = x_in.shape[3], x_in.shape[4] + anchor_idx = getattr(window, "causal_anchor_index", None) + anchor_shift = 1 if (anchor_idx is not None and anchor_idx >= 0) else 0 + # window-local video length (frames the model will denoise), excluding any appended guides + video_len = len(window.index_list) + anchor_shift + + ci = getattr(window, "canonical_index", 0) + active = active_guides(prepared, ci) + if not active: + _dbg(f"window {window.index_list[0]}-{window.index_list[-1]} canonical={ci} | no guide (all strengths 0)") + return # every guide is strength 0 for this window: inject nothing, leave conds untouched + + win = f"{window.index_list[0]}-{window.index_list[-1]}" + coords = [] + entries = [] + strengths = [] + for i, p in enumerate(active): + s = resolve_strength(p, ci) + strengths.append(s) + c = _guide_pixel_coords(model, p["latent"], p["rel_index"], device) + coords.append(c) + pixel_mask = p["pixel_mask"] + entries.append({ + "pre_filter_count": p["latent_shape"][0] * H * W, + "strength": s, + "pixel_mask": pixel_mask.to(device) if isinstance(pixel_mask, torch.Tensor) else None, + "latent_shape": [p["latent_shape"][0], H, W], + }) + if LTEX_DEBUG: + _dbg(f"window {win} canonical={ci} | guide#{i} rel_index={p['rel_index']} " + f"coord_t={c[0, 0, 0, 0].item():.1f} strength={s:.3f}") + pw_coords = torch.cat(coords, dim=2) + + # keyframe_idxs: append our coords after any existing (already window-local) guide coords. + existing_kf = new_cond_item.get("keyframe_idxs") + if existing_kf is not None and hasattr(existing_kf, "cond") and existing_kf.cond is not None: + merged = torch.cat([existing_kf.cond.to(device), pw_coords], dim=2) + new_cond_item["keyframe_idxs"] = existing_kf._copy_with(merged) + else: + new_cond_item["keyframe_idxs"] = comfy.conds.CONDRegular(pw_coords) + + # guide_attention_entries: append after any existing entries. + existing_entries = new_cond_item.get("guide_attention_entries") + if existing_entries is not None and hasattr(existing_entries, "cond") and existing_entries.cond: + merged_entries = [*existing_entries.cond, *entries] + new_cond_item["guide_attention_entries"] = existing_entries._copy_with(merged_entries) + else: + new_cond_item["guide_attention_entries"] = comfy.conds.CONDConstant(entries) + + # denoise_mask: append a keep-mask (1 - strength) for our tokens; build a video portion if none. + B = x_in.shape[0] + pw_mask_parts = [ + torch.full((B, 1, p["latent_shape"][0], H, W), max(0.0, 1.0 - s), + device=device, dtype=x_in.dtype) + for p, s in zip(active, strengths) + ] + pw_mask = torch.cat(pw_mask_parts, dim=dim) + existing_dm = new_cond_item.get("denoise_mask") + if existing_dm is not None and hasattr(existing_dm, "cond") and isinstance(existing_dm.cond, torch.Tensor): + base = existing_dm.cond.to(device) + if base.shape[3] != H or base.shape[4] != W: + base = base.expand(-1, -1, -1, H, W) + merged_mask = torch.cat([base, pw_mask], dim=dim) + new_cond_item["denoise_mask"] = existing_dm._copy_with(merged_mask) + else: + video_mask = torch.ones((B, 1, video_len, H, W), device=device, dtype=x_in.dtype) + merged_mask = torch.cat([video_mask, pw_mask], dim=dim) + new_cond_item["denoise_mask"] = comfy.conds.CONDRegular(merged_mask) diff --git a/pyproject.toml b/pyproject.toml index 6cd3c6f..4b65375 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "ltexperiments" -description = "Context window sampling for LTXV / LTX-2 (LTXAV) video+audio models, including IC-LoRA guide splitting across windows. Self-contained; uses ComfyUI's context_handler and wrapper hooks (no core edits)." -version = "0.1.0" +description = "Context window sampling for LTXV / LTX-2 (LTXAV) video and audio models, including IC-LoRA guide splitting and per-window guide insertion across context windows. Self-contained and uses ComfyUI's context_handler and wrapper hooks without editing core." +version = "0.2.0" license = {file = "LICENSE"} [project.urls]