chris.dumas/WhatDreamsCost-Dumas
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

perf: streamline prompt relay and guide ui

Browse Source
This commit is contained in:
OpenClaw Agent
2026-07-09 15:51:58 +00:00
parent a536ac5a14
commit ab7d14bd89
2 changed files with 172 additions and 152 deletions
Download Patch File Download Diff File Expand all files Collapse all files

270
prompt_relay.py
View File

@@ -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