refactor: extract director branch and audio helpers
This commit is contained in:
585
ltx_director.py
585
ltx_director.py
@@ -1340,6 +1340,262 @@ def _build_motion_guide_data(tdata: dict, use_custom_motion: bool, start_frame:
|
||||
log.warning("[LTXDirector] Could not build motion_guide_data: %s", e)
|
||||
|
||||
return motion_guide_data
|
||||
|
||||
|
||||
def _build_reference_mode_outputs(
|
||||
reference_mode: str,
|
||||
vae,
|
||||
global_prompt: str,
|
||||
local_prompts: str,
|
||||
segment_lengths: str,
|
||||
duration_frames: int,
|
||||
epsilon: float,
|
||||
char_images: list,
|
||||
char_slot_images: list,
|
||||
guide_data: dict,
|
||||
latent_w: int,
|
||||
latent_h: int,
|
||||
latent_grid_h: int,
|
||||
latent_grid_w: int,
|
||||
clean_latent_frames: int,
|
||||
reference_strength: float,
|
||||
optional_latent,
|
||||
ref_images,
|
||||
resize_method: str,
|
||||
divisible_by: int,
|
||||
prepare_tensor_image,
|
||||
model,
|
||||
clip,
|
||||
conditioning_neg,
|
||||
_dev,
|
||||
):
|
||||
if reference_mode == "Licon MSR (Prefix)":
|
||||
if vae is None:
|
||||
raise ValueError("Licon MSR (Prefix) ref option requires connecting the VAE to LTX Director!")
|
||||
|
||||
prompt_text = (global_prompt or "") + " " + (local_prompts or "")
|
||||
tag_pairs = [("@character1", "@char1"), ("@character2", "@char2"), ("@character3", "@char3")]
|
||||
referenced_slots = [i for i, tags in enumerate(tag_pairs) if any(t in prompt_text for t in tags)]
|
||||
selected = []
|
||||
for slot in referenced_slots:
|
||||
if slot < len(char_slot_images):
|
||||
selected.extend(char_slot_images[slot])
|
||||
if not selected:
|
||||
selected = char_images
|
||||
log.info("[LTXDirector] MSR slideshow subjects: referenced slots=%s, images=%d", referenced_slots, len(selected))
|
||||
|
||||
identity_images = [prepare_tensor_image(c, latent_w, latent_h) for c in selected]
|
||||
scene_images = list(guide_data["images"])
|
||||
if scene_images:
|
||||
bg_slide = scene_images[0]
|
||||
if bg_slide.shape[1] != latent_h or bg_slide.shape[2] != latent_w:
|
||||
bg_slide = _resize_image(bg_slide, latent_w, latent_h, resize_method, divisible_by)
|
||||
else:
|
||||
bg_slide = torch.zeros((1, latent_h, latent_w, 3), dtype=torch.float32)
|
||||
|
||||
slideshow_sources = identity_images + [bg_slide]
|
||||
base_count = MSR_PREFIX_FRAMES // len(slideshow_sources)
|
||||
remainder = MSR_PREFIX_FRAMES % len(slideshow_sources)
|
||||
slideshow_tensors = []
|
||||
for index, src_img in enumerate(slideshow_sources):
|
||||
repeats = base_count + (1 if index < remainder else 0)
|
||||
slideshow_tensors.extend([src_img[0]] * repeats)
|
||||
slideshow_video = torch.stack(slideshow_tensors)
|
||||
|
||||
keyframe_images = slideshow_sources
|
||||
prefix_latents = ((MSR_PREFIX_FRAMES - 1) // 8) + 1
|
||||
tail_latents = prefix_latents
|
||||
total_latent_frames = clean_latent_frames + tail_latents
|
||||
tail_pixels = tail_latents * 8
|
||||
|
||||
local_part = local_prompts.strip() if local_prompts.strip() else global_prompt
|
||||
clean_lengths = segment_lengths.strip() if segment_lengths.strip() else str(duration_frames)
|
||||
injected_local = f"{local_part} | {global_prompt}"
|
||||
injected_lengths = f"{clean_lengths},{tail_pixels}"
|
||||
|
||||
dummy_full = {"samples": torch.zeros([1, 128, total_latent_frames, latent_grid_h, latent_grid_w], device=_dev)}
|
||||
patched, conditioning = _encode_relay(
|
||||
model, clip, dummy_full, global_prompt, injected_local, injected_lengths, epsilon,
|
||||
)
|
||||
|
||||
if optional_latent is None:
|
||||
latent = {
|
||||
"samples": torch.zeros([1, 128, clean_latent_frames, latent_grid_h, latent_grid_w], device=_dev),
|
||||
"noise_mask": torch.ones((1, 1, clean_latent_frames, latent_grid_h, latent_grid_w), dtype=torch.float32, device=_dev),
|
||||
}
|
||||
else:
|
||||
latent = optional_latent
|
||||
|
||||
out_guide_data = {
|
||||
"images": [],
|
||||
"insert_frames": [],
|
||||
"strengths": [],
|
||||
"frame_rate": guide_data.get("frame_rate"),
|
||||
"msr": {
|
||||
"slideshow": slideshow_video,
|
||||
"keyframes": keyframe_images,
|
||||
"prefix_latents": int(prefix_latents),
|
||||
"strength": float(reference_strength),
|
||||
"downscale": float(MSR_LATENT_DOWNSCALE),
|
||||
"clean_latent_frames": int(clean_latent_frames),
|
||||
"negative": conditioning_neg,
|
||||
},
|
||||
}
|
||||
log.info(
|
||||
"[LTXDirector] Licon MSR Engine: clean=%d, tail=%d (prefix), %d keyframes for the guide node.",
|
||||
clean_latent_frames, tail_latents, len(keyframe_images),
|
||||
)
|
||||
return patched, conditioning, latent, out_guide_data
|
||||
|
||||
extra_refs = []
|
||||
if reference_mode == "Ghost Mask (End)" and ref_images is not None:
|
||||
try:
|
||||
for bi in range(ref_images.shape[0]):
|
||||
extra_refs.append(ref_images[bi:bi + 1])
|
||||
except Exception as e:
|
||||
log.warning("[LTXDirector] Could not process ref_images input: %s", e)
|
||||
|
||||
ghost_refs = char_images + extra_refs
|
||||
is_ghost = reference_mode == "Ghost Mask (End)" and bool(ghost_refs)
|
||||
n_refs = len(ghost_refs) if is_ghost else 0
|
||||
|
||||
if is_ghost:
|
||||
for i, single_ref in enumerate(ghost_refs):
|
||||
ref_tensor = prepare_tensor_image(single_ref, latent_w, latent_h)
|
||||
guide_data["images"].append(ref_tensor)
|
||||
guide_data["insert_frames"].append((clean_latent_frames + i) * 8)
|
||||
guide_data["strengths"].append(float(reference_strength))
|
||||
guide_data["ghost_pre_extend"] = int(n_refs)
|
||||
|
||||
total_latent_frames = clean_latent_frames + n_refs
|
||||
tail_pixels = n_refs * 8
|
||||
local_part = local_prompts.strip() if local_prompts.strip() else global_prompt
|
||||
clean_lengths = segment_lengths.strip() if segment_lengths.strip() else str(duration_frames)
|
||||
injected_local = f"{local_part} | {global_prompt}"
|
||||
injected_lengths = f"{clean_lengths},{tail_pixels}"
|
||||
dummy_full = {"samples": torch.zeros([1, 128, total_latent_frames, latent_grid_h, latent_grid_w], device=_dev)}
|
||||
patched, conditioning = _encode_relay(
|
||||
model, clip, dummy_full, global_prompt, injected_local, injected_lengths, epsilon,
|
||||
)
|
||||
|
||||
if optional_latent is None:
|
||||
latent = {
|
||||
"samples": torch.zeros([1, 128, clean_latent_frames, latent_grid_h, latent_grid_w], device=_dev),
|
||||
"noise_mask": torch.ones((1, 1, clean_latent_frames, latent_grid_h, latent_grid_w), dtype=torch.float32, device=_dev),
|
||||
}
|
||||
else:
|
||||
latent = optional_latent
|
||||
|
||||
log.info(
|
||||
"[LTXDirector] Ghost Mask: %d references; guide will hide them in a tail past "
|
||||
"frame %d (clean_latent_frames=%d). Enable Clean Latent Slice (length=%d) downstream.",
|
||||
n_refs, clean_latent_frames, clean_latent_frames, clean_latent_frames,
|
||||
)
|
||||
return patched, conditioning, latent, guide_data
|
||||
|
||||
if optional_latent is None:
|
||||
latent = {"samples": torch.zeros([1, 128, clean_latent_frames, latent_grid_h, latent_grid_w], device=_dev)}
|
||||
log.info(
|
||||
"[PromptRelay] Auto-generated LTXV latent: %dx%d, %d pixel frames (%d latent frames)",
|
||||
latent_w, latent_h, ((clean_latent_frames - 1) * 8) + 1, clean_latent_frames,
|
||||
)
|
||||
else:
|
||||
latent = optional_latent
|
||||
|
||||
patched, conditioning = _encode_relay(
|
||||
model, clip, latent, global_prompt, local_prompts, segment_lengths, epsilon,
|
||||
)
|
||||
return patched, conditioning, latent, guide_data
|
||||
|
||||
|
||||
def _build_audio_latent(
|
||||
audio_vae,
|
||||
audio_out,
|
||||
use_custom_audio: bool,
|
||||
override_audio: bool,
|
||||
is_retake_active: bool,
|
||||
inpaint_audio: bool,
|
||||
tdata: dict,
|
||||
start_frame: int,
|
||||
ltxv_length: int,
|
||||
frame_rate_f: float,
|
||||
):
|
||||
if audio_vae is None:
|
||||
return {}
|
||||
|
||||
def get_empty_latent():
|
||||
inner = getattr(audio_vae, "first_stage_model", audio_vae)
|
||||
z_channels = audio_vae.latent_channels
|
||||
audio_freq = inner.latent_frequency_bins
|
||||
num_audio_latents = inner.num_of_latents_from_frames(ltxv_length, frame_rate_f)
|
||||
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 not (use_custom_audio or override_audio or is_retake_active):
|
||||
audio_latent = get_empty_latent()
|
||||
log.info("[PromptRelay] Auto-generated empty audio latent.")
|
||||
return audio_latent
|
||||
|
||||
if audio_out is None:
|
||||
raise ValueError("No audio waveform to encode.")
|
||||
|
||||
waveform = audio_out["waveform"]
|
||||
if waveform.ndim == 2:
|
||||
waveform = waveform.unsqueeze(0)
|
||||
if waveform.ndim != 3:
|
||||
raise ValueError(f"Expected custom audio waveform with 2 or 3 dims, got shape {tuple(waveform.shape)}")
|
||||
|
||||
if hasattr(audio_vae, "first_stage_model"):
|
||||
latent_samples = audio_vae.encode(waveform.movedim(1, -1))
|
||||
else:
|
||||
latent_samples = audio_vae.encode({
|
||||
"waveform": waveform,
|
||||
"sample_rate": audio_out["sample_rate"],
|
||||
})
|
||||
|
||||
if latent_samples.numel() == 0:
|
||||
raise ValueError("Encoded audio latent is empty (0 elements).")
|
||||
|
||||
B, _C, F_len, H_len = latent_samples.shape
|
||||
if is_retake_active:
|
||||
gap_mask = torch.zeros((B, F_len, H_len), dtype=torch.float32, device=latent_samples.device)
|
||||
retake_start = float(tdata.get("retakeStart", 0))
|
||||
retake_len = float(tdata.get("retakeLength", 0))
|
||||
overlap_start = max(start_frame, retake_start)
|
||||
overlap_end = min(start_frame + ltxv_length, retake_start + retake_len)
|
||||
if overlap_end > overlap_start:
|
||||
rel_start = overlap_start - start_frame
|
||||
rel_len = overlap_end - overlap_start
|
||||
start_idx, end_idx = _time_range_to_latent_indices(rel_start, rel_len, 0, ltxv_length, frame_rate_f, F_len)
|
||||
gap_mask[:, start_idx:end_idx, :] = 1.0
|
||||
else:
|
||||
gap_mask = torch.ones((B, F_len, H_len), dtype=torch.float32, device=latent_samples.device)
|
||||
audio_segs_key = "motionSegments" if override_audio else "audioSegments"
|
||||
file_key = "videoFile" if override_audio else "audioFile"
|
||||
for seg in tdata.get(audio_segs_key, []):
|
||||
if not seg.get(file_key):
|
||||
continue
|
||||
seg_start = float(seg.get("start", 0))
|
||||
seg_len = float(seg.get("length", 1))
|
||||
if seg_start + seg_len <= start_frame or seg_start >= start_frame + ltxv_length:
|
||||
continue
|
||||
offset = max(0, start_frame - seg_start)
|
||||
seg_len = max(1.0, seg_len - offset)
|
||||
seg_start = max(0, seg_start - start_frame)
|
||||
start_idx, end_idx = _time_range_to_latent_indices(seg_start, seg_len, 0, ltxv_length, frame_rate_f, F_len)
|
||||
gap_mask[:, start_idx:end_idx, :] = 0.0
|
||||
|
||||
if inpaint_audio:
|
||||
mask = gap_mask
|
||||
else:
|
||||
mask = torch.zeros((B, F_len, H_len), dtype=torch.float32, device=latent_samples.device)
|
||||
|
||||
audio_latent = {"samples": latent_samples, "type": "audio", "noise_mask": mask}
|
||||
log.info("[PromptRelay] Generated custom audio latent with dynamic noise mask.")
|
||||
return audio_latent
|
||||
|
||||
|
||||
class LTXDirector(io.ComfyNode):
|
||||
@@ -1619,295 +1875,54 @@ class LTXDirector(io.ComfyNode):
|
||||
if not (global_prompt or "").strip() and local_prompts:
|
||||
global_prompt = local_prompts.split("|")[0].strip()
|
||||
|
||||
_dev = comfy.model_management.intermediate_device()
|
||||
|
||||
if reference_mode == "Licon MSR (Prefix)":
|
||||
# ============ LICON MSR DIRECTOR ENGINE ============
|
||||
# Character slots drive identity via an IC-LoRA reference slideshow + per-character
|
||||
# keyframes, handed to the downstream LTXDirectorGuide through guide_data["msr"].
|
||||
if vae is None:
|
||||
raise ValueError("Licon MSR (Prefix) ref option requires connecting the VAE to LTX Director!")
|
||||
|
||||
# Honour @charN tags: select only the slots the prompt references; else all filled slots.
|
||||
_prompt_text = (global_prompt or "") + " " + (local_prompts or "")
|
||||
_tag_pairs = [("@character1", "@char1"), ("@character2", "@char2"), ("@character3", "@char3")]
|
||||
_referenced_slots = [i for i, tags in enumerate(_tag_pairs) if any(t in _prompt_text for t in tags)]
|
||||
_selected = []
|
||||
for _slot in _referenced_slots:
|
||||
if _slot < len(char_slot_images):
|
||||
_selected.extend(char_slot_images[_slot])
|
||||
if not _selected:
|
||||
_selected = char_images
|
||||
log.info("[LTXDirector] MSR slideshow subjects: referenced slots=%s, images=%d", _referenced_slots, len(_selected))
|
||||
|
||||
identity_images = [
|
||||
prepare_tensor_image(c, latent_w, latent_h) for c in _selected
|
||||
]
|
||||
|
||||
# Scene background = first timeline image (already in guide_data), else black.
|
||||
scene_images = list(guide_data["images"])
|
||||
if scene_images:
|
||||
bg_slide = scene_images[0]
|
||||
if bg_slide.shape[1] != latent_h or bg_slide.shape[2] != latent_w:
|
||||
bg_slide = _resize_image(bg_slide, latent_w, latent_h, resize_method, divisible_by)
|
||||
else:
|
||||
bg_slide = torch.zeros((1, latent_h, latent_w, 3), dtype=torch.float32)
|
||||
|
||||
# Build the slideshow: characters then background, distributed across MSR_PREFIX_FRAMES.
|
||||
slideshow_sources = identity_images + [bg_slide]
|
||||
base_count = MSR_PREFIX_FRAMES // len(slideshow_sources)
|
||||
remainder = MSR_PREFIX_FRAMES % len(slideshow_sources)
|
||||
slideshow_tensors = []
|
||||
for index, src_img in enumerate(slideshow_sources):
|
||||
repeats = base_count + (1 if index < remainder else 0)
|
||||
slideshow_tensors.extend([src_img[0]] * repeats)
|
||||
slideshow_video = torch.stack(slideshow_tensors)
|
||||
|
||||
keyframe_images = slideshow_sources
|
||||
prefix_latents = ((MSR_PREFIX_FRAMES - 1) // 8) + 1
|
||||
tail_latents = prefix_latents
|
||||
total_latent_frames = clean_latent_frames + tail_latents
|
||||
tail_pixels = tail_latents * 8
|
||||
|
||||
# Relay conditioning over the runtime length: clean -> locals, tail -> global.
|
||||
local_part = local_prompts.strip() if local_prompts.strip() else global_prompt
|
||||
clean_lengths = segment_lengths.strip() if segment_lengths.strip() else str(duration_frames)
|
||||
injected_local = f"{local_part} | {global_prompt}"
|
||||
injected_lengths = f"{clean_lengths},{tail_pixels}"
|
||||
|
||||
dummy_full = {"samples": torch.zeros(
|
||||
[1, 128, total_latent_frames, latent_grid_h, latent_grid_w], device=_dev,
|
||||
)}
|
||||
patched, conditioning = _encode_relay(
|
||||
model, clip, dummy_full, global_prompt, injected_local, injected_lengths, epsilon,
|
||||
)
|
||||
|
||||
# Clean base latent (the true visible region); the guide node pads + appends keyframes per stage.
|
||||
if optional_latent is None:
|
||||
latent = {
|
||||
"samples": torch.zeros([1, 128, clean_latent_frames, latent_grid_h, latent_grid_w], device=_dev),
|
||||
"noise_mask": torch.ones((1, 1, clean_latent_frames, latent_grid_h, latent_grid_w), dtype=torch.float32, device=_dev),
|
||||
}
|
||||
else:
|
||||
latent = optional_latent
|
||||
|
||||
# Hand the RAW full-res references to the guide node via guide_data["msr"].
|
||||
guide_data = {
|
||||
"images": [], "insert_frames": [], "strengths": [], "frame_rate": frame_rate_f,
|
||||
"msr": {
|
||||
"slideshow": slideshow_video,
|
||||
"keyframes": keyframe_images,
|
||||
"prefix_latents": int(prefix_latents),
|
||||
"strength": float(reference_strength),
|
||||
"downscale": float(MSR_LATENT_DOWNSCALE),
|
||||
"clean_latent_frames": int(clean_latent_frames),
|
||||
"negative": conditioning_neg,
|
||||
},
|
||||
}
|
||||
log.info(
|
||||
"[LTXDirector] Licon MSR Engine: clean=%d, tail=%d (prefix), %d keyframes for the guide node.",
|
||||
clean_latent_frames, tail_latents, len(keyframe_images),
|
||||
)
|
||||
|
||||
else:
|
||||
# ============ OFF / Ghost Mask (End) ============
|
||||
# OFF runs over the clean latent. Ghost Mask hides one reference per character in a
|
||||
# tail PAST the clean region, each in its own frame so identities never blend. To stay
|
||||
# consistent with the audio latent (which is sized to the CLEAN length), the Director —
|
||||
# exactly like Licon MSR — outputs a CLEAN video latent and lets the GUIDE pre-extend
|
||||
# it. This keeps video and audio matched at the AV-concat stage; extending the video
|
||||
# inside the Director instead would desync them and crash the joint LTX2 sampler.
|
||||
# Extra Ghost reference images (e.g. an object) from the optional "ref_images" input.
|
||||
# Split a (possibly batched) IMAGE into single frames and append them AFTER the @char
|
||||
# references as their own block — independent of how many characters are loaded.
|
||||
extra_refs = []
|
||||
if reference_mode == "Ghost Mask (End)" and ref_images is not None:
|
||||
try:
|
||||
for _bi in range(ref_images.shape[0]):
|
||||
extra_refs.append(ref_images[_bi:_bi + 1])
|
||||
except Exception as _e:
|
||||
log.warning("[LTXDirector] Could not process ref_images input: %s", _e)
|
||||
|
||||
ghost_refs = char_images + extra_refs
|
||||
is_ghost = (reference_mode == "Ghost Mask (End)" and bool(ghost_refs))
|
||||
n_refs = len(ghost_refs) if is_ghost else 0
|
||||
|
||||
if is_ghost:
|
||||
# Build the reference list now (so the guide can encode + anchor them), but keep
|
||||
# the emitted latent clean. Each reference anchors at (clean_latent_frames + i).
|
||||
for i, single_ref in enumerate(ghost_refs):
|
||||
ref_tensor = prepare_tensor_image(single_ref, latent_w, latent_h)
|
||||
guide_data["images"].append(ref_tensor)
|
||||
guide_data["insert_frames"].append((clean_latent_frames + i) * 8)
|
||||
guide_data["strengths"].append(float(reference_strength))
|
||||
# Tell the guide to pre-extend the (clean) latent by this many frames before it
|
||||
# processes the references, so the tail anchors land in-range. The guide pads;
|
||||
# the Director does not — keeping video/audio lengths matched downstream.
|
||||
guide_data["ghost_pre_extend"] = int(n_refs)
|
||||
|
||||
# Relay conditioning over the EXTENDED length (clean + tail), to match the latent
|
||||
# length the guide will hand the sampler. Clean region -> locals, tail -> global.
|
||||
total_latent_frames = clean_latent_frames + n_refs
|
||||
tail_pixels = n_refs * 8
|
||||
local_part = local_prompts.strip() if local_prompts.strip() else global_prompt
|
||||
clean_lengths = segment_lengths.strip() if segment_lengths.strip() else str(duration_frames)
|
||||
injected_local = f"{local_part} | {global_prompt}"
|
||||
injected_lengths = f"{clean_lengths},{tail_pixels}"
|
||||
dummy_full = {"samples": torch.zeros(
|
||||
[1, 128, total_latent_frames, latent_grid_h, latent_grid_w], device=_dev,
|
||||
)}
|
||||
patched, conditioning = _encode_relay(
|
||||
model, clip, dummy_full, global_prompt, injected_local, injected_lengths, epsilon,
|
||||
)
|
||||
|
||||
if optional_latent is None:
|
||||
latent = {
|
||||
"samples": torch.zeros([1, 128, clean_latent_frames, latent_grid_h, latent_grid_w], device=_dev),
|
||||
"noise_mask": torch.ones((1, 1, clean_latent_frames, latent_grid_h, latent_grid_w), dtype=torch.float32, device=_dev),
|
||||
}
|
||||
else:
|
||||
latent = optional_latent
|
||||
|
||||
log.info(
|
||||
"[LTXDirector] Ghost Mask: %d references; guide will hide them in a tail past "
|
||||
"frame %d (clean_latent_frames=%d). Enable Clean Latent Slice (length=%d) downstream.",
|
||||
n_refs, clean_latent_frames, clean_latent_frames, clean_latent_frames,
|
||||
)
|
||||
|
||||
else:
|
||||
# OFF: plain clean latent + relay.
|
||||
if optional_latent is None:
|
||||
samples = torch.zeros(
|
||||
[1, 128, clean_latent_frames, latent_grid_h, latent_grid_w], device=_dev,
|
||||
)
|
||||
latent = {"samples": samples}
|
||||
log.info(
|
||||
"[PromptRelay] Auto-generated LTXV latent: %dx%d, %d pixel frames (%d latent frames)",
|
||||
latent_w, latent_h, ltxv_length, clean_latent_frames,
|
||||
)
|
||||
else:
|
||||
latent = optional_latent
|
||||
|
||||
patched, conditioning = _encode_relay(
|
||||
model, clip, latent, global_prompt, local_prompts, segment_lengths, epsilon,
|
||||
)
|
||||
_dev = comfy.model_management.intermediate_device()
|
||||
patched, conditioning, latent, guide_data = _build_reference_mode_outputs(
|
||||
reference_mode=reference_mode,
|
||||
vae=vae,
|
||||
global_prompt=global_prompt,
|
||||
local_prompts=local_prompts,
|
||||
segment_lengths=segment_lengths,
|
||||
duration_frames=duration_frames,
|
||||
epsilon=epsilon,
|
||||
char_images=char_images,
|
||||
char_slot_images=char_slot_images,
|
||||
guide_data=guide_data,
|
||||
latent_w=latent_w,
|
||||
latent_h=latent_h,
|
||||
latent_grid_h=latent_grid_h,
|
||||
latent_grid_w=latent_grid_w,
|
||||
clean_latent_frames=clean_latent_frames,
|
||||
reference_strength=reference_strength,
|
||||
optional_latent=optional_latent,
|
||||
ref_images=ref_images,
|
||||
resize_method=resize_method,
|
||||
divisible_by=divisible_by,
|
||||
prepare_tensor_image=prepare_tensor_image,
|
||||
model=model,
|
||||
clip=clip,
|
||||
conditioning_neg=conditioning_neg,
|
||||
_dev=_dev,
|
||||
)
|
||||
|
||||
# --- Build Audio Output ---
|
||||
audio_out = _build_combined_audio(tdata, start_frame, ltxv_length, frame_rate_f, override_audio=override_audio)
|
||||
|
||||
# --- Audio Latent Generation ---
|
||||
audio_latent = {}
|
||||
|
||||
if audio_vae is not None:
|
||||
# Helper to generate empty latent
|
||||
def get_empty_latent():
|
||||
# Support both raw AudioVAE objects and ComfyUI VAE wrappers.
|
||||
inner = getattr(audio_vae, "first_stage_model", audio_vae)
|
||||
z_channels = audio_vae.latent_channels
|
||||
audio_freq = inner.latent_frequency_bins
|
||||
num_audio_latents = inner.num_of_latents_from_frames(ltxv_length, frame_rate_f)
|
||||
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 or override_audio or is_retake_active:
|
||||
try:
|
||||
if audio_out is not None:
|
||||
# 1. Encode audio waveform into latent space
|
||||
waveform = audio_out["waveform"]
|
||||
if waveform.ndim == 2:
|
||||
waveform = waveform.unsqueeze(0)
|
||||
if waveform.ndim != 3:
|
||||
raise ValueError(
|
||||
f"Expected custom audio waveform with 2 or 3 dims, got shape {tuple(waveform.shape)}"
|
||||
)
|
||||
|
||||
# Wrapped ComfyUI VAE expects (batch, samples, channels);
|
||||
# raw AudioVAE expects a dict with waveform in (batch, channels, samples).
|
||||
if hasattr(audio_vae, "first_stage_model"):
|
||||
latent_samples = audio_vae.encode(waveform.movedim(1, -1))
|
||||
else:
|
||||
latent_samples = audio_vae.encode({
|
||||
"waveform": waveform,
|
||||
"sample_rate": audio_out["sample_rate"],
|
||||
})
|
||||
|
||||
if latent_samples.numel() == 0:
|
||||
raise ValueError("Encoded audio latent is empty (0 elements).")
|
||||
|
||||
# 2. Create a 3D gap mask [B, F, H] to avoid accidental broadcasting to the 5D video latent
|
||||
# which also has 128 channels. A 4D audio mask [1, 128, F, H] confuses ComfyUI's KSampler
|
||||
# into masking the video latent as well, causing black frames.
|
||||
B, C, F_len, H_len = latent_samples.shape
|
||||
|
||||
if is_retake_active:
|
||||
gap_mask = torch.zeros((B, F_len, H_len), dtype=torch.float32, device=latent_samples.device)
|
||||
|
||||
retake_start = float(tdata.get("retakeStart", 0))
|
||||
retake_len = float(tdata.get("retakeLength", 0))
|
||||
|
||||
overlap_start = max(start_frame, retake_start)
|
||||
overlap_end = min(start_frame + ltxv_length, retake_start + retake_len)
|
||||
|
||||
if overlap_end > overlap_start:
|
||||
rel_start = overlap_start - start_frame
|
||||
rel_len = overlap_end - overlap_start
|
||||
start_idx, end_idx = _time_range_to_latent_indices(
|
||||
rel_start, rel_len, 0, ltxv_length, frame_rate_f, F_len
|
||||
)
|
||||
gap_mask[:, start_idx:end_idx, :] = 1.0
|
||||
else:
|
||||
gap_mask = torch.ones((B, F_len, H_len), dtype=torch.float32, device=latent_samples.device)
|
||||
|
||||
audio_segs_key = "motionSegments" if override_audio else "audioSegments"
|
||||
file_key = "videoFile" if override_audio else "audioFile"
|
||||
for seg in tdata.get(audio_segs_key, []):
|
||||
if not seg.get(file_key):
|
||||
continue
|
||||
|
||||
seg_start = float(seg.get("start", 0))
|
||||
seg_len = float(seg.get("length", 1))
|
||||
|
||||
if seg_start + seg_len <= start_frame or seg_start >= start_frame + ltxv_length:
|
||||
continue
|
||||
|
||||
offset = max(0, start_frame - seg_start)
|
||||
seg_len = max(1.0, seg_len - offset)
|
||||
seg_start = max(0, seg_start - start_frame)
|
||||
start_idx, end_idx = _time_range_to_latent_indices(
|
||||
seg_start, seg_len, 0, ltxv_length, frame_rate_f, F_len
|
||||
)
|
||||
gap_mask[:, start_idx:end_idx, :] = 0.0
|
||||
|
||||
if inpaint_audio:
|
||||
# Generate new audio in the gaps, preserve custom audio segments
|
||||
mask = gap_mask
|
||||
else:
|
||||
# Preserve the entire audio latent (no generation).
|
||||
# We use a 3D zeros mask to prevent video blackouts.
|
||||
mask = torch.zeros((B, F_len, H_len), dtype=torch.float32, device=latent_samples.device)
|
||||
|
||||
audio_latent = {
|
||||
"samples": latent_samples,
|
||||
"type": "audio",
|
||||
"noise_mask": mask
|
||||
}
|
||||
log.info("[PromptRelay] Generated custom audio latent with dynamic noise mask.")
|
||||
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.error("[PromptRelay] Could not generate empty audio latent: %s", e)
|
||||
raise e
|
||||
try:
|
||||
audio_latent = _build_audio_latent(
|
||||
audio_vae=audio_vae,
|
||||
audio_out=audio_out,
|
||||
use_custom_audio=use_custom_audio,
|
||||
override_audio=override_audio,
|
||||
is_retake_active=is_retake_active,
|
||||
inpaint_audio=inpaint_audio,
|
||||
tdata=tdata,
|
||||
start_frame=start_frame,
|
||||
ltxv_length=ltxv_length,
|
||||
frame_rate_f=frame_rate_f,
|
||||
)
|
||||
except Exception as e:
|
||||
log.error("[PromptRelay] Failed to generate custom audio latent: %s", e)
|
||||
raise e
|
||||
|
||||
motion_guide_data = _build_motion_guide_data(
|
||||
tdata=tdata,
|
||||
|
||||
Reference in New Issue
Block a user