From f751d76cfadeee6dd2caff52fd336c2cb6833358 Mon Sep 17 00:00:00 2001 From: kijai <40791699+kijai@users.noreply.github.com> Date: Thu, 30 Apr 2026 01:38:18 +0300 Subject: [PATCH] Add some more knobs to play with --- advanced_options.py | 58 +++++++++++++++++++++++++++++++++++++++++++++ nodes.py | 23 +++++++++++++----- prompt_relay.py | 48 ++++++++++++++++++++++++++++++++----- 3 files changed, 117 insertions(+), 12 deletions(-) create mode 100644 advanced_options.py diff --git a/advanced_options.py b/advanced_options.py new file mode 100644 index 0000000..72a2ce9 --- /dev/null +++ b/advanced_options.py @@ -0,0 +1,58 @@ +from comfy_api.latest import io + + +RelayOptions = io.Custom("RELAY_OPTIONS") + + + +class PromptRelayAdvancedOptions(io.ComfyNode): + """Per-stream temporal-penalty knobs for Prompt Relay encoders.""" + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="PromptRelayAdvancedOptions", + display_name="Prompt Relay Advanced Options", + category="conditioning/prompt_relay", + is_experimental=True, + description=( + "Optional per-stream tuning. Connect to the relay_options input on a Prompt Relay encoder. Audio fields only affect LTX2." + ), + inputs=[ + io.Float.Input( + "video_strength", default=1.0, min=0.0, max=10.0, step=0.05, + tooltip="Multiplier on the video temporal penalty. 0 disables video segmentation. Most useful in the 0–1 range to soften boundaries; values >1 saturate quickly at the default epsilon (1e-3) because the bias is already large enough to zero out distant tokens. To make >1 visibly meaningful, raise epsilon to ~0.1 or higher.", + ), + io.Float.Input( + "video_window_scale", default=1.0, min=0.0, max=4.0, step=0.05, + tooltip="Scales the flat anchor zone (default L/2 - 2 frames). <1 starts the soft falloff sooner; >1 widens the rigid zone. 0 collapses the anchor to a point — falloff begins immediately at the segment midpoint (sharper than default, not softer).", + ), + io.Float.Input( + "audio_epsilon", default=0.0, min=0.0, max=0.99, step=1e-4, + tooltip="Epsilon for the audio stream. 0 = inherit from the encoder's main epsilon.", + ), + io.Float.Input( + "audio_strength", default=1.0, min=0.0, max=10.0, step=0.05, + tooltip="Multiplier on the audio temporal penalty. 0 lets audio bleed across visual cuts. Most useful in the 0–1 range; values >1 saturate quickly at the default epsilon (1e-3) — raise audio_epsilon to ~0.1 or higher to make >1 visibly meaningful.", + ), + io.Float.Input( + "audio_window_scale", default=1.0, min=0.0, max=4.0, step=0.05, + tooltip="Scales the flat anchor zone width for the audio stream. <1 starts the soft falloff sooner; >1 widens the rigid zone. 0 collapses the anchor to a point — falloff begins immediately at the segment midpoint (sharper than default, not softer).", + ), + ], + outputs=[ + RelayOptions.Output(display_name="relay_options"), + ], + ) + + @classmethod + def execute(cls, video_strength, video_window_scale, + audio_epsilon, audio_strength, audio_window_scale) -> io.NodeOutput: + opts = { + "video_strength": video_strength, + "video_window_scale": video_window_scale, + "audio_epsilon": audio_epsilon if audio_epsilon > 0 else None, + "audio_strength": audio_strength, + "audio_window_scale": audio_window_scale, + } + return io.NodeOutput(opts) diff --git a/nodes.py b/nodes.py index dd8e330..b8749a6 100644 --- a/nodes.py +++ b/nodes.py @@ -11,6 +11,7 @@ from .prompt_relay import ( ) from .patches import detect_model_type, apply_patches +from .advanced_options import PromptRelayAdvancedOptions, RelayOptions log = logging.getLogger(__name__) @@ -52,7 +53,7 @@ def _convert_to_latent_lengths(pixel_lengths, temporal_stride, latent_frames): return result -def _encode_relay(model, clip, latent, global_prompt, local_prompts, segment_lengths, epsilon): +def _encode_relay(model, clip, latent, global_prompt, local_prompts, segment_lengths, epsilon, relay_options=None): locals_list = [p.strip() for p in local_prompts.split("|") if p.strip()] if not locals_list: raise ValueError("At least one local prompt is required (separate with |)") @@ -84,7 +85,7 @@ def _encode_relay(model, clip, latent, global_prompt, local_prompts, segment_len latent_frames, tokens_per_frame, effective_lengths, ) - q_token_idx = build_segments(token_ranges, effective_lengths, epsilon) + q_token_idx = build_segments(token_ranges, effective_lengths, epsilon, relay_options) mask_fn = create_mask_fn(q_token_idx, tokens_per_frame, latent_frames) patched = model.clone() @@ -127,6 +128,10 @@ class PromptRelayEncode(io.ComfyNode): "epsilon", default=1e-3, min=1e-6, max=0.99, step=1e-4, tooltip="Penalty decay parameter. Values below ~0.1 all produce sharp boundaries (paper default 0.001). For softer transitions, try 0.5 or higher.", ), + RelayOptions.Input( + "relay_options", optional=True, + tooltip="Optional advanced per-stream tuning. Connect a Prompt Relay Advanced Options node.", + ), ], outputs=[ io.Model.Output(display_name="model"), @@ -135,9 +140,9 @@ class PromptRelayEncode(io.ComfyNode): ) @classmethod - def execute(cls, model, clip, latent, global_prompt, local_prompts, segment_lengths, epsilon) -> io.NodeOutput: + def execute(cls, model, clip, latent, global_prompt, local_prompts, segment_lengths, epsilon, relay_options=None) -> io.NodeOutput: patched, conditioning = _encode_relay( - model, clip, latent, global_prompt, local_prompts, segment_lengths, epsilon, + model, clip, latent, global_prompt, local_prompts, segment_lengths, epsilon, relay_options, ) return io.NodeOutput(patched, conditioning) @@ -192,6 +197,10 @@ class PromptRelayEncodeTimeline(io.ComfyNode): "time_units", options=["frames", "seconds"], default="frames", optional=True, tooltip="Display the ruler, segment ranges, length input, and total in frames or seconds. Internal storage is always pixel-space frames.", ), + RelayOptions.Input( + "relay_options", optional=True, + tooltip="Optional advanced per-stream tuning. Connect a Prompt Relay Advanced Options node.", + ), ], outputs=[ io.Model.Output(display_name="model"), @@ -201,9 +210,9 @@ class PromptRelayEncodeTimeline(io.ComfyNode): @classmethod - def execute(cls, model, clip, latent, global_prompt, max_frames, timeline_data, local_prompts, segment_lengths, epsilon, fps=24.0, time_units="frames") -> io.NodeOutput: + def execute(cls, model, clip, latent, global_prompt, max_frames, timeline_data, local_prompts, segment_lengths, epsilon, fps=24.0, time_units="frames", relay_options=None) -> io.NodeOutput: patched, conditioning = _encode_relay( - model, clip, latent, global_prompt, local_prompts, segment_lengths, epsilon, + model, clip, latent, global_prompt, local_prompts, segment_lengths, epsilon, relay_options, ) return io.NodeOutput(patched, conditioning) @@ -211,9 +220,11 @@ class PromptRelayEncodeTimeline(io.ComfyNode): NODE_CLASS_MAPPINGS = { "PromptRelayEncode": PromptRelayEncode, "PromptRelayEncodeTimeline": PromptRelayEncodeTimeline, + "PromptRelayAdvancedOptions": PromptRelayAdvancedOptions, } NODE_DISPLAY_NAME_MAPPINGS = { "PromptRelayEncode": "Prompt Relay Encode", "PromptRelayEncodeTimeline": "Prompt Relay Encode (Timeline)", + "PromptRelayAdvancedOptions": "Prompt Relay Advanced Options", } diff --git a/prompt_relay.py b/prompt_relay.py index a3d2b2a..adfffe0 100644 --- a/prompt_relay.py +++ b/prompt_relay.py @@ -13,7 +13,8 @@ def build_temporal_cost(q_token_idx, Lq, Lk, device, dtype, tokens_per_frame): for seg in q_token_idx: local = seg["local_token_idx"].to(device=device) d = (query_frames.float()[:, None] - seg["midpoint"]).abs() - cost = (torch.relu(d - seg["window"]) ** 2) / (2 * seg["sigma"] ** 2) + strength = seg.get("strength", 1.0) + cost = strength * (torch.relu(d - seg["window"]) ** 2) / (2 * seg["sigma"] ** 2) offset[:, local] = cost.to(offset.dtype) return offset @@ -27,7 +28,10 @@ def build_temporal_cost_scaled(q_token_idx, Lq, Lk, device, dtype, latent_frames for seg in q_token_idx: local = seg["local_token_idx"].to(device=device) d = (query_frames[:, None] - seg["midpoint"]).abs() - cost = (torch.relu(d - seg["window"]) ** 2) / (2 * seg["sigma"] ** 2) + sigma_a = seg.get("sigma_audio", seg["sigma"]) + window_a = seg.get("window_audio", seg["window"]) + strength_a = seg.get("strength_audio", 1.0) + cost = strength_a * (torch.relu(d - window_a) ** 2) / (2 * sigma_a ** 2) offset[:, local] = cost.to(offset.dtype) return offset @@ -76,11 +80,39 @@ def create_mask_fn(q_token_idx, fallback_tokens_per_frame, latent_frames): return mask_fn -def build_segments(token_ranges, segment_lengths, epsilon=1e-3): - """Per-segment metadata (local_token_idx, midpoint, window, sigma) for the temporal penalty.""" +def build_segments(token_ranges, segment_lengths, epsilon=1e-3, relay_options=None): + """Per-segment metadata for the temporal penalty. + + relay_options (optional dict) overrides per-stream knobs: + video_strength, video_window_scale, + audio_epsilon, audio_strength, audio_window_scale + Audio knobs only affect architectures whose cross-attention takes the scaled + (non-integer-frame) path — currently LTX audio_attn2. + """ # Paper uses constant sigma = 1/ln(1/epsilon) regardless of segment length sigma = 1.0 / math.log(1.0 / epsilon) if 0 < epsilon < 1 else 0.1448 + opts = relay_options or {} + v_strength = opts.get("video_strength", 1.0) + v_window_scale = opts.get("video_window_scale", 1.0) + a_epsilon = opts.get("audio_epsilon") + a_strength = opts.get("audio_strength", 1.0) + a_window_scale = opts.get("audio_window_scale", 1.0) + + if a_epsilon is not None and 0 < a_epsilon < 1: + sigma_audio = 1.0 / math.log(1.0 / a_epsilon) + else: + sigma_audio = sigma + + if relay_options: + log.info( + "[PromptRelay] Advanced options active — video: strength=%.3f window_scale=%.3f | " + "audio: epsilon=%s strength=%.3f window_scale=%.3f", + v_strength, v_window_scale, + f"{a_epsilon:.4f}" if a_epsilon is not None else "inherit", + a_strength, a_window_scale, + ) + q_token_idx = [] frame_cursor = 0 @@ -89,12 +121,16 @@ def build_segments(token_ranges, segment_lengths, epsilon=1e-3): frame_cursor += L continue midpoint = (2 * frame_cursor + L) // 2 - window = max(L // 2 - 2, 0) + base_window = max(L // 2 - 2, 0) q_token_idx.append({ "local_token_idx": torch.arange(tok_start, tok_end), "midpoint": midpoint, - "window": window, + "window": max(base_window * v_window_scale, 0.0), "sigma": sigma, + "strength": v_strength, + "window_audio": max(base_window * a_window_scale, 0.0), + "sigma_audio": sigma_audio, + "strength_audio": a_strength, }) frame_cursor += L