perf: streamline prompt relay and guide ui
This commit is contained in:
OpenClaw Agent
@@ -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";
|
||||
}
|
||||
},
|
||||
});
|
||||
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");
|
||||
},
|
||||
});
|
||||
|
||||
270
prompt_relay.py
270
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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user