diff --git a/context_windows.py b/context_windows.py index ae517fb..ac74801 100644 --- a/context_windows.py +++ b/context_windows.py @@ -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 diff --git a/guide_cond.py b/guide_cond.py index 3c12ce9..1e86110 100644 --- a/guide_cond.py +++ b/guide_cond.py @@ -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 diff --git a/nodes.py b/nodes.py index b0e44ef..b7b06be 100644 --- a/nodes.py +++ b/nodes.py @@ -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].