Add some more knobs to play with

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)

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",
 }

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