From ab7d14bd89a6fab7be0b562425c5ec5a1a1983cd Mon Sep 17 00:00:00 2001 From: OpenClaw Agent Date: Thu, 9 Jul 2026 15:51:58 +0000 Subject: [PATCH] perf: streamline prompt relay and guide ui --- js/ltx_director_guide.js | 54 ++++---- prompt_relay.py | 270 ++++++++++++++++++++------------------- 2 files changed, 172 insertions(+), 152 deletions(-) diff --git a/js/ltx_director_guide.js b/js/ltx_director_guide.js index e46e335..717d245 100644 --- a/js/ltx_director_guide.js +++ b/js/ltx_director_guide.js @@ -1,24 +1,30 @@ -import { app } from "../../scripts/app.js"; - -// LTX Director Guide is a pure pass-through processor node. -// All configuration (images, insert frames, strengths) comes from -// the guide_data output of Prompt Relay Encode (Timeline). -app.registerExtension({ - name: "Comfy.LTXDirectorGuideCS", - async nodeCreated(node) { - if (node.comfyClass !== "LTXDirectorGuideCS") return; - - // Hide retake_mode widget on LiteGraph as it is dynamically auto-detected from the timeline data. - const w = node.widgets?.find(x => x.name === "retake_mode"); - if (w) { - w.hidden = true; - if (!w.options) w.options = {}; - w.options.hidden = true; - if (!window.LiteGraph || !window.LiteGraph.vueNodesMode) { - w.computeSize = () => [0, -4]; - w.draw = () => { }; - } - if (w.element) w.element.style.display = "none"; - } - }, -}); \ No newline at end of file +import { app } from "../../scripts/app.js"; + +// LTX Director Guide is a pure pass-through processor node. +// All configuration (images, insert frames, strengths) comes from +// the guide_data output of Prompt Relay Encode (Timeline). +function hideWidget(node, widgetName) { + const widget = node.widgets?.find((entry) => entry.name === widgetName); + if (!widget) return; + + widget.hidden = true; + widget.options = widget.options || {}; + widget.options.hidden = true; + + if (!window.LiteGraph || !window.LiteGraph.vueNodesMode) { + widget.computeSize = () => [0, -4]; + widget.draw = () => {}; + } + + if (widget.element) widget.element.style.display = "none"; +} + +app.registerExtension({ + name: "Comfy.LTXDirectorGuideCS", + async nodeCreated(node) { + if (node.comfyClass !== "LTXDirectorGuideCS") return; + + // Hide retake_mode widget on LiteGraph as it is dynamically auto-detected from the timeline data. + hideWidget(node, "retake_mode"); + }, +}); diff --git a/prompt_relay.py b/prompt_relay.py index ee575fe..15c19f6 100644 --- a/prompt_relay.py +++ b/prompt_relay.py @@ -1,71 +1,87 @@ -import logging -import math -import torch - -log = logging.getLogger(__name__) +import logging +import math +import os +import torch + +log = logging.getLogger(__name__) +_DEBUG_PROMPT_RELAY = os.environ.get("PROMPT_RELAY_DEBUG", "").strip().lower() in {"1", "true", "yes", "on"} +_DEBUG_LOG_PATH = os.path.join(os.path.dirname(__file__), "debug_prompt_relay.log") + + +def _build_segment_cost(query_frames, seg, dtype): + local = seg["local_token_idx"] + distance = (query_frames[:, None] - seg["midpoint"]).abs() + cost = seg["strength"] * (torch.relu(distance - seg["window"]) ** 2) / (2 * seg["sigma"] ** 2) + return local, cost.to(dtype) + + +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.float32) / float(tokens_per_frame)).floor() + + for seg in q_token_idx: + local, cost = _build_segment_cost(query_frames, seg, offset.dtype) + offset[:, local.to(device=device)] = cost + + return offset + + +def build_temporal_cost_scaled(q_token_idx, Lq, Lk, device, dtype, latent_frames, is_audio=False): + """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: + if is_audio: + active_seg = { + "local_token_idx": seg["local_token_idx"], + "midpoint": seg["midpoint"], + "window": seg.get("window_audio", seg["window"]), + "sigma": seg.get("sigma_audio", seg["sigma"]), + "strength": seg.get("strength_audio", 1.0), + } + else: + active_seg = seg + local, cost = _build_segment_cost(query_frames, active_seg, offset.dtype) + offset[:, local.to(device=device)] = cost + + return offset + + +def debug_log(msg): + if not _DEBUG_PROMPT_RELAY: + return + try: + with open(_DEBUG_LOG_PATH, "a", encoding="utf-8") as f: + f.write(msg + "\n") + except Exception: + pass -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() - 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 - - -def build_temporal_cost_scaled(q_token_idx, Lq, Lk, device, dtype, latent_frames, is_audio=False): - """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() - if is_audio: - sigma_val = seg.get("sigma_audio", seg["sigma"]) - window_val = seg.get("window_audio", seg["window"]) - strength_val = seg.get("strength_audio", 1.0) - else: - sigma_val = seg["sigma"] - window_val = seg["window"] - strength_val = seg.get("strength", 1.0) - cost = strength_val * (torch.relu(d - window_val) ** 2) / (2 * sigma_val ** 2) - offset[:, local] = cost.to(offset.dtype) - - return offset - - -def debug_log(msg): - try: - import os - log_path = os.path.join(os.path.dirname(__file__), "debug_prompt_relay.log") - with open(log_path, "a", encoding="utf-8") as f: - f.write(msg + "\n") - except Exception: - pass - - -def create_mask_fn(q_token_idx, fallback_tokens_per_frame, latent_frames): +def create_mask_fn(q_token_idx, fallback_tokens_per_frame, latent_frames): """Closure: mask_fn(Lq, Lk, dtype, device, transformer_options) -> additive mask or None. Takes shapes/dtype/device instead of tensors so callers can compute the mask - without first materializing q/k projections — required so PromptRelay can - wrap an existing cross-attn forward (e.g. KJNodes NAG) instead of replacing it. - """ - cache = {} - max_token_idx = max(int(seg["local_token_idx"].max().item()) for seg in q_token_idx) + 1 - - def mask_fn(Lq, Lk, dtype, device, transformer_options): - debug_log(f"mask_fn check: Lq={Lq} Lk={Lk} max_token_idx={max_token_idx} cond_or_uncond={transformer_options.get('cond_or_uncond', [])}") - if Lq == Lk: - debug_log("mask_fn: Lq == Lk, returning None") + without first materializing q/k projections — required so PromptRelay can + wrap an existing cross-attn forward (e.g. KJNodes NAG) instead of replacing it. + """ + cache = {} + max_token_idx = max(int(seg["local_token_idx"].max().item()) for seg in q_token_idx) + 1 + latent_frames = max(1, int(latent_frames)) + fallback_tokens_per_frame = max(1, int(fallback_tokens_per_frame)) + + def _get_video_query_layout(Lq, grid_sizes): + if grid_sizes is not None: + return int(grid_sizes[1]) * int(grid_sizes[2]) + if Lq % latent_frames == 0: + return Lq // latent_frames + return fallback_tokens_per_frame + + def mask_fn(Lq, Lk, dtype, device, transformer_options): + debug_log(f"mask_fn check: Lq={Lq} Lk={Lk} max_token_idx={max_token_idx} cond_or_uncond={transformer_options.get('cond_or_uncond', [])}") + if Lq == Lk: + debug_log("mask_fn: Lq == Lk, returning None") return None # Only apply on conditional pass — not unconditional (negative prompt) @@ -78,22 +94,16 @@ def create_mask_fn(q_token_idx, fallback_tokens_per_frame, latent_frames): attn_type = transformer_options.get("promptrelay_attn_type", "attn2") is_audio = (attn_type == "audio_attn2") - if is_audio: - mode = "scaled" - video_lq = -1 - else: - if grid_sizes is not None: - video_tpf = int(grid_sizes[1]) * int(grid_sizes[2]) - else: - if Lq % latent_frames == 0: - video_tpf = Lq // latent_frames - else: - video_tpf = 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: - debug_log(f"mask_fn: Lk == video_lq ({Lk == video_lq}) or Lk < max_token_idx ({Lk < max_token_idx}), returning None") + if is_audio: + mode = "scaled" + video_lq = -1 + else: + video_tpf = _get_video_query_layout(Lq, grid_sizes) + 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: + debug_log(f"mask_fn: Lk == video_lq ({Lk == video_lq}) or Lk < max_token_idx ({Lk < max_token_idx}), returning None") return None mode = "video" if Lq == video_lq else "scaled" @@ -113,7 +123,7 @@ def create_mask_fn(q_token_idx, fallback_tokens_per_frame, latent_frames): return cache[key].to(dtype) - return mask_fn + return mask_fn def build_segments(token_ranges, segment_lengths, epsilon=1e-3, relay_options=None): @@ -173,7 +183,7 @@ def build_segments(token_ranges, segment_lengths, epsilon=1e-3, relay_options=No return q_token_idx -def get_raw_tokenizer(clip): +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): @@ -183,56 +193,60 @@ def get_raw_tokenizer(clip): 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): + 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 _get_tokenizer_input_ids(raw_tokenizer, text): + tokenized = raw_tokenizer(text) + if isinstance(tokenized, dict) and "input_ids" in tokenized: + return tokenized["input_ids"] + if hasattr(tokenized, "input_ids"): + return tokenized.input_ids + if isinstance(tokenized, list): + return tokenized + return [] + + +def _tokenized_length(raw_tokenizer, text, eos_adjustment): + return max(0, len(_get_tokenizer_input_ids(raw_tokenizer, text)) - eos_adjustment) + + +def _tokenizer_has_eos(raw_tokenizer): + if getattr(raw_tokenizer, "add_eos", False): + return True + try: + ids = _get_tokenizer_input_ids(raw_tokenizer, "test") + if not ids: + return False + eos_id = getattr(raw_tokenizer, "eos_token_id", None) + return (eos_id is not None and ids[-1] == eos_id) or ids[-1] == 1 + except Exception: + return False + + +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) - - # Detect if the tokenizer appends EOS dynamically - has_eos = getattr(raw_tokenizer, "add_eos", False) - if not has_eos: - try: - test_res = raw_tokenizer("test") - if isinstance(test_res, dict) and "input_ids" in test_res: - ids = test_res["input_ids"] - elif hasattr(test_res, "input_ids"): - ids = test_res.input_ids - elif isinstance(test_res, list): - ids = test_res - else: - ids = [] - - if ids: - eos_id = getattr(raw_tokenizer, "eos_token_id", None) - if eos_id is not None and ids[-1] == eos_id: - has_eos = True - elif ids[-1] == 1: - has_eos = True - except Exception: - pass - - 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 + """ + prefixed_locals = [" " + lp for lp in local_prompts] + full_prompt = global_prompt + "".join(prefixed_locals) + + eos_adj = 1 if _tokenizer_has_eos(raw_tokenizer) else 0 + prev_len = _tokenized_length(raw_tokenizer, global_prompt, eos_adj) + token_ranges = [] + built = global_prompt + + for plp in prefixed_locals: + built += plp + cur_len = _tokenized_length(raw_tokenizer, built, 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