Gate conditioning-carried guides with manual override key

This commit is contained in:
drozbay
2026-06-08 10:58:01 -06:00
parent 51cdcf8dbc
commit b4f5aff476
3 changed files with 49 additions and 46 deletions
+10 -10
View File
@@ -38,7 +38,7 @@ from comfy.context_windows import (
create_weights_pyramid,
)
from . import guides
from . import perwindow_guides as pw
from . import perwindow_guides as perwindow_guides
if TYPE_CHECKING:
from comfy.model_base import BaseModel
@@ -196,10 +196,10 @@ class WindowingState:
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))
active = perwindow_guides.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)
s = torch.cat([s, perwindow_guides.window_guide_latent(active, self.dim, s.device)], dim=self.dim)
ng += perwindow_guides.total_frame_count(active)
sliced.append(s)
guide_frame_counts.append(ng)
return sliced, guide_frame_counts
@@ -348,8 +348,8 @@ 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
pw_raw = perwindow_guides.get_perwindow_guides(conds)
perwindow_prepared = perwindow_guides.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:
@@ -409,11 +409,11 @@ class IndexListContextHandler(ContextHandlerABC):
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)
has_pw = perwindow_guides.PERWINDOW_KEY in actual_cond and bool(getattr(self, "_perwindow_prepared", None))
resized_actual_cond.pop(perwindow_guides.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:
if key == perwindow_guides.PERWINDOW_KEY:
continue
try:
cond_item = actual_cond[key]
@@ -484,7 +484,7 @@ class IndexListContextHandler(ContextHandlerABC):
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)
perwindow_guides.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
+9 -10
View File
@@ -29,10 +29,9 @@ GUIDE_COORDS_KEY = "guide_cond_keyframe_idxs"
GUIDE_STRENGTH_KEY = "guide_cond_strength"
# Standalone patchifier for node-side coord computation (matches LTXVModel.patchifier config).
_PATCHIFIER = SymmetricPatchifier(1, start_end=True)
_PATCHED = False
LTXCONDGUIDE_ENABLED = False # Set to False by default. You must enable this manually for now.
LTXCONDGUIDE_PATCHIFIER = SymmetricPatchifier(1, start_end=True)
LTXCONDGUIDE_PATCHED = False
# --------------------------------------------------------------------------------------------
# Node-side helpers (model-free): coordinates + conditioning accumulation
@@ -53,7 +52,7 @@ def compute_guide_coords(guide_latent, frame_idx, scale_factors, causal_fix, lat
if latent_downscale_factor > 1:
from comfy_extras.nodes_lt import LTXVAddGuide
dilated, dilated_mask = LTXVAddGuide.dilate_latent(guide_latent, latent_downscale_factor)
_, latent_coords = _PATCHIFIER.patchify(dilated)
_, latent_coords = LTXCONDGUIDE_PATCHIFIER.patchify(dilated)
pixel_coords = latent_to_pixel_coords(latent_coords, scale_factors, causal_fix=causal_fix)
pixel_coords[:, 0] += frame_idx
spatial_end_offset = (latent_downscale_factor - 1) * torch.tensor(
@@ -61,11 +60,11 @@ def compute_guide_coords(guide_latent, frame_idx, scale_factors, causal_fix, lat
).view(1, -1, 1, 1)
pixel_coords[:, 1:, :, 1:] += spatial_end_offset.to(pixel_coords.dtype)
# Holes are marked < 0 in the dilated mask; keep only the real tokens.
mask_tokens, _ = _PATCHIFIER.patchify(dilated_mask)
mask_tokens, _ = LTXCONDGUIDE_PATCHIFIER.patchify(dilated_mask)
keep = mask_tokens[0, :, 0] > 0
return pixel_coords[:, :, keep, :]
_, latent_coords = _PATCHIFIER.patchify(guide_latent)
_, latent_coords = LTXCONDGUIDE_PATCHIFIER.patchify(guide_latent)
pixel_coords = latent_to_pixel_coords(latent_coords, scale_factors, causal_fix=causal_fix)
pixel_coords[:, 0] += frame_idx
return pixel_coords
@@ -236,8 +235,8 @@ def _make_patched_extra_conds(orig_extra_conds):
def install():
"""Install the gated patches once. Idempotent across module reloads (guarded by a marker
attribute on the patched callables)."""
global _PATCHED
if _PATCHED:
global LTXCONDGUIDE_PATCHED
if LTXCONDGUIDE_PATCHED or not LTXCONDGUIDE_ENABLED:
return
if not getattr(ltx_model.LTXBaseModel._forward, "_ltex_guide_cond", False):
@@ -249,4 +248,4 @@ def install():
if not getattr(model_cls.extra_conds, "_ltex_guide_cond", False):
model_cls.extra_conds = _make_patched_extra_conds(model_cls.extra_conds)
_PATCHED = True
LTXCONDGUIDE_PATCHED = True
+30 -26
View File
@@ -10,9 +10,9 @@ import node_helpers
import comfy.utils
import comfy.model_management
from . import context_windows as cw
from . import perwindow_guides as pw
from . import guide_cond as gc
from . import context_windows as context_windows
from . import perwindow_guides as perwindow_guides
from . import guide_cond as guide_cond
CATEGORY = "LTExperiments"
@@ -30,7 +30,7 @@ def _encode_guide(vae, latent_width, latent_height, images):
def _existing_perwindow_guides(cond):
for t in cond:
v = t[1].get(pw.PERWINDOW_KEY)
v = t[1].get(perwindow_guides.PERWINDOW_KEY)
if v:
return list(v)
return []
@@ -49,14 +49,14 @@ class LTEx_ContextWindowsManualNode(io.ComfyNode):
io.Int.Input("context_length", min=1, default=16, tooltip="The length of the context window."),
io.Int.Input("context_overlap", min=0, default=4, tooltip="The overlap of the context window."),
io.Combo.Input("context_schedule", options=[
cw.ContextSchedules.STATIC_STANDARD,
cw.ContextSchedules.UNIFORM_STANDARD,
cw.ContextSchedules.UNIFORM_LOOPED,
cw.ContextSchedules.BATCHED,
], default=cw.ContextSchedules.STATIC_STANDARD, tooltip="Step-dependent scheduling algorithm for context windows."),
context_windows.ContextSchedules.STATIC_STANDARD,
context_windows.ContextSchedules.UNIFORM_STANDARD,
context_windows.ContextSchedules.UNIFORM_LOOPED,
context_windows.ContextSchedules.BATCHED,
], default=context_windows.ContextSchedules.STATIC_STANDARD, tooltip="Step-dependent scheduling algorithm for context windows."),
io.Int.Input("context_stride", min=1, default=1, tooltip="The stride of the context window; only applicable to uniform schedules."),
io.Boolean.Input("closed_loop", default=False, tooltip="Whether to close the context window loop; only applicable to looped schedules."),
io.Combo.Input("fuse_method", options=cw.ContextFuseMethods.LIST_STATIC, default=cw.ContextFuseMethods.PYRAMID, tooltip="The method to use to fuse the context windows."),
io.Combo.Input("fuse_method", options=context_windows.ContextFuseMethods.LIST_STATIC, default=context_windows.ContextFuseMethods.PYRAMID, tooltip="The method to use to fuse the context windows."),
io.Int.Input("dim", min=0, max=5, default=0, tooltip="The dimension to apply the context windows to."),
io.Boolean.Input("freenoise", default=False, tooltip="Whether to apply FreeNoise noise shuffling, improves window blending."),
io.String.Input("cond_retain_index_list", default="", tooltip="List of latent indices to retain in the conditioning tensors for each window. For concat-style I2V models (e.g. Wan I2V, HunyuanVideo I2V, Cosmos I2V, SVD) the encoded start image lives in the c_concat conditioning channels; setting this to '0' will retain that start image content at sub-pos 0 of every window."),
@@ -74,9 +74,9 @@ class LTEx_ContextWindowsManualNode(io.ComfyNode):
def execute(cls, model: io.Model.Type, context_length: int, context_overlap: int, context_schedule: str, context_stride: int, closed_loop: bool, fuse_method: str, dim: int, freenoise: bool,
cond_retain_index_list: list[int]=[], split_conds_to_windows: bool=False, latent_retain_index_list: list[int]=[], causal_window_fix: bool=True) -> io.Model:
model = model.clone()
model.model_options["context_handler"] = cw.IndexListContextHandler(
context_schedule=cw.get_matching_context_schedule(context_schedule),
fuse_method=cw.get_matching_fuse_method(fuse_method),
model.model_options["context_handler"] = context_windows.IndexListContextHandler(
context_schedule=context_windows.get_matching_context_schedule(context_schedule),
fuse_method=context_windows.get_matching_fuse_method(fuse_method),
context_length=context_length,
context_overlap=context_overlap,
context_stride=context_stride,
@@ -89,9 +89,9 @@ class LTEx_ContextWindowsManualNode(io.ComfyNode):
causal_window_fix=causal_window_fix,
)
# make memory usage calculation only take into account the context window latents
cw.create_prepare_sampling_wrapper(model)
context_windows.create_prepare_sampling_wrapper(model)
if freenoise: # no other use for this wrapper at this time
cw.create_sampler_sample_wrapper(model)
context_windows.create_sampler_sample_wrapper(model)
return io.NodeOutput(model)
@@ -107,14 +107,14 @@ class LTEx_LTXVContextWindowsNode(LTEx_ContextWindowsManualNode):
io.Int.Input("context_length", min=1, max=nodes.MAX_RESOLUTION, step=8, default=145, tooltip="The length of the context window in real frames. Must be 8*n + 1."),
io.Int.Input("context_overlap", min=0, step=8, default=40, tooltip="The overlap of the context window in real frames."),
io.Combo.Input("context_schedule", options=[
cw.ContextSchedules.STATIC_STANDARD,
cw.ContextSchedules.UNIFORM_STANDARD,
cw.ContextSchedules.UNIFORM_LOOPED,
cw.ContextSchedules.BATCHED,
], default=cw.ContextSchedules.UNIFORM_STANDARD, tooltip="Step-dependent scheduling algorithm for context windows."),
context_windows.ContextSchedules.STATIC_STANDARD,
context_windows.ContextSchedules.UNIFORM_STANDARD,
context_windows.ContextSchedules.UNIFORM_LOOPED,
context_windows.ContextSchedules.BATCHED,
], default=context_windows.ContextSchedules.UNIFORM_STANDARD, tooltip="Step-dependent scheduling algorithm for context windows."),
io.Int.Input("context_stride", min=1, default=1, tooltip="The stride of the context window; only applicable to uniform schedules.", advanced=True),
io.Boolean.Input("closed_loop", default=False, tooltip="Whether to close the context window loop; only applicable to looped schedules.", advanced=True),
io.Combo.Input("fuse_method", options=cw.ContextFuseMethods.LIST_STATIC, default=cw.ContextFuseMethods.PYRAMID, tooltip="The method to use to fuse the context windows."),
io.Combo.Input("fuse_method", options=context_windows.ContextFuseMethods.LIST_STATIC, default=context_windows.ContextFuseMethods.PYRAMID, tooltip="The method to use to fuse the context windows."),
io.Boolean.Input("freenoise", default=True, tooltip="Whether to apply FreeNoise noise shuffling, improves window blending.", advanced=True),
io.Boolean.Input("retain_first_frame", default=False, tooltip="Retain the first latent frame in every context window (may help retain initial reference)."),
io.Boolean.Input("split_conds_to_windows", default=False, tooltip="Whether to split multiple conditionings (created by ConditionCombine) to each window based on region index.", advanced=True),
@@ -171,7 +171,7 @@ class LTEx_LTXVAddPerWindowGuideNode(io.ComfyNode):
outs = []
for cond in (positive, negative):
guides = [*_existing_perwindow_guides(cond), guide]
outs.append(node_helpers.conditioning_set_values(cond, {pw.PERWINDOW_KEY: guides}))
outs.append(node_helpers.conditioning_set_values(cond, {perwindow_guides.PERWINDOW_KEY: guides}))
return io.NodeOutput(outs[0], outs[1])
@@ -218,9 +218,13 @@ class LTEx_LTXVAddGuideConditioningNode(io.ComfyNode):
@classmethod
def execute(cls, positive, negative, vae, latent, image, frame_idx: int, strength: float, attention_mask=None, iclora_parameters=None) -> io.NodeOutput:
from guide_cond import LTXCONDGUIDE_ENABLED
if not LTXCONDGUIDE_ENABLED:
raise RuntimeError("The guide on conditioning functionality is gated behind the LTXCONDGUIDE_ENABLED variable in guide_cond.py. Set it to True to enable this node.")
from comfy_extras.nodes_lt import LTXVAddGuide, _append_guide_attention_entry # reuse core's encode + frame-index snapping + attn entries
gc.install() # idempotent; gated monkeypatch is a no-op unless guide keys are present
guide_cond.install() # idempotent; gated monkeypatch is a no-op unless guide keys are present
scale_factors = vae.downscale_index_formula
latent_image = latent["samples"]
@@ -259,10 +263,10 @@ class LTEx_LTXVAddGuideConditioningNode(io.ComfyNode):
raise ValueError("Guide frames exceed the length of the latent sequence.")
t = t.to(comfy.model_management.intermediate_device())
coords = gc.compute_guide_coords(t, frame_idx, scale_factors, causal_fix, latent_downscale_factor)
coords = guide_cond.compute_guide_coords(t, frame_idx, scale_factors, causal_fix, latent_downscale_factor)
positive = gc.append_guide(positive, t, coords, float(strength))
negative = gc.append_guide(negative, t, coords, float(strength))
positive = guide_cond.append_guide(positive, t, coords, float(strength))
negative = guide_cond.append_guide(negative, t, coords, float(strength))
# Per-guide attention control: reuse core's entry helper. No latent dilation here, so
# pre_filter_count is just the guide's token count and latent_shape is its [F, H, W].