Init

.gitignore

@@ -0,0 +1 @@
+*.pyc

README.md

@@ -0,0 +1,9 @@
+# ComfyUI Prompt-Relay
+
+WORK IN PROGRESS
+
+
+
+Original project:
+
+https://gordonchen19.github.io/Prompt-Relay/

__init__.py

@@ -0,0 +1,3 @@
+from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
+
+__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]

nodes.py

@@ -0,0 +1,107 @@
+import logging
+
+from comfy_api.latest import io
+
+from .prompt_relay import (
+    get_raw_tokenizer,
+    map_token_indices,
+    build_segments,
+    create_mask_fn,
+    distribute_segment_lengths,
+)
+
+from .patches import detect_model_type, apply_patches
+
+log = logging.getLogger(__name__)
+
+
+class PromptRelayEncode(io.ComfyNode):
+    """Encodes temporal local prompts and patches the model for Prompt Relay."""
+
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="PromptRelayEncode",
+            display_name="Prompt Relay Encode",
+            category="conditioning/prompt_relay",
+            description=(
+                "Encodes a global prompt combined with temporal local prompts and patches the model "
+                "for Prompt Relay temporal control. Local prompts are separated by |. "
+                "Use a standard CLIPTextEncode for the negative prompt."
+            ),
+            inputs=[
+                io.Model.Input("model"),
+                io.Clip.Input("clip"),
+                io.Latent.Input("latent", tooltip="Empty latent video — dimensions are read from its shape."),
+                io.String.Input(
+                    "global_prompt", multiline=True, default="",
+                    tooltip="Conditions the entire video. Anchors persistent characters, objects, and scene context.",
+                ),
+                io.String.Input(
+                    "local_prompts", multiline=True, default="",
+                    tooltip="Ordered prompts for each temporal segment, separated by |",
+                ),
+                io.String.Input(
+                    "segment_lengths", default="",
+                    tooltip="Comma-separated pixel space frame counts per segment. Leave empty to auto-distribute evenly.",
+                ),
+                io.Float.Input(
+                    "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.",
+                ),
+            ],
+            outputs=[
+                io.Model.Output(display_name="model"),
+                io.Conditioning.Output(display_name="positive"),
+            ],
+        )
+
+    @classmethod
+    def execute(cls, model, clip, latent, global_prompt, local_prompts, segment_lengths, epsilon) -> io.NodeOutput:
+        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 |)")
+
+        arch, patch_size, temporal_stride = detect_model_type(model)
+
+        parsed_lengths = None
+        if segment_lengths.strip():
+            pixel_lengths = [int(x.strip()) for x in segment_lengths.split(",") if x.strip()]
+            parsed_lengths = [max(1, round(p / temporal_stride)) for p in pixel_lengths]
+
+        raw_tokenizer = get_raw_tokenizer(clip)
+        full_prompt, token_ranges = map_token_indices(raw_tokenizer, global_prompt, locals_list)
+
+        log.info("[PromptRelay] Global: tokens [0:%d] (%d tokens)", token_ranges[0][0], token_ranges[0][0])
+        for i, (s, e) in enumerate(token_ranges):
+            log.info("[PromptRelay] Segment %d: tokens [%d:%d] (%d tokens)", i, s, e, e - s)
+
+        conditioning = clip.encode_from_tokens_scheduled(clip.tokenize(full_prompt))
+
+        samples = latent["samples"]
+        latent_frames = samples.shape[2]
+        tokens_per_frame = (samples.shape[3] // patch_size[1]) * (samples.shape[4] // patch_size[2])
+
+        effective_lengths = distribute_segment_lengths(len(locals_list), latent_frames, parsed_lengths)
+
+        log.info(
+            "[PromptRelay] Latent: %d frames, %d tokens/frame, segments: %s",
+            latent_frames, tokens_per_frame, effective_lengths,
+        )
+
+        q_token_idx = build_segments(token_ranges, effective_lengths, epsilon)
+        mask_fn = create_mask_fn(q_token_idx, tokens_per_frame, latent_frames)
+
+        patched = model.clone()
+        apply_patches(patched, arch, mask_fn)
+
+        return io.NodeOutput(patched, conditioning)
+
+
+NODE_CLASS_MAPPINGS = {
+    "PromptRelayEncode": PromptRelayEncode,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "PromptRelayEncode": "Prompt Relay Encode",
+}

patches.py

@@ -0,0 +1,167 @@
+import types
+import torch
+import comfy.ldm.modules.attention
+
+
+def _masked_attention(q, k, v, heads, mask, transformer_options={}, **kwargs):
+    # Bypass wrap_attn (sage/etc may ignore masks) by calling attention_pytorch directly.
+    return comfy.ldm.modules.attention.attention_pytorch(
+        q, k, v, heads, mask=mask,
+        _inside_attn_wrapper=True,
+        transformer_options=transformer_options,
+        **kwargs,
+    )
+
+
+def _wan_t2v_forward(self, mask_fn, x, context, transformer_options={}, **kwargs):
+    q = self.norm_q(self.q(x))
+    k = self.norm_k(self.k(context))
+    v = self.v(context)
+
+    mask = mask_fn(q, k, transformer_options)
+    if mask is not None:
+        x = _masked_attention(q, k, v, heads=self.num_heads, mask=mask,
+                              transformer_options=transformer_options)
+    else:
+        x = comfy.ldm.modules.attention.optimized_attention(
+            q, k, v, heads=self.num_heads, transformer_options=transformer_options,
+        )
+    return self.o(x)
+
+
+def _wan_i2v_forward(self, mask_fn, x, context, context_img_len, transformer_options={}, **kwargs):
+    context_img = context[:, :context_img_len]
+    context_text = context[:, context_img_len:]
+
+    q = self.norm_q(self.q(x))
+
+    k_img = self.norm_k_img(self.k_img(context_img))
+    v_img = self.v_img(context_img)
+    img_x = comfy.ldm.modules.attention.optimized_attention(
+        q, k_img, v_img, heads=self.num_heads, transformer_options=transformer_options,
+    )
+
+    k = self.norm_k(self.k(context_text))
+    v = self.v(context_text)
+
+    mask = mask_fn(q, k, transformer_options)
+    if mask is not None:
+        x = _masked_attention(q, k, v, heads=self.num_heads, mask=mask,
+                              transformer_options=transformer_options)
+    else:
+        x = comfy.ldm.modules.attention.optimized_attention(
+            q, k, v, heads=self.num_heads, transformer_options=transformer_options,
+        )
+
+    return self.o(x + img_x)
+
+
+def _ltx_forward(self, mask_fn, x, context=None, mask=None, pe=None, k_pe=None, transformer_options={}):
+    from comfy.ldm.lightricks.model import apply_rotary_emb
+
+    is_self_attn = context is None
+    context = x if is_self_attn else context
+
+    q = self.q_norm(self.to_q(x))
+    k = self.k_norm(self.to_k(context))
+    v = self.to_v(context)
+
+    if pe is not None:
+        q = apply_rotary_emb(q, pe)
+        k = apply_rotary_emb(k, pe if k_pe is None else k_pe)
+
+    if not is_self_attn:
+        temporal_mask = mask_fn(q, k, transformer_options)
+        if temporal_mask is not None:
+            mask = temporal_mask if mask is None else mask + temporal_mask
+
+    if mask is None:
+        out = comfy.ldm.modules.attention.optimized_attention(
+            q, k, v, self.heads, attn_precision=self.attn_precision,
+            transformer_options=transformer_options,
+        )
+    else:
+        out = _masked_attention(q, k, v, self.heads, mask=mask,
+                                attn_precision=self.attn_precision,
+                                transformer_options=transformer_options)
+
+    if self.to_gate_logits is not None:
+        gate_logits = self.to_gate_logits(x)
+        b, t, _ = out.shape
+        out = out.view(b, t, self.heads, self.dim_head)
+        out = out * (2.0 * torch.sigmoid(gate_logits)).unsqueeze(-1)
+        out = out.view(b, t, self.heads * self.dim_head)
+
+    return self.to_out(out)
+
+
+class _CrossAttnPatch:
+    """Descriptor that binds (impl, mask_fn) as a method onto a cross-attn module."""
+
+    def __init__(self, impl, mask_fn):
+        self.impl = impl
+        self.mask_fn = mask_fn
+
+    def __get__(self, obj, objtype=None):
+        impl, mask_fn = self.impl, self.mask_fn
+
+        def wrapped(self_module, *args, **kwargs):
+            return impl(self_module, mask_fn, *args, **kwargs)
+
+        return types.MethodType(wrapped, obj)
+
+
+def detect_model_type(model):
+    """Return (arch, patch_size, temporal_stride) for latent geometry.
+
+    temporal_stride is the VAE's pixel→latent temporal compression factor,
+    used to convert user-facing pixel frame counts to latent frames.
+    """
+    diff_model = model.model.diffusion_model
+
+    if hasattr(diff_model, "patch_size") and not hasattr(diff_model, "patchifier"):
+        return "wan", tuple(diff_model.patch_size), 4
+
+    if hasattr(diff_model, "patchifier"):
+        return "ltx", (1, 1, 1), int(diff_model.vae_scale_factors[0])
+
+    raise ValueError(
+        f"Unsupported model type: {type(diff_model).__name__}. "
+        f"Currently supports Wan and LTX models."
+    )
+
+
+def _check_unpatched(model_clone, key):
+    if key in getattr(model_clone, "object_patches", {}):
+        raise RuntimeError(
+            f"PromptRelay: cross-attention forward at '{key}' is already patched by "
+            "another node (e.g. KJNodes NAG). Stacking is not supported — remove the "
+            "conflicting node."
+        )
+
+
+def apply_patches(model_clone, arch, mask_fn):
+    diffusion_model = model_clone.get_model_object("diffusion_model")
+
+    if arch == "wan":
+        from comfy.ldm.wan.model import WanI2VCrossAttention
+        for idx, block in enumerate(diffusion_model.blocks):
+            key = f"diffusion_model.blocks.{idx}.cross_attn.forward"
+            _check_unpatched(model_clone, key)
+            cross_attn = block.cross_attn
+            impl = _wan_i2v_forward if isinstance(cross_attn, WanI2VCrossAttention) else _wan_t2v_forward
+            model_clone.add_object_patch(key, _CrossAttnPatch(impl, mask_fn).__get__(cross_attn, cross_attn.__class__))
+        return
+
+    if arch == "ltx":
+        for idx, block in enumerate(diffusion_model.transformer_blocks):
+            for attr in ("attn2", "audio_attn2"):
+                module = getattr(block, attr, None)
+                if module is None:
+                    continue
+                key = f"diffusion_model.transformer_blocks.{idx}.{attr}.forward"
+                _check_unpatched(model_clone, key)
+                model_clone.add_object_patch(key, _CrossAttnPatch(_ltx_forward, mask_fn).__get__(module, module.__class__))
+        return
+
+    raise ValueError(f"Unknown model arch: {arch}")

prompt_relay.py

@@ -0,0 +1,165 @@
+import logging
+import math
+import torch
+
+log = logging.getLogger(__name__)
+
+
+def build_temporal_cost(q_token_idx, Lq, Lk, device, dtype, tokens_per_frame):
+    """Gaussian penalty matrix [Lq, Lk] for video cross-attention (integer frame indexing)."""
+    offset = torch.zeros(Lq, Lk, device=device, dtype=dtype)
+    query_frames = torch.arange(Lq, device=device, dtype=torch.long) // 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)
+        offset[:, local] = cost.to(offset.dtype)
+
+    return offset
+
+
+def build_temporal_cost_scaled(q_token_idx, Lq, Lk, device, dtype, latent_frames):
+    """Penalty matrix for queries that don't map to integer frames (e.g. LTXAV audio tokens)."""
+    offset = torch.zeros(Lq, Lk, device=device, dtype=dtype)
+    query_frames = torch.arange(Lq, device=device, dtype=torch.float32) * latent_frames / Lq
+
+    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)
+        offset[:, local] = cost.to(offset.dtype)
+
+    return offset
+
+
+def create_mask_fn(q_token_idx, fallback_tokens_per_frame, latent_frames):
+    """Closure: mask_fn(q, k, transformer_options) -> additive mask or None."""
+    cache = {}
+    max_token_idx = max(int(seg["local_token_idx"].max().item()) for seg in q_token_idx) + 1
+
+    def mask_fn(q, k, transformer_options):
+        Lq, Lk = q.shape[1], k.shape[1]
+
+        if Lq == Lk:
+            return None
+
+        # Only apply on conditional pass — not unconditional (negative prompt)
+        cond_or_uncond = transformer_options.get("cond_or_uncond", [])
+        if 1 in cond_or_uncond and 0 not in cond_or_uncond:
+            return None
+
+        grid_sizes = transformer_options.get("grid_sizes", None)
+        video_tpf = int(grid_sizes[1]) * int(grid_sizes[2]) if grid_sizes is not None else fallback_tokens_per_frame
+        video_lq = latent_frames * video_tpf
+
+        # Skip cross-modal attention — text keys are padded to a fixed length ≥ max_token_idx and != video_lq
+        if Lk == video_lq or Lk < max_token_idx:
+            return None
+
+        mode = "video" if Lq == video_lq else "scaled"
+
+        key = (Lq, Lk, mode, q.device)
+        if key not in cache:
+            if mode == "video":
+                cost = build_temporal_cost(q_token_idx, Lq, Lk, q.device, q.dtype, video_tpf)
+            else:
+                cost = build_temporal_cost_scaled(q_token_idx, Lq, Lk, q.device, q.dtype, latent_frames)
+            log.info(
+                "[PromptRelay] Built penalty matrix (%s): Lq=%d, Lk=%d, nonzero=%d/%d",
+                mode, Lq, Lk, (cost > 0).sum().item(), cost.numel(),
+            )
+            cache[key] = -cost
+
+        return cache[key].to(q.dtype)
+
+    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."""
+    # 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
+
+    q_token_idx = []
+    frame_cursor = 0
+
+    for (tok_start, tok_end), L in zip(token_ranges, segment_lengths):
+        if L <= 0:
+            frame_cursor += L
+            continue
+        midpoint = (2 * frame_cursor + L) // 2
+        window = max(L // 2 - 2, 0)
+        q_token_idx.append({
+            "local_token_idx": torch.arange(tok_start, tok_end),
+            "midpoint": midpoint,
+            "window": window,
+            "sigma": sigma,
+        })
+        frame_cursor += L
+
+    return q_token_idx
+
+
+def get_raw_tokenizer(clip):
+    """Extract the raw SPiece/HF tokenizer from a ComfyUI CLIP object."""
+    tokenizer_wrapper = clip.tokenizer
+    for attr_name in dir(tokenizer_wrapper):
+        if attr_name.startswith("_"):
+            continue
+        inner = getattr(tokenizer_wrapper, attr_name, None)
+        if inner is not None and hasattr(inner, "tokenizer"):
+            return inner.tokenizer
+
+    raise RuntimeError(
+        f"Could not find raw tokenizer on CLIP object. "
+        f"Known attributes: {[a for a in dir(tokenizer_wrapper) if not a.startswith('_')]}"
+    )
+
+
+def map_token_indices(raw_tokenizer, global_prompt, local_prompts):
+    """Tokenize global + space-prefixed locals; return (full_prompt, per-local token ranges).
+
+    Uses incremental tokenization to avoid SentencePiece context-dependency issues.
+    """
+    prefixed_locals = [" " + lp for lp in local_prompts]
+    full_prompt = global_prompt + "".join(prefixed_locals)
+    has_eos = getattr(raw_tokenizer, "add_eos", False)
+    eos_adj = 1 if has_eos else 0
+
+    prev_len = len(raw_tokenizer(global_prompt)["input_ids"]) - eos_adj
+    token_ranges = []
+    built = global_prompt
+
+    for plp in prefixed_locals:
+        built += plp
+        cur_len = len(raw_tokenizer(built)["input_ids"]) - eos_adj
+        if cur_len <= prev_len:
+            raise ValueError(f"Local prompt produced no tokens: '{plp.strip()}'")
+        token_ranges.append((prev_len, cur_len))
+        prev_len = cur_len
+
+    return full_prompt, token_ranges
+
+
+def distribute_segment_lengths(num_segments, latent_frames, specified_lengths=None):
+    """Validate or auto-distribute segment frame counts, capped to fit within latent_frames."""
+    if specified_lengths:
+        if len(specified_lengths) != num_segments:
+            raise ValueError(
+                f"Number of segment_lengths ({len(specified_lengths)}) "
+                f"must match number of local prompts ({num_segments})"
+            )
+        lengths = specified_lengths
+    else:
+        # ceil division — matches reference implementation
+        step = -(-latent_frames // num_segments)
+        lengths = [step] * num_segments
+
+    effective = []
+    cursor = 0
+    for L in lengths:
+        end = min(cursor + L, latent_frames)
+        effective.append(max(end - cursor, 0))
+        cursor = end
+    return effective