354 lines
15 KiB
Python
354 lines
15 KiB
Python
"""
|
|
Monkey-patch for multitalk_loop.py prompt selection.
|
|
|
|
Replaces the original prompt selection block (iteration-based) with a
|
|
frame-based prompt_schedule-aware version. Falls back to original behavior
|
|
when no prompt_schedule is present in text_embeds.
|
|
"""
|
|
|
|
import logging
|
|
import textwrap
|
|
|
|
log = logging.getLogger("ComfyUI-MultiTalkPromptSchedule")
|
|
|
|
_PATCHED = False
|
|
|
|
|
|
def _select_prompt(text_embeds, audio_start_idx, iteration_count, log_fn):
|
|
"""Frame-based prompt selection with fallback to iteration-based.
|
|
|
|
Returns:
|
|
positive: list of prompt embeddings for this window
|
|
"""
|
|
prompt_schedule = text_embeds.get("prompt_schedule", None)
|
|
nag_schedule = text_embeds.get("nag_schedule", None)
|
|
|
|
if prompt_schedule is not None:
|
|
# Restore original NAG embeds before swapping (prevents cache corruption)
|
|
if "_orig_nag_prompt_embeds" not in text_embeds:
|
|
import copy as _copy
|
|
text_embeds["_orig_nag_prompt_embeds"] = _copy.deepcopy(text_embeds.get("nag_prompt_embeds"))
|
|
text_embeds["_orig_negative_prompt_embeds"] = _copy.deepcopy(text_embeds.get("negative_prompt_embeds"))
|
|
else:
|
|
text_embeds["nag_prompt_embeds"] = text_embeds["_orig_nag_prompt_embeds"]
|
|
text_embeds["negative_prompt_embeds"] = text_embeds["_orig_negative_prompt_embeds"]
|
|
|
|
# Frame-based prompt selection (MultiTalkPromptSchedule node)
|
|
prompt_index = len(prompt_schedule) - 1 # default to last
|
|
for idx, (sched_start, sched_end) in enumerate(prompt_schedule):
|
|
if audio_start_idx < sched_end:
|
|
prompt_index = idx
|
|
break
|
|
positive = [text_embeds["prompt_embeds"][prompt_index]]
|
|
|
|
# Per-scene NAG negative if available
|
|
if nag_schedule is not None and prompt_index in nag_schedule:
|
|
nag_index = nag_schedule[prompt_index]
|
|
nag_embed = [text_embeds["all_nag_embeds"][nag_index]]
|
|
text_embeds["nag_prompt_embeds"] = nag_embed
|
|
text_embeds["negative_prompt_embeds"] = nag_embed
|
|
log_fn(f"Using per-scene NAG negative for scene {prompt_index}")
|
|
elif nag_schedule is not None:
|
|
if "all_nag_embeds" in text_embeds:
|
|
text_embeds["nag_prompt_embeds"] = [text_embeds["all_nag_embeds"][0]]
|
|
text_embeds["negative_prompt_embeds"] = [text_embeds["all_nag_embeds"][0]]
|
|
|
|
sched_start, sched_end = prompt_schedule[prompt_index]
|
|
log_fn(
|
|
f"Using scheduled prompt {prompt_index} "
|
|
f"(frames {sched_start}-{sched_end}) "
|
|
f"at audio_start_idx={audio_start_idx}"
|
|
)
|
|
elif len(text_embeds["prompt_embeds"]) > 1:
|
|
# Fallback: iteration-based (original Kijai behavior, pipe separator)
|
|
prompt_index = min(
|
|
iteration_count, len(text_embeds["prompt_embeds"]) - 1
|
|
)
|
|
positive = [text_embeds["prompt_embeds"][prompt_index]]
|
|
log_fn(f"Using prompt index: {prompt_index}")
|
|
else:
|
|
positive = text_embeds["prompt_embeds"]
|
|
|
|
return positive
|
|
|
|
|
|
def apply_prompt_schedule_patch():
|
|
"""Monkey-patch multitalk_loop to support prompt_schedule in text_embeds."""
|
|
global _PATCHED
|
|
if _PATCHED:
|
|
return
|
|
|
|
import importlib
|
|
|
|
loop_mod = None
|
|
# Try both naming conventions (hyphen vs underscore)
|
|
for module_name in [
|
|
"custom_nodes.ComfyUI-WanVideoWrapper.multitalk.multitalk_loop",
|
|
"custom_nodes.ComfyUI_WanVideoWrapper.multitalk.multitalk_loop",
|
|
]:
|
|
try:
|
|
loop_mod = importlib.import_module(module_name)
|
|
break
|
|
except ImportError:
|
|
continue
|
|
|
|
if loop_mod is None:
|
|
log.warning(
|
|
"[MultiTalkPromptSchedule] ComfyUI-WanVideoWrapper not found. "
|
|
"Patch not applied — the node will still load but scheduling "
|
|
"won't work until WanVideoWrapper is installed."
|
|
)
|
|
return
|
|
|
|
import inspect
|
|
import re
|
|
|
|
original_fn = loop_mod.multitalk_loop
|
|
|
|
# Get the source to check if the patch is already present
|
|
# (e.g. user already manually patched)
|
|
try:
|
|
src = inspect.getsource(original_fn)
|
|
if "prompt_schedule" in src:
|
|
log.info(
|
|
"[MultiTalkPromptSchedule] multitalk_loop already contains "
|
|
"prompt_schedule logic. Skipping monkey-patch."
|
|
)
|
|
_PATCHED = True
|
|
return
|
|
except (OSError, TypeError):
|
|
# Can't inspect source (compiled, etc.) — proceed with patch
|
|
pass
|
|
|
|
# ── Build the patched function ──────────────────────────────────────
|
|
# We wrap the original function and intercept the prompt selection.
|
|
# The trick: multitalk_loop is a method (self, **kwargs). We wrap it
|
|
# so that before it runs, we inject a _select_prompt callback into
|
|
# text_embeds that the original code path will pick up.
|
|
#
|
|
# Strategy: We pre-process text_embeds so that when the original code
|
|
# hits `if len(text_embeds["prompt_embeds"]) > 1:`, the embeds are
|
|
# already set to the correct prompt for the FIRST window. For
|
|
# subsequent windows we can't intercept since the loop is internal.
|
|
#
|
|
# Better strategy: We actually patch the source code of the function
|
|
# at the AST/text level to inject our prompt_schedule block.
|
|
|
|
src = inspect.getsource(original_fn)
|
|
|
|
# The original prompt selection block we need to replace:
|
|
original_block = textwrap.dedent("""\
|
|
# Use the appropriate prompt for this section
|
|
if len(text_embeds["prompt_embeds"]) > 1:
|
|
prompt_index = min(iteration_count, len(text_embeds["prompt_embeds"]) - 1)
|
|
positive = [text_embeds["prompt_embeds"][prompt_index]]
|
|
log.info(f"Using prompt index: {prompt_index}")
|
|
else:
|
|
positive = text_embeds["prompt_embeds"]""")
|
|
|
|
patched_block = textwrap.dedent("""\
|
|
# Use the appropriate prompt for this section
|
|
# [Patched by ComfyUI-MultiTalkPromptSchedule]
|
|
log.info(f"[MultiTalkPromptSchedule PATCH] text_embeds keys: {list(text_embeds.keys())}")
|
|
prompt_schedule = text_embeds.get("prompt_schedule", None)
|
|
nag_schedule = text_embeds.get("nag_schedule", None)
|
|
if prompt_schedule is not None:
|
|
# Restore original NAG embeds before swapping (prevents cache corruption between runs)
|
|
if "_orig_nag_prompt_embeds" not in text_embeds:
|
|
import copy as _copy
|
|
text_embeds["_orig_nag_prompt_embeds"] = _copy.deepcopy(text_embeds.get("nag_prompt_embeds"))
|
|
text_embeds["_orig_negative_prompt_embeds"] = _copy.deepcopy(text_embeds.get("negative_prompt_embeds"))
|
|
else:
|
|
text_embeds["nag_prompt_embeds"] = text_embeds["_orig_nag_prompt_embeds"]
|
|
text_embeds["negative_prompt_embeds"] = text_embeds["_orig_negative_prompt_embeds"]
|
|
prompt_index = len(prompt_schedule) - 1
|
|
for idx, (sched_start, sched_end) in enumerate(prompt_schedule):
|
|
if audio_start_idx < sched_end:
|
|
prompt_index = idx
|
|
break
|
|
positive = [text_embeds["prompt_embeds"][prompt_index]]
|
|
if nag_schedule is not None and prompt_index in nag_schedule:
|
|
nag_index = nag_schedule[prompt_index]
|
|
nag_embed = [text_embeds["all_nag_embeds"][nag_index]]
|
|
text_embeds["nag_prompt_embeds"] = nag_embed
|
|
text_embeds["negative_prompt_embeds"] = nag_embed
|
|
log.info(f"Using per-scene NAG negative for scene {prompt_index}")
|
|
elif nag_schedule is not None:
|
|
if "all_nag_embeds" in text_embeds:
|
|
text_embeds["nag_prompt_embeds"] = [text_embeds["all_nag_embeds"][0]]
|
|
text_embeds["negative_prompt_embeds"] = [text_embeds["all_nag_embeds"][0]]
|
|
sched_start, sched_end = prompt_schedule[prompt_index]
|
|
log.info(f"Using scheduled prompt {prompt_index} (frames {sched_start}-{sched_end}) at audio_start_idx={audio_start_idx}")
|
|
elif len(text_embeds["prompt_embeds"]) > 1:
|
|
prompt_index = min(iteration_count, len(text_embeds["prompt_embeds"]) - 1)
|
|
positive = [text_embeds["prompt_embeds"][prompt_index]]
|
|
log.info(f"Using prompt index: {prompt_index}")
|
|
else:
|
|
positive = text_embeds["prompt_embeds"]""")
|
|
|
|
# Normalize whitespace for matching: the source may have different
|
|
# indentation (8 spaces inside the while loop).
|
|
def _normalize_indent(block):
|
|
"""Strip common leading whitespace so we can match flexibly."""
|
|
lines = block.splitlines()
|
|
return [l.strip() for l in lines if l.strip()]
|
|
|
|
original_lines = _normalize_indent(original_block)
|
|
|
|
# Find the block in the source by matching stripped lines
|
|
src_lines = src.splitlines()
|
|
match_start = None
|
|
for i in range(len(src_lines)):
|
|
if src_lines[i].strip() == original_lines[0]:
|
|
# Check if subsequent lines match
|
|
matched = True
|
|
for j, orig_line in enumerate(original_lines):
|
|
if i + j >= len(src_lines):
|
|
matched = False
|
|
break
|
|
if src_lines[i + j].strip() != orig_line:
|
|
matched = False
|
|
break
|
|
if matched:
|
|
match_start = i
|
|
break
|
|
|
|
if match_start is None:
|
|
log.warning(
|
|
"[MultiTalkPromptSchedule] Could not find the original prompt "
|
|
"selection block in multitalk_loop source. The upstream code may "
|
|
"have changed. Patch not applied."
|
|
)
|
|
return
|
|
|
|
# Determine the indentation of the matched block
|
|
indent = ""
|
|
first_line = src_lines[match_start]
|
|
indent = first_line[: len(first_line) - len(first_line.lstrip())]
|
|
|
|
# Build the replacement with proper indentation
|
|
patched_lines = []
|
|
for line in patched_block.splitlines():
|
|
if line.strip():
|
|
patched_lines.append(indent + line)
|
|
else:
|
|
patched_lines.append("")
|
|
|
|
# Replace in source
|
|
match_end = match_start + len(original_lines)
|
|
new_src_lines = (
|
|
src_lines[:match_start] + patched_lines + src_lines[match_end:]
|
|
)
|
|
|
|
# Patch 2: deepcopy audio_embedding to prevent cache corruption.
|
|
# The original `audio_embedding = multitalk_audio_embeds` is just a
|
|
# reference. When audio is short, the loop pads audio_embedding[i]
|
|
# in-place (torch.cat) which corrupts the cached node output.
|
|
# On the next run, embeddings are already padded → wrong sizes → frozen video.
|
|
audio_orig = "audio_embedding = multitalk_audio_embeds"
|
|
audio_patched = (
|
|
"# [Patched by ComfyUI-MultiTalkPromptSchedule] deepcopy to prevent cache corruption\n"
|
|
"audio_embedding = [t.clone() for t in multitalk_audio_embeds]"
|
|
)
|
|
patched_count = 0
|
|
for i, line in enumerate(new_src_lines):
|
|
if audio_orig in line:
|
|
indent2 = line[: len(line) - len(line.lstrip())]
|
|
new_src_lines[i] = "\n".join(indent2 + l for l in audio_patched.splitlines())
|
|
patched_count += 1
|
|
if patched_count > 0:
|
|
log.info(
|
|
"[MultiTalkPromptSchedule] Also patched audio_embedding "
|
|
"to deepcopy (prevents cache corruption on re-run)."
|
|
)
|
|
else:
|
|
log.warning(
|
|
"[MultiTalkPromptSchedule] Could not find "
|
|
"'audio_embedding = multitalk_audio_embeds' to patch."
|
|
)
|
|
|
|
new_src = "\n".join(new_src_lines)
|
|
|
|
# Remove the `def multitalk_loop(self, **kwargs):` indentation
|
|
# since we'll exec it as a module-level function.
|
|
# Actually, we need to dedent the function definition first.
|
|
new_src = textwrap.dedent(new_src)
|
|
|
|
# Compile and exec in the module's namespace
|
|
try:
|
|
code = compile(new_src, loop_mod.__file__, "exec")
|
|
exec(code, loop_mod.__dict__)
|
|
log.info(
|
|
"[MultiTalkPromptSchedule] Successfully patched multitalk_loop "
|
|
"with prompt_schedule support."
|
|
)
|
|
_PATCHED = True
|
|
|
|
# CRITICAL: nodes_sampler.py imports multitalk_loop with
|
|
# `from .multitalk.multitalk_loop import multitalk_loop`
|
|
# which creates a separate reference. We must update it too,
|
|
# otherwise the sampler keeps calling the original function.
|
|
#
|
|
# We scan ALL loaded modules because ComfyUI may register them
|
|
# under unexpected names in sys.modules (hyphen vs underscore, etc.)
|
|
import sys
|
|
patched_fn = loop_mod.multitalk_loop
|
|
patched_count = 0
|
|
for mod_key, mod in list(sys.modules.items()):
|
|
if mod is None or mod is loop_mod:
|
|
continue
|
|
try:
|
|
if hasattr(mod, 'multitalk_loop') and callable(getattr(mod, 'multitalk_loop')):
|
|
old_fn = getattr(mod, 'multitalk_loop')
|
|
if old_fn is not patched_fn:
|
|
setattr(mod, 'multitalk_loop', patched_fn)
|
|
log.info(
|
|
f"[MultiTalkPromptSchedule] Patched multitalk_loop "
|
|
f"reference in module: {mod_key}"
|
|
)
|
|
patched_count += 1
|
|
except Exception:
|
|
continue
|
|
if patched_count == 0:
|
|
log.warning(
|
|
"[MultiTalkPromptSchedule] No other modules with multitalk_loop "
|
|
"found. The sampler may not be loaded yet — registering an "
|
|
"import hook to patch it when it loads."
|
|
)
|
|
# Deferred patch: if nodes_sampler loads AFTER us, patch it then
|
|
_deferred_patch_fn = patched_fn
|
|
_deferred_loop_mod = loop_mod
|
|
|
|
class _MultitalkImportHook:
|
|
def find_module(self, name, path=None):
|
|
if 'nodes_sampler' in name:
|
|
return self
|
|
return None
|
|
|
|
def load_module(self, name):
|
|
# Let the normal import proceed
|
|
sys.meta_path.remove(self)
|
|
mod = importlib.import_module(name)
|
|
# Now patch it
|
|
if hasattr(mod, 'multitalk_loop'):
|
|
mod.multitalk_loop = _deferred_patch_fn
|
|
log.info(
|
|
f"[MultiTalkPromptSchedule] Deferred patch applied "
|
|
f"to {name}"
|
|
)
|
|
sys.meta_path.append(self)
|
|
return mod
|
|
|
|
sys.meta_path.append(_MultitalkImportHook())
|
|
else:
|
|
log.info(
|
|
f"[MultiTalkPromptSchedule] Patched {patched_count} module(s) "
|
|
f"with updated multitalk_loop reference."
|
|
)
|
|
|
|
except Exception as e:
|
|
log.error(
|
|
f"[MultiTalkPromptSchedule] Failed to compile patched "
|
|
f"multitalk_loop: {e}. Patch not applied."
|
|
)
|