291 lines
19 KiB
Python
291 lines
19 KiB
Python
"""LTExperiments context-window nodes (V3 io).
|
|
|
|
Patches the model with this pack's forked context-window handler (see context_windows.py) for
|
|
LTXV / LTX-2 sampling.
|
|
"""
|
|
from comfy_api.latest import ComfyExtension, io
|
|
import torch
|
|
import nodes
|
|
import node_helpers
|
|
import comfy.utils
|
|
import comfy.model_management
|
|
|
|
from . import context_windows as context_windows
|
|
from . import perwindow_guides as perwindow_guides
|
|
from . import guide_cond as guide_cond
|
|
|
|
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(perwindow_guides.PERWINDOW_KEY)
|
|
if v:
|
|
return list(v)
|
|
return []
|
|
|
|
|
|
class LTEx_ContextWindowsManualNode(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id="LTEx_ContextWindowsManual",
|
|
display_name="LTEx Context Windows (Manual)",
|
|
category=CATEGORY,
|
|
description="Manually set context windows.",
|
|
inputs=[
|
|
io.Model.Input("model", tooltip="The model to apply context windows to during sampling."),
|
|
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=[
|
|
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=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."),
|
|
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."),
|
|
io.String.Input("latent_retain_index_list", default="", tooltip="List of latent indices to retain in the noise latent itself for each window. Use for workflows where reference content (e.g. a start image) lives directly in the noise latent rather than in separate conditioning channels (e.g. inplace-style I2V like LTXV, AnimateDiff). Independent of cond_retain_index_list."),
|
|
io.Boolean.Input("causal_window_fix", default=True, tooltip="Whether to add a causal fix frame to non-0-indexed context windows."),
|
|
],
|
|
outputs=[
|
|
io.Model.Output(tooltip="The model with context windows applied during sampling."),
|
|
],
|
|
is_experimental=True,
|
|
)
|
|
|
|
@classmethod
|
|
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"] = 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,
|
|
closed_loop=closed_loop,
|
|
dim=dim,
|
|
freenoise=freenoise,
|
|
cond_retain_index_list=cond_retain_index_list,
|
|
split_conds_to_windows=split_conds_to_windows,
|
|
latent_retain_index_list=latent_retain_index_list,
|
|
causal_window_fix=causal_window_fix,
|
|
)
|
|
# make memory usage calculation only take into account the context window latents
|
|
context_windows.create_prepare_sampling_wrapper(model)
|
|
if freenoise: # no other use for this wrapper at this time
|
|
context_windows.create_sampler_sample_wrapper(model)
|
|
return io.NodeOutput(model)
|
|
|
|
|
|
class LTEx_LTXVContextWindowsNode(LTEx_ContextWindowsManualNode):
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
schema = super().define_schema()
|
|
schema.node_id = "LTEx_LTXVContextWindows"
|
|
schema.display_name = "LTEx LTXV Context Windows"
|
|
schema.description = "Set context windows for LTXV / LTX-2 models."
|
|
schema.inputs = [
|
|
io.Model.Input("model", tooltip="The model to apply context windows to during sampling."),
|
|
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=[
|
|
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=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),
|
|
]
|
|
return schema
|
|
|
|
@classmethod
|
|
def execute(cls, model: io.Model.Type, context_length: int, context_overlap: int, context_schedule: str, fuse_method: str, freenoise: bool,
|
|
retain_first_frame: bool=False, split_conds_to_windows: bool=False, context_stride: int=1, closed_loop: bool=False) -> io.Model:
|
|
context_length = max(((context_length - 1) // 8) + 1, 1) # at least length 1
|
|
context_overlap = max(context_overlap // 8, 0) # at least overlap 0
|
|
retain_index_list = "0" if retain_first_frame else ""
|
|
return super().execute(model, context_length, context_overlap, context_schedule, context_stride, closed_loop, fuse_method, dim=2, freenoise=freenoise,
|
|
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, {perwindow_guides.PERWINDOW_KEY: guides}))
|
|
return io.NodeOutput(outs[0], outs[1])
|
|
|
|
|
|
class LTEx_LTXVAddGuideConditioningNode(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
from comfy_extras.nodes_lt import ICLoRAParameters # lazy: dodge custom-node load-order issues
|
|
return io.Schema(
|
|
node_id="LTEx_LTXVAddGuideConditioning",
|
|
display_name="LTEx LTXVAddGuide (Cond Only)",
|
|
category=CATEGORY,
|
|
description="Proof of concept: carry an LTXV / LTX-2 guide on the conditioning instead of "
|
|
"appending it to the working latent (the LTXVAddGuide approach). The guide is "
|
|
"injected as extra tokens inside the model forward and stripped before output, "
|
|
"so no LTXVCropGuides node is needed and the latent is never modified. Works for "
|
|
"LTXV and LTXAV (video stream). Installs a gated monkeypatch of the LTX forward "
|
|
"that is a no-op unless this node is used. Drop-in alternative to LTXVAddGuide "
|
|
"for A/B testing the conditioning-carried path.",
|
|
inputs=[
|
|
io.Conditioning.Input("positive"),
|
|
io.Conditioning.Input("negative"),
|
|
io.Vae.Input("vae"),
|
|
io.Latent.Input("latent", tooltip="The latent being sampled. Used only to size the guide encode; not modified."),
|
|
io.Image.Input("image", tooltip="Guide image or video. Must be 8*n + 1 frames or it is cropped to the nearest."),
|
|
io.Int.Input("frame_idx", default=0, min=-nodes.MAX_RESOLUTION, max=nodes.MAX_RESOLUTION,
|
|
tooltip="Pixel-frame index to place the guide at. For 9+ frame guides, frame_idx must be divisible by 8 (rounded down). Negative values count from the end."),
|
|
io.Float.Input("strength", default=1.0, min=0.0, max=1.0, step=0.01,
|
|
tooltip="How strongly the guide pins its frames. 1.0 locks them (clean), lower loosens; 0 ignores the guide."),
|
|
io.Mask.Input("attention_mask", optional=True,
|
|
tooltip="Optional pixel-space spatial mask. Controls per-region guide influence via "
|
|
"self-attention, multiplied by strength. Same semantics as LTXVAddGuide."),
|
|
ICLoRAParameters.Input("iclora_parameters", optional=True,
|
|
tooltip="Optional IC-LoRA parameters (e.g. reference_downscale_factor > 1) from a "
|
|
"Get IC-LoRA Parameters node. The low-res reference is spread across the "
|
|
"full-res grid via dilated RoPE coordinates; the latent is still untouched."),
|
|
],
|
|
outputs=[
|
|
io.Conditioning.Output(display_name="positive"),
|
|
io.Conditioning.Output(display_name="negative"),
|
|
io.Latent.Output(display_name="latent_passthrough", tooltip="Input latent, returned unmodified"),
|
|
],
|
|
is_experimental=True,
|
|
)
|
|
|
|
@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
|
|
|
|
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"]
|
|
_, _, latent_length, latent_height, latent_width = latent_image.shape
|
|
|
|
# IC-LoRA references encode at 1/N spatial resolution and are spread across the full-res grid.
|
|
latent_downscale_factor = LTXVAddGuide.get_reference_downscale_factor(iclora_parameters)
|
|
if latent_downscale_factor > 1 and (latent_width % latent_downscale_factor or latent_height % latent_downscale_factor):
|
|
raise ValueError(
|
|
f"Latent spatial size {latent_width}x{latent_height} must be divisible by "
|
|
f"reference_downscale_factor {latent_downscale_factor} from the IC-LoRA parameters."
|
|
)
|
|
|
|
# Mirror LTXVAddGuide's causal first-frame handling for mid-video multi-frame guides.
|
|
time_scale_factor = scale_factors[0]
|
|
num_frames_to_keep = ((image.shape[0] - 1) // time_scale_factor) * time_scale_factor + 1
|
|
resolved_frame_idx = frame_idx
|
|
if frame_idx < 0:
|
|
resolved_frame_idx = max((latent_length - 1) * time_scale_factor + 1 + frame_idx, 0)
|
|
causal_fix = resolved_frame_idx == 0 or num_frames_to_keep == 1
|
|
|
|
if not causal_fix:
|
|
image = torch.cat([image[:1], image], dim=0)
|
|
|
|
image, t = LTXVAddGuide.encode(vae, latent_width, latent_height, image, scale_factors, latent_downscale_factor)
|
|
|
|
if not causal_fix:
|
|
t = t[:, :, 1:, :, :]
|
|
image = image[1:]
|
|
|
|
# Snap/resolve frame_idx (num_keyframes is 0 here since nothing is added to the latent).
|
|
frame_idx, latent_idx = LTXVAddGuide.get_latent_index(
|
|
positive, latent_length, len(image), frame_idx, scale_factors, latent_shape=latent_image.shape
|
|
)
|
|
if latent_idx + t.shape[2] > latent_length:
|
|
raise ValueError("Guide frames exceed the length of the latent sequence.")
|
|
|
|
t = t.to(comfy.model_management.intermediate_device())
|
|
coords = guide_cond.compute_guide_coords(t, frame_idx, scale_factors, causal_fix, latent_downscale_factor)
|
|
|
|
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].
|
|
pre_filter_count = t.shape[2] * t.shape[3] * t.shape[4]
|
|
positive, negative = _append_guide_attention_entry(
|
|
positive, negative, pre_filter_count, list(t.shape[2:]), strength=strength, attention_mask=attention_mask,
|
|
)
|
|
return io.NodeOutput(positive, negative, latent)
|
|
|
|
|
|
class LTExperimentsExtension(ComfyExtension):
|
|
async def get_node_list(self) -> list[type[io.ComfyNode]]:
|
|
return [
|
|
LTEx_LTXVContextWindowsNode,
|
|
LTEx_LTXVAddPerWindowGuideNode,
|
|
LTEx_LTXVAddGuideConditioningNode,
|
|
]
|
|
|
|
|
|
def comfy_entrypoint():
|
|
return LTExperimentsExtension()
|