764 lines
29 KiB
Python
764 lines
29 KiB
Python
import logging
|
|
import math
|
|
import types
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
import comfy.utils
|
|
import folder_paths
|
|
import comfy.ldm.modules.attention
|
|
|
|
from comfy_api.latest import io
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
# ==============================================================================
|
|
# 1. CORE MATH & TOKENIZATION (Berasal dari prompt_relay.py)
|
|
# ==============================================================================
|
|
|
|
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):
|
|
"""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()
|
|
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
|
|
|
|
|
|
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 if q_token_idx else 0
|
|
|
|
def mask_fn(q, k, transformer_options):
|
|
Lq, Lk = q.shape[1], k.shape[1]
|
|
|
|
if Lq == Lk:
|
|
return None
|
|
|
|
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
|
|
|
|
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, relay_options=None):
|
|
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
|
|
|
|
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
|
|
base_window = max(L // 2 - 2, 0)
|
|
q_token_idx.append({
|
|
"local_token_idx": torch.arange(tok_start, tok_end),
|
|
"midpoint": midpoint,
|
|
"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
|
|
|
|
return q_token_idx
|
|
|
|
|
|
def get_raw_tokenizer(clip):
|
|
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("Could not find raw tokenizer on CLIP object.")
|
|
|
|
|
|
def map_token_indices(raw_tokenizer, global_prompt, local_prompts):
|
|
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):
|
|
if specified_lengths:
|
|
if len(specified_lengths) != num_segments:
|
|
raise ValueError("Number of segment_lengths must match number of local prompts")
|
|
lengths = specified_lengths
|
|
else:
|
|
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
|
|
|
|
|
|
# ==============================================================================
|
|
# 2. MODEL PATCHING (Berasal dari patches.py)
|
|
# ==============================================================================
|
|
|
|
def _masked_attention(q, k, v, heads, mask, transformer_options={}, **kwargs):
|
|
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:
|
|
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):
|
|
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__}.")
|
|
|
|
|
|
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. 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}")
|
|
|
|
|
|
# ==============================================================================
|
|
# 3. ENCODE RELAY LOGIC (Penggabungan & Setup)
|
|
# ==============================================================================
|
|
|
|
def _convert_to_latent_lengths(pixel_lengths, temporal_stride, latent_frames):
|
|
if not pixel_lengths:
|
|
return []
|
|
total_pixel = sum(pixel_lengths)
|
|
if total_pixel <= 0:
|
|
return [1] * len(pixel_lengths)
|
|
|
|
naive_total = max(1, round(total_pixel / temporal_stride))
|
|
target_total = min(latent_frames, naive_total)
|
|
|
|
if target_total >= latent_frames - 1:
|
|
target_total = latent_frames
|
|
|
|
exact = [p * target_total / total_pixel for p in pixel_lengths]
|
|
result = [int(e) for e in exact]
|
|
diff = target_total - sum(result)
|
|
if diff > 0:
|
|
order = sorted(range(len(exact)), key=lambda i: -(exact[i] - int(exact[i])))
|
|
for k in range(diff):
|
|
result[order[k % len(order)]] += 1
|
|
|
|
for i in range(len(result)):
|
|
if result[i] < 1:
|
|
max_idx = max(range(len(result)), key=lambda j: result[j])
|
|
if result[max_idx] > 1:
|
|
result[max_idx] -= 1
|
|
result[i] = 1
|
|
|
|
return result
|
|
|
|
|
|
def _encode_relay(model, clip, latent, global_prompt, local_prompts, segment_lengths, epsilon):
|
|
for name, val in (("global_prompt", global_prompt),
|
|
("local_prompts", local_prompts),
|
|
("segment_lengths", segment_lengths)):
|
|
if val is None:
|
|
raise ValueError(f"PromptRelay: '{name}' arrived as 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 |)")
|
|
|
|
arch, patch_size, temporal_stride = detect_model_type(model)
|
|
|
|
samples = latent["samples"]
|
|
latent_frames = samples.shape[2]
|
|
tokens_per_frame = (samples.shape[3] // patch_size[1]) * (samples.shape[4] // patch_size[2])
|
|
|
|
parsed_lengths = None
|
|
if segment_lengths.strip():
|
|
pixel_lengths = [int(x.strip()) for x in segment_lengths.split(",") if x.strip()]
|
|
parsed_lengths = _convert_to_latent_lengths(pixel_lengths, temporal_stride, latent_frames)
|
|
|
|
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))
|
|
|
|
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, None)
|
|
mask_fn = create_mask_fn(q_token_idx, tokens_per_frame, latent_frames)
|
|
|
|
patched = model.clone()
|
|
apply_patches(patched, arch, mask_fn)
|
|
|
|
patched.model_options["pr_segments"] = q_token_idx
|
|
patched.model_options["pr_latent_frames"] = latent_frames
|
|
patched.model_options["pr_tokens_per_frame"] = tokens_per_frame
|
|
|
|
return patched, conditioning
|
|
|
|
|
|
# ==============================================================================
|
|
# 4. TEMPORAL LORA LOGIC
|
|
# ==============================================================================
|
|
|
|
class _TemporalGatedLoRA(nn.Module):
|
|
def __init__(self, original, lora_down, lora_up, scale, strength, gate_weights, tokens_per_frame):
|
|
super().__init__()
|
|
self.original = original
|
|
self._lora_down = lora_down.cpu()
|
|
self._lora_up = lora_up.cpu()
|
|
self.scale = scale
|
|
self.strength = strength
|
|
self._gate_weights = gate_weights.cpu()
|
|
self.tokens_per_frame = tokens_per_frame
|
|
|
|
def forward(self, x, *args, **kwargs):
|
|
base = self.original(x, *args, **kwargs)
|
|
|
|
lora_down = self._lora_down.to(device=x.device, dtype=x.dtype)
|
|
lora_up = self._lora_up.to(device=x.device, dtype=x.dtype)
|
|
|
|
delta = (x @ lora_down.T) @ lora_up.T
|
|
delta = delta * (self.scale * self.strength)
|
|
|
|
if x.dim() < 2:
|
|
return base + delta
|
|
|
|
seq_len = x.shape[-2]
|
|
n_frames = len(self._gate_weights)
|
|
n_video_tokens = n_frames * self.tokens_per_frame
|
|
|
|
if seq_len >= n_video_tokens and n_video_tokens > 0:
|
|
gate_weights = self._gate_weights.to(device=x.device, dtype=x.dtype)
|
|
|
|
frame_idx = torch.arange(n_video_tokens, device=x.device) // self.tokens_per_frame
|
|
frame_idx = frame_idx.clamp(0, n_frames - 1)
|
|
video_gates = gate_weights[frame_idx]
|
|
|
|
if seq_len > n_video_tokens:
|
|
extra = torch.zeros(seq_len - n_video_tokens, device=x.device, dtype=x.dtype)
|
|
token_gates = torch.cat([video_gates, extra])
|
|
else:
|
|
token_gates = video_gates
|
|
|
|
view_shape = (1,) * (x.dim() - 2) + (seq_len, 1)
|
|
delta = delta * token_gates.view(view_shape)
|
|
|
|
return base + delta
|
|
|
|
|
|
def _build_gate_weights(seg, latent_frames):
|
|
midpoint = float(seg["midpoint"])
|
|
window = float(seg["window"])
|
|
sigma = float(seg["sigma"])
|
|
|
|
weights = []
|
|
for frame in range(latent_frames):
|
|
d = abs(frame - midpoint)
|
|
cost = (max(d - window, 0.0) ** 2) / (2.0 * sigma ** 2)
|
|
weights.append(math.exp(-cost))
|
|
|
|
return torch.tensor(weights, dtype=torch.float32)
|
|
|
|
|
|
def _get_module(diffusion_model, local_path):
|
|
m = diffusion_model
|
|
for part in local_path.split("."):
|
|
m = getattr(m, part)
|
|
return m
|
|
|
|
|
|
def _apply_gated_lora(model_clone, lora_file, segment, strength, tokens_per_frame, latent_frames):
|
|
try:
|
|
import comfy.loras
|
|
except ImportError:
|
|
log.error("[PromptRelay] comfy.loras not available in this ComfyUI build — cannot apply LoRA gate.")
|
|
return
|
|
|
|
key_map = {}
|
|
try:
|
|
comfy.loras.model_lora_keys_unet(model_clone.model, key_map)
|
|
except Exception as e:
|
|
log.error("[PromptRelay] Failed to build LoRA key map: %s", e)
|
|
return
|
|
|
|
try:
|
|
patches = comfy.loras.load_lora(lora_file, key_map)
|
|
except Exception as e:
|
|
log.error("[PromptRelay] Failed to parse LoRA file: %s", e)
|
|
return
|
|
|
|
gate_weights = _build_gate_weights(segment, latent_frames)
|
|
diffusion_model = model_clone.get_model_object("diffusion_model")
|
|
|
|
applied = 0
|
|
skipped = 0
|
|
|
|
for param_key, patch_data in patches.items():
|
|
if not param_key.endswith(".weight"):
|
|
skipped += 1
|
|
continue
|
|
|
|
module_path = param_key[:-7]
|
|
|
|
if not module_path.startswith("diffusion_model."):
|
|
skipped += 1
|
|
continue
|
|
|
|
local_path = module_path[len("diffusion_model."):]
|
|
|
|
try:
|
|
alpha = patch_data[0]
|
|
lora_up = patch_data[1]
|
|
lora_down = patch_data[2]
|
|
except (TypeError, IndexError):
|
|
skipped += 1
|
|
continue
|
|
|
|
if lora_up is None or lora_down is None:
|
|
skipped += 1
|
|
continue
|
|
|
|
if lora_down.dim() != 2 or lora_up.dim() != 2:
|
|
skipped += 1
|
|
continue
|
|
|
|
rank = lora_down.shape[0]
|
|
scale = float(alpha) / rank if alpha is not None else 1.0
|
|
|
|
if module_path in model_clone.object_patches:
|
|
target = model_clone.object_patches[module_path]
|
|
else:
|
|
try:
|
|
target = _get_module(diffusion_model, local_path)
|
|
except AttributeError:
|
|
skipped += 1
|
|
continue
|
|
|
|
wrapped = _TemporalGatedLoRA(
|
|
original=target,
|
|
lora_down=lora_down,
|
|
lora_up=lora_up,
|
|
scale=scale,
|
|
strength=strength,
|
|
gate_weights=gate_weights,
|
|
tokens_per_frame=tokens_per_frame,
|
|
)
|
|
|
|
model_clone.add_object_patch(module_path, wrapped)
|
|
applied += 1
|
|
|
|
log.info(
|
|
"[PromptRelay] Segment midpoint=%.1f window=%.1f — applied %d LoRA patches, skipped %d",
|
|
float(segment["midpoint"]), float(segment["window"]), applied, skipped,
|
|
)
|
|
|
|
|
|
# ==============================================================================
|
|
# 5. COMFYUI NODE CLASSES
|
|
# ==============================================================================
|
|
|
|
class PromptRelayEncodeTimeline(io.ComfyNode):
|
|
"""WYSIWYG timeline variant — segments and lengths come from a visual editor in the node UI."""
|
|
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="PromptRelayEncodeTimeline",
|
|
display_name="Prompt Relay Encode (Timeline)",
|
|
category="conditioning/prompt_relay",
|
|
description=(
|
|
"Same as Prompt Relay Encode, but local prompts and segment lengths are edited "
|
|
"visually as draggable blocks on a timeline. The max_frames input only sets the "
|
|
"timeline scale (pixel space) — actual frame count is still read from the latent."
|
|
),
|
|
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.Int.Input(
|
|
"max_frames", default=129, min=1, max=10000, step=1,
|
|
tooltip="Total timeline length in pixel-space frames. Used by the editor for visual scale only.",
|
|
),
|
|
io.String.Input(
|
|
"timeline_data", default="",
|
|
tooltip="JSON state of the timeline editor (auto-managed; do not edit by hand).",
|
|
),
|
|
io.String.Input(
|
|
"local_prompts", multiline=True, default="",
|
|
tooltip="Auto-populated from the timeline editor.",
|
|
),
|
|
io.String.Input(
|
|
"segment_lengths", default="",
|
|
tooltip="Auto-populated from the timeline editor (pixel-space frame counts).",
|
|
),
|
|
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.",
|
|
),
|
|
io.Float.Input(
|
|
"fps", default=24.0, min=0.1, max=240.0, step=0.1, optional=True,
|
|
tooltip="Frames per second — only affects how time is displayed in the timeline editor when time_units is set to 'seconds'.",
|
|
),
|
|
io.Combo.Input(
|
|
"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.",
|
|
),
|
|
],
|
|
outputs=[
|
|
io.Model.Output(display_name="model"),
|
|
io.Conditioning.Output(display_name="positive"),
|
|
io.Float.Output(display_name="fps_out"),
|
|
io.Int.Output(display_name="max_frames_out"),
|
|
],
|
|
)
|
|
|
|
@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:
|
|
patched, conditioning = _encode_relay(
|
|
model, clip, latent, global_prompt, local_prompts, segment_lengths, epsilon,
|
|
)
|
|
return io.NodeOutput(patched, conditioning, float(fps), int(max_frames))
|
|
|
|
|
|
class PromptRelayLoraGate(io.ComfyNode):
|
|
"""Applies a LoRA to one temporal segment defined by a Prompt Relay node."""
|
|
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="PromptRelayLoraGate",
|
|
display_name="Prompt Relay LoRA Gate",
|
|
category="conditioning/prompt_relay",
|
|
description=(
|
|
"Applies a LoRA with Gaussian temporal gating to one Prompt Relay segment. "
|
|
"Chain multiple gates — one per segment. Bypass unused gates. "
|
|
"segment_index is zero-based and must match the segment order in the Prompt Relay node."
|
|
),
|
|
inputs=[
|
|
io.Model.Input("model"),
|
|
io.Clip.Input("clip", tooltip="CLIP model. (Note: Patched CLIP only affects text encoded downstream, like negative prompts).", optional=True),
|
|
io.Combo.Input(
|
|
"lora_name",
|
|
folder_paths.get_filename_list("loras"),
|
|
tooltip="LoRA to apply to this segment.",
|
|
),
|
|
io.Int.Input(
|
|
"segment_index", default=0, min=0, max=99, step=1,
|
|
tooltip="Zero-based index of the Prompt Relay segment this LoRA targets. Segment 0 is the first local prompt.",
|
|
),
|
|
io.Float.Input(
|
|
"strength", default=0.8, min=-10.0, max=10.0, step=0.01,
|
|
tooltip="LoRA strength for the UNet at the segment centre. Fades to zero outside the segment window.",
|
|
),
|
|
io.Float.Input(
|
|
"clip_strength", default=1.0, min=-10.0, max=10.0, step=0.01,
|
|
tooltip="LoRA strength for the CLIP model (applied globally to this CLIP stream).",
|
|
),
|
|
],
|
|
outputs=[
|
|
io.Model.Output(display_name="model"),
|
|
io.Clip.Output(display_name="clip"),
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, model, lora_name, segment_index, strength, clip_strength, clip=None) -> io.NodeOutput:
|
|
if strength == 0.0 and clip_strength == 0.0:
|
|
return io.NodeOutput(model, clip)
|
|
|
|
lora_path = folder_paths.get_full_path("loras", lora_name)
|
|
if lora_path is None:
|
|
raise ValueError(f"[PromptRelay LoraGate] LoRA file not found: {lora_name}")
|
|
|
|
import comfy.utils
|
|
import comfy.sd
|
|
|
|
# Load the file once for both Model and CLIP
|
|
lora_file = comfy.utils.load_torch_file(lora_path, safe_load=True)
|
|
|
|
# --- 1. Patch CLIP globally (Standard ComfyUI method) ---
|
|
patched_clip = clip
|
|
if clip is not None and clip_strength != 0.0:
|
|
_, patched_clip = comfy.sd.load_lora_for_models(None, clip, lora_file, 0.0, clip_strength)
|
|
|
|
# --- 2. Patch UNet temporally ---
|
|
patched_model = model
|
|
if strength != 0.0:
|
|
segments = model.model_options.get("pr_segments")
|
|
latent_frames = model.model_options.get("pr_latent_frames")
|
|
tokens_per_frame = model.model_options.get("pr_tokens_per_frame")
|
|
|
|
if segments is None:
|
|
log.warning(
|
|
"[PromptRelay LoraGate] No Prompt Relay segment metadata on this model. "
|
|
"Ensure the model comes from a Prompt Relay Encode node. Passing through model unchanged."
|
|
)
|
|
elif segment_index >= len(segments):
|
|
log.warning(
|
|
"[PromptRelay LoraGate] segment_index %d is out of range — only %d segment(s) defined. "
|
|
"Passing through model unchanged.",
|
|
segment_index, len(segments),
|
|
)
|
|
else:
|
|
patched_model = model.clone()
|
|
_apply_gated_lora(
|
|
patched_model,
|
|
lora_file,
|
|
segments[segment_index],
|
|
strength,
|
|
tokens_per_frame,
|
|
latent_frames,
|
|
)
|
|
|
|
# Forward metadata to the next node in the chain
|
|
patched_model.model_options["pr_segments"] = segments
|
|
patched_model.model_options["pr_latent_frames"] = latent_frames
|
|
patched_model.model_options["pr_tokens_per_frame"] = tokens_per_frame
|
|
|
|
return io.NodeOutput(patched_model, patched_clip)
|
|
|
|
|
|
# ==============================================================================
|
|
# 6. MAPPINGS
|
|
# ==============================================================================
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"PromptRelayEncodeTimeline": PromptRelayEncodeTimeline,
|
|
"PromptRelayLoraGate": PromptRelayLoraGate,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"PromptRelayEncodeTimeline": "Prompt Relay Encode (Timeline)",
|
|
"PromptRelayLoraGate": "Prompt Relay LoRA Gate",
|
|
} |