v1.3.0 New Nodes: LTX Director and LTX Director Guide
This commit is contained in:
27
__init__.py
27
__init__.py
@@ -3,26 +3,43 @@ from .multi_image_loader import MultiImageLoader
|
|||||||
from .ltx_sequencer import LTXSequencer
|
from .ltx_sequencer import LTXSequencer
|
||||||
from .speech_length_calculator import SpeechLengthCalculator
|
from .speech_length_calculator import SpeechLengthCalculator
|
||||||
from .load_audio_ui import LoadAudioUI
|
from .load_audio_ui import LoadAudioUI
|
||||||
from .load_video_ui import LoadVideoUI # Imported new node
|
from .load_video_ui import LoadVideoUI
|
||||||
|
from .ltx_director import LTXDirector
|
||||||
|
from .ltx_director_guide import LTXDirectorGuide
|
||||||
|
from comfy_api.latest import ComfyExtension, io
|
||||||
|
from typing_extensions import override
|
||||||
|
|
||||||
|
class PromptRelay(ComfyExtension):
|
||||||
|
@override
|
||||||
|
async def get_node_list(self) -> list[type[io.ComfyNode]]:
|
||||||
|
return [
|
||||||
|
LTXDirector,
|
||||||
|
LTXDirectorGuide
|
||||||
|
]
|
||||||
|
|
||||||
|
async def comfy_entrypoint() -> PromptRelay:
|
||||||
|
return PromptRelay()
|
||||||
|
|
||||||
# Register the node classes
|
|
||||||
NODE_CLASS_MAPPINGS = {
|
NODE_CLASS_MAPPINGS = {
|
||||||
"LTXKeyframer": LTXKeyframer,
|
"LTXKeyframer": LTXKeyframer,
|
||||||
"MultiImageLoader": MultiImageLoader,
|
"MultiImageLoader": MultiImageLoader,
|
||||||
"LTXSequencer": LTXSequencer,
|
"LTXSequencer": LTXSequencer,
|
||||||
"SpeechLengthCalculator": SpeechLengthCalculator,
|
"SpeechLengthCalculator": SpeechLengthCalculator,
|
||||||
"LoadAudioUI": LoadAudioUI,
|
"LoadAudioUI": LoadAudioUI,
|
||||||
"LoadVideoUI": LoadVideoUI # Registered new node
|
"LoadVideoUI": LoadVideoUI,
|
||||||
|
"LTXDirector": LTXDirector,
|
||||||
|
"LTXDirectorGuide": LTXDirectorGuide,
|
||||||
}
|
}
|
||||||
|
|
||||||
# Provide clean display names for the ComfyUI interface
|
|
||||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||||
"LTXKeyframer": "LTX Keyframer",
|
"LTXKeyframer": "LTX Keyframer",
|
||||||
"MultiImageLoader": "Multi Image Loader",
|
"MultiImageLoader": "Multi Image Loader",
|
||||||
"LTXSequencer": "LTX Sequencer",
|
"LTXSequencer": "LTX Sequencer",
|
||||||
"SpeechLengthCalculator": "Speech Length Calculator",
|
"SpeechLengthCalculator": "Speech Length Calculator",
|
||||||
"LoadAudioUI": "Load Audio UI",
|
"LoadAudioUI": "Load Audio UI",
|
||||||
"LoadVideoUI": "Load Video UI" # Registered display name
|
"LoadVideoUI": "Load Video UI",
|
||||||
|
"LTXDirector": "LTX Director",
|
||||||
|
"LTXDirectorGuide": "LTX Director Guide",
|
||||||
}
|
}
|
||||||
|
|
||||||
WEB_DIRECTORY = "./js"
|
WEB_DIRECTORY = "./js"
|
||||||
|
|||||||
3638
example_workflows/LTX Director Example Workflow Compa Subgraph.json
Normal file
3638
example_workflows/LTX Director Example Workflow Compa Subgraph.json
Normal file
File diff suppressed because it is too large
Load Diff
Binary file not shown.
|
After Width: | Height: | Size: 285 KiB |
2337
example_workflows/LTX Director Example Workflow.json
Normal file
2337
example_workflows/LTX Director Example Workflow.json
Normal file
File diff suppressed because it is too large
Load Diff
BIN
example_workflows/LTX Director Example Workflow.png
Normal file
BIN
example_workflows/LTX Director Example Workflow.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 825 KiB |
3870
js/ltx_director.js
Normal file
3870
js/ltx_director.js
Normal file
File diff suppressed because it is too large
Load Diff
17
js/ltx_director_guide.js
Normal file
17
js/ltx_director_guide.js
Normal file
@@ -0,0 +1,17 @@
|
|||||||
|
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).
|
||||||
|
// No dynamic widgets or sync logic needed.
|
||||||
|
app.registerExtension({
|
||||||
|
name: "Comfy.LTXDirectorGuide",
|
||||||
|
async nodeCreated(node) {
|
||||||
|
if (node.comfyClass !== "LTXDirectorGuide") return;
|
||||||
|
// Nothing to initialize — the node has no configurable widgets.
|
||||||
|
},
|
||||||
|
});
|
||||||
|
|
||||||
|
// Nothing to initialize — the node has no configurable widgets.
|
||||||
|
},
|
||||||
|
});
|
||||||
643
ltx_director.py
Normal file
643
ltx_director.py
Normal file
@@ -0,0 +1,643 @@
|
|||||||
|
import logging
|
||||||
|
import json
|
||||||
|
import base64
|
||||||
|
import io as _io
|
||||||
|
import math
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
import torch
|
||||||
|
import av
|
||||||
|
from PIL import Image
|
||||||
|
|
||||||
|
import os
|
||||||
|
import folder_paths
|
||||||
|
import comfy.model_management
|
||||||
|
|
||||||
|
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__)
|
||||||
|
|
||||||
|
# Custom socket type shared with LTXSequencer
|
||||||
|
GuideData = io.Custom("GUIDE_DATA")
|
||||||
|
|
||||||
|
|
||||||
|
def _load_image_tensor(seg: dict) -> torch.Tensor:
|
||||||
|
"""Decode an image from the ComfyUI input folder (if imageFile provided) or fallback to base64
|
||||||
|
to a ComfyUI-style image tensor of shape [1, H, W, 3], float32 in [0, 1]."""
|
||||||
|
if seg.get("imageFile"):
|
||||||
|
file_path = os.path.join(folder_paths.get_input_directory(), seg["imageFile"])
|
||||||
|
if os.path.exists(file_path):
|
||||||
|
img = Image.open(file_path).convert("RGB")
|
||||||
|
arr = np.array(img, dtype=np.float32) / 255.0
|
||||||
|
return torch.from_numpy(arr).unsqueeze(0)
|
||||||
|
|
||||||
|
b64_str = seg.get("imageB64", "")
|
||||||
|
if not b64_str or b64_str.startswith("/view?"):
|
||||||
|
return torch.zeros((1, 512, 512, 3), dtype=torch.float32)
|
||||||
|
|
||||||
|
if "," in b64_str:
|
||||||
|
b64_str = b64_str.split(",", 1)[1]
|
||||||
|
|
||||||
|
try:
|
||||||
|
img_bytes = base64.b64decode(b64_str)
|
||||||
|
img = Image.open(_io.BytesIO(img_bytes)).convert("RGB")
|
||||||
|
arr = np.array(img, dtype=np.float32) / 255.0
|
||||||
|
return torch.from_numpy(arr).unsqueeze(0)
|
||||||
|
except:
|
||||||
|
return torch.zeros((1, 512, 512, 3), dtype=torch.float32)
|
||||||
|
|
||||||
|
|
||||||
|
def _resize_image(tensor: torch.Tensor, target_w: int, target_h: int, method: str, divisible_by: int) -> torch.Tensor:
|
||||||
|
"""Resize a [1, H, W, 3] float32 tensor to target dimensions using the given method,
|
||||||
|
then snap the final dimensions to be divisible by `divisible_by`."""
|
||||||
|
from PIL import Image as _PilImage
|
||||||
|
import torchvision.transforms.functional as TF
|
||||||
|
|
||||||
|
def snap(val, div):
|
||||||
|
return max(div, (val // div) * div)
|
||||||
|
|
||||||
|
tw = snap(target_w, divisible_by)
|
||||||
|
th = snap(target_h, divisible_by)
|
||||||
|
|
||||||
|
img_np = (tensor[0].cpu().numpy() * 255.0).clip(0, 255).astype(np.uint8)
|
||||||
|
pil = _PilImage.fromarray(img_np)
|
||||||
|
src_w, src_h = pil.size
|
||||||
|
|
||||||
|
if method == "stretch to fit":
|
||||||
|
resized = pil.resize((tw, th), _PilImage.LANCZOS)
|
||||||
|
|
||||||
|
elif method == "maintain aspect ratio":
|
||||||
|
ratio = min(tw / src_w, th / src_h)
|
||||||
|
new_w = int(src_w * ratio)
|
||||||
|
new_h = int(src_h * ratio)
|
||||||
|
new_w = snap(new_w, divisible_by)
|
||||||
|
new_h = snap(new_h, divisible_by)
|
||||||
|
resized = pil.resize((new_w, new_h), _PilImage.LANCZOS)
|
||||||
|
|
||||||
|
elif method == "pad":
|
||||||
|
ratio = min(tw / src_w, th / src_h)
|
||||||
|
new_w = snap(int(src_w * ratio), divisible_by)
|
||||||
|
new_h = snap(int(src_h * ratio), divisible_by)
|
||||||
|
inner = pil.resize((new_w, new_h), _PilImage.LANCZOS)
|
||||||
|
resized = _PilImage.new("RGB", (tw, th), (0, 0, 0))
|
||||||
|
resized.paste(inner, ((tw - new_w) // 2, (th - new_h) // 2))
|
||||||
|
|
||||||
|
elif method == "crop":
|
||||||
|
ratio = max(tw / src_w, th / src_h)
|
||||||
|
new_w = int(src_w * ratio)
|
||||||
|
new_h = int(src_h * ratio)
|
||||||
|
inner = pil.resize((new_w, new_h), _PilImage.LANCZOS)
|
||||||
|
left = (new_w - tw) // 2
|
||||||
|
top = (new_h - th) // 2
|
||||||
|
resized = inner.crop((left, top, left + tw, top + th))
|
||||||
|
|
||||||
|
else:
|
||||||
|
resized = pil.resize((tw, th), _PilImage.LANCZOS)
|
||||||
|
|
||||||
|
arr = np.array(resized, dtype=np.float32) / 255.0
|
||||||
|
return torch.from_numpy(arr).unsqueeze(0)
|
||||||
|
|
||||||
|
|
||||||
|
def _compress_image(tensor: torch.Tensor, crf: int) -> torch.Tensor:
|
||||||
|
"""Apply H.264 compression artefacts to a [1, H, W, 3] float32 tensor (ComfyUI image format).
|
||||||
|
crf=0 means no compression. Uses PyAV to encode/decode a single frame in-memory."""
|
||||||
|
if crf == 0:
|
||||||
|
return tensor
|
||||||
|
img = tensor[0] # [H, W, 3]
|
||||||
|
# Dimensions must be even for H.264
|
||||||
|
h = (img.shape[0] // 2) * 2
|
||||||
|
w = (img.shape[1] // 2) * 2
|
||||||
|
img_np = (img[:h, :w] * 255.0).byte().cpu().numpy() # uint8 [H, W, 3]
|
||||||
|
|
||||||
|
try:
|
||||||
|
buf = _io.BytesIO()
|
||||||
|
container = av.open(buf, mode="w", format="mp4")
|
||||||
|
stream = container.add_stream("libx264", rate=1)
|
||||||
|
stream.width = w
|
||||||
|
stream.height = h
|
||||||
|
stream.pix_fmt = "yuv420p"
|
||||||
|
stream.options = {"crf": str(crf), "preset": "ultrafast"}
|
||||||
|
frame = av.VideoFrame.from_ndarray(img_np, format="rgb24")
|
||||||
|
for pkt in stream.encode(frame):
|
||||||
|
container.mux(pkt)
|
||||||
|
for pkt in stream.encode(None):
|
||||||
|
container.mux(pkt)
|
||||||
|
container.close()
|
||||||
|
|
||||||
|
buf.seek(0)
|
||||||
|
container_r = av.open(buf, mode="r")
|
||||||
|
decoded = None
|
||||||
|
for frame_r in container_r.decode(video=0):
|
||||||
|
decoded = frame_r.to_ndarray(format="rgb24") # [H, W, 3]
|
||||||
|
break
|
||||||
|
container_r.close()
|
||||||
|
|
||||||
|
if decoded is None:
|
||||||
|
return tensor
|
||||||
|
arr = torch.from_numpy(decoded.astype(np.float32) / 255.0).to(tensor.device, tensor.dtype)
|
||||||
|
# Re-embed into original tensor shape (may have been cropped by even-rounding)
|
||||||
|
out = tensor.clone()
|
||||||
|
out[0, :h, :w] = arr
|
||||||
|
return out
|
||||||
|
except Exception as e:
|
||||||
|
log.warning("[PromptRelay] img_compression encode/decode failed: %s", e)
|
||||||
|
return tensor
|
||||||
|
|
||||||
|
|
||||||
|
def _build_combined_audio(timeline_data_str: str, duration_frames: int, frame_rate: float) -> dict:
|
||||||
|
"""Parses timeline JSON, loads/trims audio directly from memory using PyAV,
|
||||||
|
and aligns to a global timeline yielding ComfyUI's format.
|
||||||
|
Output length explicitly mimics the timeline's duration_frames length."""
|
||||||
|
target_sr = 44100
|
||||||
|
total_samples = max(1, int(math.ceil(duration_frames / frame_rate * target_sr)))
|
||||||
|
empty_audio = {"waveform": torch.zeros((1, 2, total_samples), dtype=torch.float32), "sample_rate": target_sr}
|
||||||
|
|
||||||
|
if not timeline_data_str:
|
||||||
|
return empty_audio
|
||||||
|
|
||||||
|
try:
|
||||||
|
data = json.loads(timeline_data_str)
|
||||||
|
audio_segs = data.get("audioSegments", [])
|
||||||
|
except Exception:
|
||||||
|
return empty_audio
|
||||||
|
|
||||||
|
if not audio_segs:
|
||||||
|
return empty_audio
|
||||||
|
|
||||||
|
out_waveform = torch.zeros((2, total_samples), dtype=torch.float32)
|
||||||
|
|
||||||
|
for seg in audio_segs:
|
||||||
|
buffer = None
|
||||||
|
if seg.get("audioFile"):
|
||||||
|
file_path = os.path.join(folder_paths.get_input_directory(), seg["audioFile"])
|
||||||
|
if os.path.exists(file_path):
|
||||||
|
with open(file_path, "rb") as f:
|
||||||
|
buffer = _io.BytesIO(f.read())
|
||||||
|
|
||||||
|
if not buffer and seg.get("audioB64"):
|
||||||
|
b64 = seg.get("audioB64")
|
||||||
|
if "," in b64:
|
||||||
|
b64 = b64.split(",", 1)[1]
|
||||||
|
try:
|
||||||
|
audio_bytes = base64.b64decode(b64)
|
||||||
|
buffer = _io.BytesIO(audio_bytes)
|
||||||
|
except:
|
||||||
|
pass
|
||||||
|
|
||||||
|
if not buffer:
|
||||||
|
continue
|
||||||
|
|
||||||
|
try:
|
||||||
|
clip_frames = []
|
||||||
|
|
||||||
|
# Use PyAV to decode directly from memory buffer
|
||||||
|
with av.open(buffer) as container:
|
||||||
|
stream = container.streams.audio[0]
|
||||||
|
|
||||||
|
# Setup resampler to ensure output is 44.1kHz, Stereo, Float32 Planar
|
||||||
|
resampler = av.AudioResampler(
|
||||||
|
format='fltp',
|
||||||
|
layout='stereo',
|
||||||
|
rate=target_sr,
|
||||||
|
)
|
||||||
|
|
||||||
|
for frame in container.decode(stream):
|
||||||
|
for resampled_frame in resampler.resample(frame):
|
||||||
|
# to_ndarray() on fltp gives shape (channels, samples)
|
||||||
|
arr = resampled_frame.to_ndarray()
|
||||||
|
clip_frames.append(torch.from_numpy(arr))
|
||||||
|
|
||||||
|
# Flush the resampler to get any remaining samples
|
||||||
|
for resampled_frame in resampler.resample(None):
|
||||||
|
arr = resampled_frame.to_ndarray()
|
||||||
|
clip_frames.append(torch.from_numpy(arr))
|
||||||
|
|
||||||
|
if not clip_frames:
|
||||||
|
continue
|
||||||
|
|
||||||
|
# Concatenate all frame blocks along the samples dimension (dim 1)
|
||||||
|
waveform = torch.cat(clip_frames, dim=1) # Shape: [2, total_clip_samples]
|
||||||
|
|
||||||
|
# Calculate interactive trim boundaries
|
||||||
|
trim_start_frames = float(seg.get("trimStart", 0))
|
||||||
|
length_frames = float(seg.get("length", 1))
|
||||||
|
start_frames = float(seg.get("start", 0))
|
||||||
|
|
||||||
|
start_sample_src = int(trim_start_frames / frame_rate * target_sr)
|
||||||
|
length_samples = int(length_frames / frame_rate * target_sr)
|
||||||
|
end_sample_src = start_sample_src + length_samples
|
||||||
|
|
||||||
|
if start_sample_src < 0: start_sample_src = 0
|
||||||
|
if end_sample_src > waveform.shape[1]:
|
||||||
|
end_sample_src = waveform.shape[1]
|
||||||
|
|
||||||
|
actual_length = end_sample_src - start_sample_src
|
||||||
|
if actual_length <= 0: continue
|
||||||
|
|
||||||
|
# Extract the correct segment of the audio
|
||||||
|
clip_waveform = waveform[:, start_sample_src:end_sample_src]
|
||||||
|
|
||||||
|
# Position onto the timeline
|
||||||
|
start_sample_dst = int(start_frames / frame_rate * target_sr)
|
||||||
|
|
||||||
|
if start_sample_dst >= out_waveform.shape[1]:
|
||||||
|
continue
|
||||||
|
|
||||||
|
end_sample_dst = start_sample_dst + actual_length
|
||||||
|
|
||||||
|
# Clip any trailing overflow so we don't index past the timeline bounds
|
||||||
|
if end_sample_dst > out_waveform.shape[1]:
|
||||||
|
actual_length = out_waveform.shape[1] - start_sample_dst
|
||||||
|
clip_waveform = clip_waveform[:, :actual_length]
|
||||||
|
end_sample_dst = start_sample_dst + actual_length
|
||||||
|
|
||||||
|
if actual_length <= 0:
|
||||||
|
continue
|
||||||
|
|
||||||
|
# Additive composite (allows clips overlapping to sum together naturally)
|
||||||
|
out_waveform[:, start_sample_dst:end_sample_dst] += clip_waveform
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
log.warning("[PromptRelay] Audio process error for segment %s: %s", seg.get("fileName"), e)
|
||||||
|
continue
|
||||||
|
|
||||||
|
return {"waveform": out_waveform.unsqueeze(0), "sample_rate": target_sr}
|
||||||
|
|
||||||
|
|
||||||
|
def _convert_to_latent_lengths(pixel_lengths, temporal_stride, latent_frames):
|
||||||
|
"""Convert pixel-space segment lengths to integer latent-space lengths using the
|
||||||
|
largest-remainder method. Targets the full `latent_frames` when the pixel sum looks
|
||||||
|
like full coverage (within one stride of latent_frames * stride). Otherwise targets
|
||||||
|
round(total_pixel / temporal_stride) so partial-coverage timelines stay partial.
|
||||||
|
"""
|
||||||
|
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)
|
||||||
|
# Within one frame of full → user clearly intended full coverage; pin to latent_frames.
|
||||||
|
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
|
||||||
|
|
||||||
|
# Ensure every segment has ≥ 1 latent frame (steal from the largest if needed).
|
||||||
|
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. "
|
||||||
|
"Likely causes: a stale workflow JSON saved with null, the timeline "
|
||||||
|
"editor's web extension failing to load, or an upstream node returning None. "
|
||||||
|
"Set the field to an empty string or fix the upstream connection."
|
||||||
|
)
|
||||||
|
|
||||||
|
# Split prompts but do NOT filter out empty ones yet, so we can detect them
|
||||||
|
locals_list = [p.strip() for p in local_prompts.split("|")]
|
||||||
|
|
||||||
|
# Check if any specific segment is empty
|
||||||
|
for p in locals_list:
|
||||||
|
if not p:
|
||||||
|
raise ValueError("There is a segment on the timeline missing a prompt!")
|
||||||
|
|
||||||
|
if not locals_list or (len(locals_list) == 1 and not locals_list[0]):
|
||||||
|
raise ValueError("At least one local prompt is required.")
|
||||||
|
|
||||||
|
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(float(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)
|
||||||
|
|
||||||
|
return patched, conditioning
|
||||||
|
|
||||||
|
|
||||||
|
class LTXDirector(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="LTXDirector",
|
||||||
|
display_name="LTX Director",
|
||||||
|
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 duration_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.Vae.Input("audio_vae", optional=True, tooltip="Optional. Connect an Audio VAE to generate audio latents."),
|
||||||
|
io.Latent.Input("optional_latent", optional=True, tooltip="Optional. Connect a latent to override the auto-generated one."),
|
||||||
|
io.String.Input(
|
||||||
|
"global_prompt", multiline=True, default="",
|
||||||
|
tooltip="Conditions the entire video. Anchors persistent characters, objects, and scene context.",
|
||||||
|
),
|
||||||
|
io.Int.Input(
|
||||||
|
"duration_frames", default=120, min=1, max=10000, step=1,
|
||||||
|
tooltip="Total timeline length in pixel-space frames. Used by the editor for visual scale only.",
|
||||||
|
),
|
||||||
|
io.Float.Input(
|
||||||
|
"duration_seconds", default=5, min=0.1, max=1000.0, step=0.01,
|
||||||
|
tooltip="Total timeline duration in seconds (computed/synced from frames).",
|
||||||
|
),
|
||||||
|
io.String.Input(
|
||||||
|
"timeline_data", default="",
|
||||||
|
tooltip="JSON state of the timeline editor (auto-managed; do not edit by hand).",
|
||||||
|
),
|
||||||
|
io.Boolean.Input(
|
||||||
|
"use_custom_audio", default=False, optional=True,
|
||||||
|
tooltip="Toggle between using timeline audio (ON) and generating audio from scratch (OFF).",
|
||||||
|
),
|
||||||
|
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=0.001, min=0.0001, max=0.99, step=0.0001,
|
||||||
|
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(
|
||||||
|
"frame_rate", default=24, min=1, max=240, step=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(
|
||||||
|
"display_mode", options=["frames", "seconds"], default="seconds", optional=True,
|
||||||
|
tooltip="Display the ruler, segment ranges, length input, and total in frames or seconds. Internal storage is always pixel-space frames.",
|
||||||
|
),
|
||||||
|
io.String.Input(
|
||||||
|
"guide_strength", default="",
|
||||||
|
tooltip="Auto-populated from the timeline editor (comma-separated guide strengths for image segments).",
|
||||||
|
),
|
||||||
|
io.Int.Input(
|
||||||
|
"custom_width", default=0, min=0, max=8192, step=1, optional=True,
|
||||||
|
tooltip="Target output width for all image segments. Set to 0 to use the original image width.",
|
||||||
|
),
|
||||||
|
io.Int.Input(
|
||||||
|
"custom_height", default=0, min=0, max=8192, step=1, optional=True,
|
||||||
|
tooltip="Target output height for all image segments. Set to 0 to use the original image height.",
|
||||||
|
),
|
||||||
|
io.Combo.Input(
|
||||||
|
"resize_method",
|
||||||
|
options=["maintain aspect ratio", "stretch to fit", "pad", "crop"],
|
||||||
|
default="maintain aspect ratio",
|
||||||
|
optional=True,
|
||||||
|
tooltip="How to resize image segments to fit the target dimensions.",
|
||||||
|
),
|
||||||
|
io.Int.Input(
|
||||||
|
"divisible_by", default=32, min=1, max=256, step=1, optional=True,
|
||||||
|
tooltip="Snap the final output image dimensions to be divisible by this number (e.g. 32 for LTX).",
|
||||||
|
),
|
||||||
|
io.Int.Input(
|
||||||
|
"img_compression", default=18, min=0, max=100, step=1, optional=True,
|
||||||
|
tooltip="H.264 CRF compression to apply to each guide image. 0 = no compression, higher = more artefacts.",
|
||||||
|
),
|
||||||
|
],
|
||||||
|
outputs=[
|
||||||
|
io.Model.Output(display_name="model"),
|
||||||
|
io.Conditioning.Output(display_name="positive"),
|
||||||
|
io.Latent.Output(display_name="video_latent", tooltip="Auto-generated LTXV empty latent (only populated when no latent is connected)."),
|
||||||
|
io.Latent.Output(display_name="audio_latent", tooltip="Auto-generated audio latent (uses custom audio if enabled)."),
|
||||||
|
GuideData.Output(display_name="guide_data"),
|
||||||
|
io.Float.Output(display_name="frame_rate", tooltip="The frame rate used for the timeline."),
|
||||||
|
io.Audio.Output(display_name="combined_audio", tooltip="Combined timeline audio layout."),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def execute(cls, model, clip, global_prompt, duration_frames, duration_seconds,
|
||||||
|
timeline_data, local_prompts, segment_lengths, guide_strength="", epsilon=1e-3,
|
||||||
|
frame_rate=24, display_mode="seconds",
|
||||||
|
custom_width=768, custom_height=512, resize_method="maintain aspect ratio",
|
||||||
|
divisible_by=32, img_compression=0, audio_vae=None, optional_latent=None,
|
||||||
|
use_custom_audio=False) -> io.NodeOutput:
|
||||||
|
|
||||||
|
# --- Build guide_data from image segments FIRST (to derive output dimensions) ---
|
||||||
|
guide_data = {"images": [], "insert_frames": [], "strengths": [], "frame_rate": frame_rate}
|
||||||
|
derived_w, derived_h = custom_width, custom_height
|
||||||
|
try:
|
||||||
|
tdata = json.loads(timeline_data) if timeline_data else {}
|
||||||
|
img_segs = [
|
||||||
|
s for s in tdata.get("segments", [])
|
||||||
|
if s.get("type", "image") == "image"
|
||||||
|
and (s.get("imageFile") or s.get("imageB64"))
|
||||||
|
and int(s.get("start", 0)) < duration_frames # exclude segments fully outside duration
|
||||||
|
]
|
||||||
|
img_segs.sort(key=lambda s: s["start"])
|
||||||
|
|
||||||
|
strengths = []
|
||||||
|
if guide_strength.strip():
|
||||||
|
strengths = [float(x.strip()) for x in guide_strength.split(",") if x.strip()]
|
||||||
|
|
||||||
|
for idx, seg in enumerate(img_segs):
|
||||||
|
tensor = _load_image_tensor(seg)
|
||||||
|
|
||||||
|
# Apply resize
|
||||||
|
src_h, src_w = tensor.shape[1], tensor.shape[2]
|
||||||
|
|
||||||
|
def snap(val, div):
|
||||||
|
return max(div, (val // div) * div)
|
||||||
|
|
||||||
|
if custom_width > 0 and custom_height > 0:
|
||||||
|
# Both dimensions set — apply selected resize_method (pad, crop, stretch, maintain AR)
|
||||||
|
tensor = _resize_image(tensor, custom_width, custom_height, resize_method, divisible_by)
|
||||||
|
elif custom_width > 0:
|
||||||
|
# Width only — scale height from AR, snap both, then resize to exact dimensions
|
||||||
|
tgt_w = snap(custom_width, divisible_by)
|
||||||
|
tgt_h = snap(int(src_h * tgt_w / src_w), divisible_by)
|
||||||
|
tensor = _resize_image(tensor, tgt_w, tgt_h, "stretch to fit", divisible_by)
|
||||||
|
elif custom_height > 0:
|
||||||
|
# Height only — scale width from AR, snap both, then resize to exact dimensions
|
||||||
|
tgt_h = snap(custom_height, divisible_by)
|
||||||
|
tgt_w = snap(int(src_w * tgt_h / src_h), divisible_by)
|
||||||
|
tensor = _resize_image(tensor, tgt_w, tgt_h, "stretch to fit", divisible_by)
|
||||||
|
else:
|
||||||
|
# Both zero — keep original dimensions, just snap to divisible_by
|
||||||
|
tensor = _resize_image(tensor, src_w, src_h, "maintain aspect ratio", divisible_by)
|
||||||
|
|
||||||
|
|
||||||
|
# Apply compression
|
||||||
|
if img_compression > 0:
|
||||||
|
tensor = _compress_image(tensor, img_compression)
|
||||||
|
|
||||||
|
# Record dimensions of the first processed image for latent generation
|
||||||
|
if idx == 0:
|
||||||
|
derived_h = tensor.shape[1]
|
||||||
|
derived_w = tensor.shape[2]
|
||||||
|
|
||||||
|
strength = strengths[idx] if idx < len(strengths) else 1.0
|
||||||
|
guide_data["images"].append(tensor)
|
||||||
|
guide_data["insert_frames"].append(int(seg["start"]))
|
||||||
|
guide_data["strengths"].append(float(strength))
|
||||||
|
|
||||||
|
# If no images were loaded from the timeline, create a dummy image at strength 0
|
||||||
|
# to prevent artifacts in text-to-video mode.
|
||||||
|
if not guide_data["images"]:
|
||||||
|
w = derived_w if derived_w > 0 else 768
|
||||||
|
h = derived_h if derived_h > 0 else 512
|
||||||
|
w = (w // 32) * 32
|
||||||
|
h = (h // 32) * 32
|
||||||
|
|
||||||
|
dummy_image = torch.zeros((1, h, w, 3), dtype=torch.float32)
|
||||||
|
guide_data["images"].append(dummy_image)
|
||||||
|
guide_data["insert_frames"].append(0)
|
||||||
|
guide_data["strengths"].append(0.0)
|
||||||
|
|
||||||
|
derived_w = w
|
||||||
|
derived_h = h
|
||||||
|
except Exception as e:
|
||||||
|
log.warning("[PromptRelay] Could not build guide_data: %s", e)
|
||||||
|
|
||||||
|
# --- Auto-generate LTXV latent if none was provided ---
|
||||||
|
ltxv_length = duration_frames + 1
|
||||||
|
if optional_latent is None:
|
||||||
|
latent_w = max(32, (derived_w // 32) * 32)
|
||||||
|
latent_h = max(32, (derived_h // 32) * 32)
|
||||||
|
# LTXV temporal: ((length - 1) // 8) + 1 latent frames; invert to get pixel frames -> length
|
||||||
|
latent_t = ((ltxv_length - 1) // 8) + 1
|
||||||
|
samples = torch.zeros(
|
||||||
|
[1, 128, latent_t, latent_h // 32, latent_w // 32],
|
||||||
|
device=comfy.model_management.intermediate_device(),
|
||||||
|
)
|
||||||
|
latent = {"samples": samples}
|
||||||
|
log.info(
|
||||||
|
"[PromptRelay] Auto-generated LTXV latent: %dx%d, %d pixel frames (%d latent frames)",
|
||||||
|
latent_w, latent_h, ltxv_length, latent_t,
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
latent = optional_latent
|
||||||
|
|
||||||
|
patched, conditioning = _encode_relay(
|
||||||
|
model, clip, latent, global_prompt, local_prompts, segment_lengths, epsilon,
|
||||||
|
)
|
||||||
|
|
||||||
|
# --- Build Audio Output ---
|
||||||
|
audio_out = _build_combined_audio(timeline_data, ltxv_length, float(frame_rate))
|
||||||
|
|
||||||
|
# --- Audio Latent Generation ---
|
||||||
|
audio_latent = {}
|
||||||
|
|
||||||
|
if audio_vae is not None:
|
||||||
|
# Helper to generate empty latent
|
||||||
|
def get_empty_latent():
|
||||||
|
z_channels = audio_vae.latent_channels
|
||||||
|
audio_freq = audio_vae.first_stage_model.latent_frequency_bins
|
||||||
|
num_audio_latents = audio_vae.first_stage_model.num_of_latents_from_frames(ltxv_length, float(frame_rate))
|
||||||
|
audio_latents = torch.zeros(
|
||||||
|
(1, z_channels, num_audio_latents, audio_freq),
|
||||||
|
device=comfy.model_management.intermediate_device(),
|
||||||
|
)
|
||||||
|
return {"samples": audio_latents, "type": "audio"}
|
||||||
|
|
||||||
|
if use_custom_audio:
|
||||||
|
try:
|
||||||
|
if audio_out is not None:
|
||||||
|
# 1. Encode audio waveform into latent space
|
||||||
|
# VAE expects shape (batch, samples, channels), so we move dim 1 (channels) to the end
|
||||||
|
latent_samples = audio_vae.encode(audio_out["waveform"].movedim(1, -1))
|
||||||
|
|
||||||
|
if latent_samples.numel() == 0:
|
||||||
|
raise ValueError("Encoded audio latent is empty (0 elements).")
|
||||||
|
|
||||||
|
# 2. Create solid mask with value 0.0 (0 means keep/use conditioning, 1 means generate noise)
|
||||||
|
mask = torch.full(
|
||||||
|
(1, latent_samples.shape[-2], latent_samples.shape[-1]),
|
||||||
|
0.0,
|
||||||
|
dtype=torch.float32,
|
||||||
|
device=comfy.model_management.intermediate_device()
|
||||||
|
)
|
||||||
|
|
||||||
|
# 3. Set Latent Noise Mask
|
||||||
|
audio_latent = {
|
||||||
|
"samples": latent_samples,
|
||||||
|
"type": "audio",
|
||||||
|
"noise_mask": mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1]))
|
||||||
|
}
|
||||||
|
log.info("[PromptRelay] Generated custom audio latent with noise mask (value=0.0).")
|
||||||
|
else:
|
||||||
|
raise ValueError("No audio waveform to encode.")
|
||||||
|
except Exception as e:
|
||||||
|
log.error("[PromptRelay] Failed to generate custom audio latent: %s", e)
|
||||||
|
raise e
|
||||||
|
else:
|
||||||
|
# Generate empty latent
|
||||||
|
try:
|
||||||
|
audio_latent = get_empty_latent()
|
||||||
|
log.info("[PromptRelay] Auto-generated empty audio latent.")
|
||||||
|
except Exception as e:
|
||||||
|
log.warning("[PromptRelay] Could not generate empty audio latent: %s", e)
|
||||||
|
|
||||||
|
return io.NodeOutput(patched, conditioning, latent, audio_latent, guide_data, float(frame_rate), audio_out)
|
||||||
|
|
||||||
|
|
||||||
|
NODE_CLASS_MAPPINGS = {
|
||||||
|
"LTXDirector": LTXDirector,
|
||||||
|
}
|
||||||
|
|
||||||
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||||
|
"PromptRelayEncodeTimeline": "Prompt Relay Encode (Timeline)",
|
||||||
|
}
|
||||||
90
ltx_director_guide.py
Normal file
90
ltx_director_guide.py
Normal file
@@ -0,0 +1,90 @@
|
|||||||
|
from comfy_extras.nodes_lt import LTXVAddGuide
|
||||||
|
import torch
|
||||||
|
import comfy.utils
|
||||||
|
from comfy_api.latest import io
|
||||||
|
from .ltx_director import GuideData
|
||||||
|
|
||||||
|
|
||||||
|
class LTXDirectorGuide(LTXVAddGuide):
|
||||||
|
@classmethod
|
||||||
|
def define_schema(cls):
|
||||||
|
return io.Schema(
|
||||||
|
node_id="LTXDirectorGuide",
|
||||||
|
display_name="LTX Director Guide",
|
||||||
|
category="LTXVCustom",
|
||||||
|
description=(
|
||||||
|
"Applies guide images from a Prompt Relay Timeline node at the frame positions "
|
||||||
|
"and strengths defined on the timeline. Connect guide_data from the timeline node."
|
||||||
|
),
|
||||||
|
inputs=[
|
||||||
|
io.Conditioning.Input("positive", tooltip="Positive conditioning to add guide keyframe info to."),
|
||||||
|
io.Conditioning.Input("negative", tooltip="Negative conditioning to add guide keyframe info to."),
|
||||||
|
io.Vae.Input("vae", tooltip="Video VAE used to encode the guide images."),
|
||||||
|
io.Latent.Input("latent", tooltip="Video latent — guides are inserted into this latent."),
|
||||||
|
GuideData.Input("guide_data", tooltip="Guide data produced by Prompt Relay Encode (Timeline)."),
|
||||||
|
io.Float.Input("scale_by", default=1.0, min=0.01, max=8.0, step=0.01, tooltip="Scale the latent by this factor."),
|
||||||
|
io.Combo.Input("upscale_method", options=["nearest-exact", "bilinear", "area", "bicubic", "bislerp"], default="bicubic", tooltip="Method used to upscale/downscale the latent."),
|
||||||
|
],
|
||||||
|
outputs=[
|
||||||
|
io.Conditioning.Output(display_name="positive"),
|
||||||
|
io.Conditioning.Output(display_name="negative"),
|
||||||
|
io.Latent.Output(display_name="latent", tooltip="Video latent with guide frames applied."),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def execute(cls, positive, negative, vae, latent, guide_data, scale_by=1.0, upscale_method="bicubic") -> io.NodeOutput:
|
||||||
|
scale_factors = vae.downscale_index_formula
|
||||||
|
|
||||||
|
# Clone latents to avoid mutating upstream nodes
|
||||||
|
latent_image = latent["samples"].clone()
|
||||||
|
|
||||||
|
if "noise_mask" in latent:
|
||||||
|
noise_mask = latent["noise_mask"].clone()
|
||||||
|
else:
|
||||||
|
batch, _, latent_frames, latent_height, latent_width = latent_image.shape
|
||||||
|
noise_mask = torch.ones(
|
||||||
|
(batch, 1, latent_frames, 1, 1),
|
||||||
|
dtype=torch.float32,
|
||||||
|
device=latent_image.device,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Apply scale factor if not 1.0
|
||||||
|
if scale_by != 1.0:
|
||||||
|
B, C, F, H, W = latent_image.shape
|
||||||
|
width = round(W * scale_by)
|
||||||
|
height = round(H * scale_by)
|
||||||
|
|
||||||
|
# Reshape to 4D for common_upscale
|
||||||
|
latent_4d = latent_image.permute(0, 2, 1, 3, 4).reshape(B * F, C, H, W)
|
||||||
|
latent_resized_4d = comfy.utils.common_upscale(latent_4d, width, height, upscale_method, "disabled")
|
||||||
|
latent_image = latent_resized_4d.reshape(B, F, C, height, width).permute(0, 2, 1, 3, 4)
|
||||||
|
|
||||||
|
# Also resize noise mask if it's not a broadcasted mask
|
||||||
|
if noise_mask.shape[-1] > 1 or noise_mask.shape[-2] > 1:
|
||||||
|
mask_4d = noise_mask.permute(0, 2, 1, 3, 4).reshape(B * F, 1, H, W)
|
||||||
|
mask_resized_4d = comfy.utils.common_upscale(mask_4d, width, height, upscale_method, "disabled")
|
||||||
|
noise_mask = mask_resized_4d.reshape(B, F, 1, height, width).permute(0, 2, 1, 3, 4)
|
||||||
|
|
||||||
|
_, _, latent_length, latent_height, latent_width = latent_image.shape
|
||||||
|
|
||||||
|
images = guide_data.get("images", [])
|
||||||
|
insert_frames = guide_data.get("insert_frames", [])
|
||||||
|
strengths = guide_data.get("strengths", [])
|
||||||
|
|
||||||
|
for idx, img_tensor in enumerate(images):
|
||||||
|
f_idx = insert_frames[idx] if idx < len(insert_frames) else 0
|
||||||
|
strength = strengths[idx] if idx < len(strengths) else 1.0
|
||||||
|
|
||||||
|
image_1, t = cls.encode(vae, latent_width, latent_height, img_tensor, scale_factors)
|
||||||
|
frame_idx, latent_idx = cls.get_latent_index(positive, latent_length, len(image_1), f_idx, scale_factors)
|
||||||
|
|
||||||
|
assert latent_idx + t.shape[2] <= latent_length, (
|
||||||
|
f"Guide image {idx + 1}: conditioning frames exceed the length of the latent sequence."
|
||||||
|
)
|
||||||
|
|
||||||
|
positive, negative, latent_image, noise_mask = cls.append_keyframe(
|
||||||
|
positive, negative, frame_idx, latent_image, noise_mask, t, strength, scale_factors,
|
||||||
|
)
|
||||||
|
|
||||||
|
return io.NodeOutput(positive, negative, {"samples": latent_image, "noise_mask": noise_mask})
|
||||||
167
patches.py
Normal file
167
patches.py
Normal file
@@ -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}")
|
||||||
201
prompt_relay.py
Normal file
201
prompt_relay.py
Normal file
@@ -0,0 +1,201 @@
|
|||||||
|
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()
|
||||||
|
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
|
||||||
|
|
||||||
|
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, relay_options=None):
|
||||||
|
"""Per-segment metadata for the temporal penalty.
|
||||||
|
|
||||||
|
relay_options (optional dict) overrides per-stream knobs:
|
||||||
|
video_strength, video_window_scale,
|
||||||
|
audio_epsilon, audio_strength, audio_window_scale
|
||||||
|
Audio knobs only affect architectures whose cross-attention takes the scaled
|
||||||
|
(non-integer-frame) path — currently LTX audio_attn2.
|
||||||
|
"""
|
||||||
|
# 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
|
||||||
|
|
||||||
|
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
|
||||||
|
|
||||||
|
if relay_options:
|
||||||
|
log.info(
|
||||||
|
"[PromptRelay] Advanced options active — video: strength=%.3f window_scale=%.3f | "
|
||||||
|
"audio: epsilon=%s strength=%.3f window_scale=%.3f",
|
||||||
|
v_strength, v_window_scale,
|
||||||
|
f"{a_epsilon:.4f}" if a_epsilon is not None else "inherit",
|
||||||
|
a_strength, a_window_scale,
|
||||||
|
)
|
||||||
|
|
||||||
|
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):
|
||||||
|
"""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
|
||||||
@@ -1,7 +1,7 @@
|
|||||||
[project]
|
[project]
|
||||||
name = "WhatDreamsCost-ComfyUI"
|
name = "WhatDreamsCost-ComfyUI"
|
||||||
description = "A variety of custom ComfyUI nodes and workflows for creatives."
|
description = "A variety of custom ComfyUI nodes and workflows for creatives."
|
||||||
version = "1.2.9"
|
version = "1.3.0"
|
||||||
license = {file = "LICENSE"}
|
license = {file = "LICENSE"}
|
||||||
# classifiers = [
|
# classifiers = [
|
||||||
# # For OS-independent nodes (works on all operating systems)
|
# # For OS-independent nodes (works on all operating systems)
|
||||||
|
|||||||
Reference in New Issue
Block a user