405 lines
21 KiB
Python
405 lines
21 KiB
Python
"""The prompt enhancer's logic: what the model is asked, and how its answer is
|
|
cleaned. Pure text in, text out -- the node (nodes.py) only wires a CLIP to it.
|
|
|
|
Every failure path returns the ORIGINAL prompt. An enhancer that empties the
|
|
prompt would render a blank video and blame the model.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import inspect
|
|
import json
|
|
import random
|
|
import re
|
|
import time
|
|
|
|
from ..._core import log, shortcuts as shortcuts_mod
|
|
|
|
SYSTEM_PROMPT = """You expand short video prompts into detailed ones for a text-to-video model.
|
|
|
|
Rules:
|
|
- Keep every element the user asked for. Never drop, replace or contradict one.
|
|
- Add concrete visual detail where the prompt is vague: lighting, materials, textures, clothing, setting, time of day.
|
|
- Describe motion in the present progressive ("is walking", "is turning").
|
|
- Add a matching soundscape: ambient sound, and any sound the described actions would make.
|
|
- If the user quoted speech, reproduce it exactly. Never invent speech that was not asked for.
|
|
- Do not invent camera moves or scene cuts unless the user asked for them.
|
|
- Do not add characters the user did not mention.
|
|
|
|
Output: one paragraph of plain prose. No preamble, no headings, no markdown, no quotes around the whole answer. Output only the prompt itself."""
|
|
|
|
DYNAMIC_PROMPT = """You expand short video prompts into detailed, DYNAMIC ones for a text-to-video model.
|
|
|
|
Rules:
|
|
- Keep every element the user asked for. Never drop, replace or contradict one.
|
|
- Every shot gets at least one concrete action beat of its subject: a verb with intensity, speed or direction ("lunges forward", "spins on one heel", "slams the door shut"), not a state ("stands", "sits", "looks", "holds", "wears", "is").
|
|
- Show a change over time inside the shot: something starts, builds, peaks or ends.
|
|
- Add physically specific secondary motion where it fits: hair, cloth, water, dust, smoke, leaves, light.
|
|
- Write all motion in the present progressive ("is sprinting", "is turning").
|
|
- Add concrete visual detail where the prompt is vague: lighting, materials, setting, time of day.
|
|
- Add a matching soundscape: ambient sound, and any sound the described actions would make.
|
|
- If the user quoted speech, reproduce it exactly. Never invent speech that was not asked for.
|
|
- Do not invent camera moves or scene cuts unless the user asked for them.
|
|
- Do not add characters the user did not mention.
|
|
|
|
Output: one paragraph of plain prose. No preamble, no headings, no markdown, no quotes around the whole answer. Output only the prompt itself."""
|
|
|
|
#: what the Style choice selects; "default" must stay SYSTEM_PROMPT so an old workflow rewrites as before
|
|
STYLES = {"default": SYSTEM_PROMPT, "dynamic": DYNAMIC_PROMPT}
|
|
|
|
CHAT_SYSTEM = """REVISION: this time the user is giving feedback on your earlier rewrite.
|
|
The message has labelled parts: ORIGINAL PROMPT (what the user wrote), YOUR LATEST REWRITE (your last answer, if shown) and USER FEEDBACK (their comments, oldest first).
|
|
- Start from YOUR LATEST REWRITE (or the ORIGINAL PROMPT if none is shown) and change it as the feedback asks. Apply every comment; when two disagree, the later one wins.
|
|
- Keep everything the feedback does not ask to change.
|
|
- The feedback is instructions to you, not prompt text: never quote it, mention it or append it.
|
|
Output only the revised prompt, nothing else."""
|
|
|
|
REFERENCE_INTRO = ("Reference entries. Use an entry only where the prompt "
|
|
"calls for it, in your own words; ignore the rest.")
|
|
|
|
CHAT_LABELS = ("ORIGINAL PROMPT:", "YOUR LATEST REWRITE:", "USER FEEDBACK")
|
|
CHARS_PER_TOKEN = 4 # turns the user's TOKEN limit into a character bound
|
|
REPEAT_LIMIT = 3 # how often a sentence may repeat before it is a stuck model
|
|
|
|
# Filler that says nothing about which shortcut a prompt means.
|
|
_STOPWORDS = frozenset("""
|
|
the and with from into onto over under for but not are was were been being has have had
|
|
its his her hers their theirs them they she him you your yours our ours this that these
|
|
those there here then than very just only also some any all each every both few more most
|
|
other such own same too can will would could should may might must does did doing done
|
|
while when where which who whom whose what why how about above after again against
|
|
before below between during out off once further through until upon via per
|
|
""".split())
|
|
|
|
|
|
def vocab(text) -> set[str]:
|
|
"""Meaningful words (letters only, 3+ long, not filler), lowercased."""
|
|
return {w for w in re.findall(r"[^\W\d_]{3,}", str(text or "").lower()) if w not in _STOPWORDS}
|
|
|
|
|
|
# --- reference ---------------------------------------------------------------
|
|
|
|
def sources(shortcut_names=None, paths=None, say=None) -> list[list[dict]]:
|
|
"""Reference groups: the picked shortcuts as one group, then one per file (a
|
|
SillyTavern lorebook with `entries`, or a FunUI shortcuts file). Entry:
|
|
{block, words|keys, vocab, whole, constant}. A missing shortcut or an unreadable
|
|
file is skipped and said so -- through `say` when given."""
|
|
say = say or (lambda m: log.warning("FunUI Prompt enhancer", m))
|
|
|
|
def sc_entry(sc):
|
|
reps = [r for r in sc.replacements if r.strip()]
|
|
lines = [f"[Shortcut] {sc.name}"]
|
|
lines += reps if len(reps) == 1 else [f"Variant {i + 1}: {r}" for i, r in enumerate(reps)]
|
|
# triggers are expanded away by now, so the content counts as a mention too
|
|
return {"block": "\n".join(lines), "words": [sc.name, *sc.triggers, *reps],
|
|
"vocab": vocab(" ".join([sc.name, *reps]))}
|
|
|
|
groups = []
|
|
library = {s.name.lower(): s for s in shortcuts_mod.listing()}
|
|
picked = []
|
|
for name in (n.strip() for n in (shortcut_names or [])):
|
|
if not name:
|
|
continue
|
|
if name.lower() in library:
|
|
picked.append(sc_entry(library[name.lower()]))
|
|
else:
|
|
say(f"reference shortcut {name!r} is not in the library; skipped")
|
|
if picked:
|
|
groups.append(picked)
|
|
for path in (str(p).strip() for p in (paths or [])):
|
|
if not path:
|
|
continue
|
|
try:
|
|
with open(path, "r", encoding="utf-8") as fh:
|
|
data = json.load(fh)
|
|
if not isinstance(data, dict):
|
|
raise ValueError("not a lorebook or shortcuts file")
|
|
except (OSError, ValueError) as exc:
|
|
say(f"reference file {path} skipped: {exc}")
|
|
continue
|
|
if "shortcuts" in data:
|
|
raw = data["shortcuts"]
|
|
raw = list(raw.values()) if isinstance(raw, dict) else raw
|
|
items = [shortcuts_mod.Shortcut.from_dict(r) for r in raw if isinstance(r, dict)]
|
|
groups.append([sc_entry(s) for s in items if s.triggers and s.enabled])
|
|
continue
|
|
entries = data.get("entries", [])
|
|
entries = list(entries.values()) if isinstance(entries, dict) else entries
|
|
entries = sorted((e for e in entries if isinstance(e, dict) and str(e.get("content") or "").strip()),
|
|
key=lambda e: e.get("insertion_order", e.get("order", 0)) or 0)
|
|
groups.append([{"block": f"[Lore] {e.get('comment') or e.get('name') or 'entry'}\n"
|
|
f"{str(e['content']).strip()}",
|
|
"keys": e.get("keys", e.get("key", [])),
|
|
"whole": e.get("matchWholeWords") is not False, # ST's default is whole words
|
|
"constant": bool(e.get("constant"))} for e in entries])
|
|
return [g for g in groups if g]
|
|
|
|
|
|
def reference(text, groups, intro="") -> str:
|
|
"""Appended after the prompt the model rewrites. Per group: the entries the prompt
|
|
mentions, or the whole group when it mentions none (the model may need what the
|
|
prompt does not name). Shortcuts are mentioned by name, trigger, content or any
|
|
meaningful word of those; lore by its keywords, a constant entry joining any match."""
|
|
low = str(text or "").lower()
|
|
|
|
def said(word, whole=True, lore=False):
|
|
word = str(word).strip()
|
|
if not word:
|
|
return False
|
|
if lore and len(word) > 2 and word.startswith("/") and word.endswith("/"):
|
|
try:
|
|
return re.search(word[1:-1], text or "", re.IGNORECASE) is not None
|
|
except re.error:
|
|
return False
|
|
if not whole:
|
|
return word.lower() in low
|
|
return re.search(r"(?<!\w)" + re.escape(word.lower()) + r"(?!\w)", low) is not None
|
|
|
|
typed = vocab(low)
|
|
|
|
def mentioned(e):
|
|
if e.get("vocab") and typed & e["vocab"]:
|
|
return True
|
|
if "keys" in e:
|
|
keys = e["keys"].split(",") if isinstance(e["keys"], str) else e["keys"]
|
|
return any(said(k, e.get("whole", True), lore=True) for k in keys or [])
|
|
return any(said(w) for w in e["words"])
|
|
|
|
blocks = []
|
|
for group in groups or []:
|
|
hit = [e for e in group if mentioned(e)]
|
|
if hit:
|
|
hit = [e for e in group if e.get("constant") or e in hit]
|
|
blocks += [e["block"] for e in (hit or group)]
|
|
if not blocks:
|
|
return ""
|
|
return "\n\n" + (str(intro or "").strip() or REFERENCE_INTRO) + "\n\n" + "\n\n".join(blocks)
|
|
|
|
|
|
# --- chat --------------------------------------------------------------------
|
|
|
|
def chat_block(chat, scene=None) -> str:
|
|
"""The user's comments on earlier rewrites as the labelled part of the message:
|
|
the LATEST rewrite of this prompt and every comment since Reset, oldest first.
|
|
`chat` is a list of rounds {"rewrites": {"whole"|"<scene index>": text}, "comment"}.
|
|
Only the latest rewrite is shown (a model shown several picks one to repeat);
|
|
every comment stays, because the user expects all of them to still hold."""
|
|
rounds = [r for r in (chat or []) if isinstance(r, dict) and str(r.get("comment") or "").strip()]
|
|
if not rounds:
|
|
return ""
|
|
key = "whole" if scene is None else str(scene)
|
|
latest = ""
|
|
for r in rounds:
|
|
rewrites = r.get("rewrites") if isinstance(r.get("rewrites"), dict) else {}
|
|
prev = rewrites.get(key)
|
|
if prev is None and len(rewrites) == 1:
|
|
prev = next(iter(rewrites.values()))
|
|
if str(prev or "").strip():
|
|
latest = str(prev).strip()
|
|
parts = []
|
|
if latest:
|
|
parts.append("YOUR LATEST REWRITE:\n" + latest)
|
|
parts.append("USER FEEDBACK (oldest first):\n" + "\n".join("- " + str(r["comment"]).strip() for r in rounds))
|
|
return "\n\n" + "\n\n".join(parts)
|
|
|
|
|
|
def cut_chat_echo(text) -> str:
|
|
"""A model that treats the conversation as more prompt text repeats its labels:
|
|
a leading label is removed, and everything from the first later label is cut."""
|
|
text = str(text or "").strip()
|
|
labels = "|".join(re.escape(l.rstrip(":")) for l in CHAT_LABELS)
|
|
text = re.sub(rf"(?i)^\s*(?:{labels})[^:\n]*:\s*", "", text)
|
|
hit = re.search(rf"(?i)(?:{labels})", text)
|
|
return (text[:hit.start()] if hit else text).strip()
|
|
|
|
|
|
# --- cleaning ----------------------------------------------------------------
|
|
|
|
def extract_thinking(raw) -> str:
|
|
"""`<think>...</think>` (also an unterminated one), or everything before a lone
|
|
`</think>` when the template opened the block itself."""
|
|
text = str(raw or "")
|
|
parts = re.findall(r"(?is)<think>(.*?)(?:</think>|$)", text)
|
|
if not parts and "</think>" in text.lower():
|
|
parts = [re.split(r"(?i)</think>", text, maxsplit=1)[0]]
|
|
return "\n\n".join(p.strip() for p in parts if p.strip())
|
|
|
|
|
|
def clean(raw) -> str:
|
|
"""Strip what a chat model wraps around a prompt it was asked to output bare:
|
|
thinking traces, code fences, "Here is..." preambles, whole-answer quotes.
|
|
Quotes INSIDE the prompt survive -- speech is quoted on purpose."""
|
|
text = str(raw or "")
|
|
text = re.sub(r"(?is).*?</think>", "", text)
|
|
text = re.sub(r"(?is)<think>.*", "", text)
|
|
text = re.sub(r"```[a-zA-Z]*\n?", "", text).replace("```", "")
|
|
text = re.sub(r"(?i)^\s*(here(?:'s| is)[^:\n]*:|output\s*:|prompt\s*:)\s*", "", text.strip())
|
|
text = text.strip().strip("`").strip()
|
|
if len(text) > 1 and text[0] in "\"'“" and text[-1] in "\"'”":
|
|
text = text[1:-1].strip()
|
|
return text.strip()
|
|
|
|
|
|
def trim_runaway(text, max_length=None):
|
|
"""Cut a prompt that never ended -- by REPETITION first, length second. A model
|
|
with no trained decoder head may never emit a stop token. Length alone cannot tell
|
|
"long and detailed" from "stuck", so what is detected is a sentence repeating; the
|
|
length bound comes from the user's own token limit. Cuts land on a sentence."""
|
|
text = str(text or "").strip()
|
|
if not text:
|
|
return text, False
|
|
sentences = re.split(r"(?<=[.!?])\s+", text)
|
|
seen: dict[str, int] = {}
|
|
for i, s in enumerate(sentences):
|
|
key = " ".join(s.lower().split())
|
|
if len(key) < 12:
|
|
continue
|
|
seen[key] = seen.get(key, 0) + 1
|
|
if seen[key] >= REPEAT_LIMIT:
|
|
kept = " ".join(sentences[:i]).strip()
|
|
if kept:
|
|
return kept, "repeat"
|
|
if max_length is None:
|
|
return text, False
|
|
cap = max(64, int(max_length) * CHARS_PER_TOKEN)
|
|
if len(text) <= cap:
|
|
return text, False
|
|
head = text[:cap]
|
|
cut = max(head.rfind(". "), head.rfind("! "), head.rfind("? "))
|
|
if cut > cap // 4:
|
|
return head[:cut + 1].strip(), "length"
|
|
return head.rsplit(" ", 1)[0].strip(), "length"
|
|
|
|
|
|
# --- generation --------------------------------------------------------------
|
|
|
|
def generation_device(clip) -> str:
|
|
"""Where the text generation actually ran -- reported, not assumed. A 32B encoder
|
|
sharing a card with the DiT can end up on the CPU: minutes per call, invisible
|
|
from the text alone."""
|
|
for path in (("patcher", "load_device"), ("_device",), ("device",)):
|
|
probe = clip
|
|
for name in path:
|
|
probe = getattr(probe, name, None)
|
|
if probe is None:
|
|
break
|
|
if probe is not None:
|
|
return str(probe)
|
|
return "an unreported device"
|
|
|
|
|
|
def _accepted(fn, swallows=False):
|
|
"""Names `fn` takes, or None when it cannot be read or takes anything. A **kwargs
|
|
catch-all counts as taking anything, unless `swallows`: a tokenizer's **kwargs ignores
|
|
what it is given -- except the names its body reads back out of it (Qwen3-VL reads
|
|
`kwargs.get("image")`), which count."""
|
|
try:
|
|
params = inspect.signature(fn).parameters
|
|
except (TypeError, ValueError):
|
|
return None
|
|
names = set(params)
|
|
if any(p.kind is p.VAR_KEYWORD for p in params.values()):
|
|
if not swallows:
|
|
return None
|
|
try:
|
|
names |= set(re.findall(r"""\bkwargs\.get\(\s*["'](\w+)["']""", inspect.getsource(fn)))
|
|
names |= set(re.findall(r"""\bkwargs\[\s*["'](\w+)["']\s*\]""", inspect.getsource(fn)))
|
|
except (OSError, TypeError):
|
|
pass
|
|
return names
|
|
|
|
|
|
def generate(clip, system, user, *, seed=None, image=None, thinking=False, max_length=400,
|
|
temperature=0.7, top_p=0.92, top_k=50, min_p=0.05, repetition_penalty=1.3,
|
|
presence_penalty=0.0, do_sample=True):
|
|
"""(raw text, status). Through a CLIP that exposes ComfyUI's generate/decode pair.
|
|
No ComfyUI tokenizer takes a system prompt (they swallow unknown kwargs), so the
|
|
instructions go first in the user turn, inside the model's own chat template --
|
|
which is also what places an image marker."""
|
|
if clip is None or not hasattr(clip, "generate") or not hasattr(clip, "decode"):
|
|
return "", ("unavailable: the connected CLIP does not expose text generation; wire a "
|
|
"generation-capable text encoder")
|
|
max_length = max(32, int(max_length or 400))
|
|
started = time.time()
|
|
notes = []
|
|
try:
|
|
merged = f"{system}\n\n{user}" if str(system).strip() else str(user)
|
|
outer = getattr(clip, "tokenizer", None)
|
|
# The stock SD1Tokenizer forwards its kwargs to the per-encoder tokenizer, which then
|
|
# decides; a model that overrides it (MiniMax H3) is the one that decides.
|
|
stock = getattr(getattr(outer, "tokenize_with_weights", None), "__qualname__", "").startswith("SD1Tokenizer.")
|
|
inner = (getattr(outer, getattr(outer, "clip", None) or "", None) if stock else None) or outer
|
|
takes = _accepted(getattr(inner, "tokenize_with_weights", None), swallows=True)
|
|
# Only what this encoder's tokenizer really takes: the others swallow unknown kwargs
|
|
# silently, so a picture or chat template would be dropped while the status said "generated".
|
|
tk = {"min_length": 1} # Gemma pads to 1024 otherwise
|
|
for name, value in (("skip_template", False), ("thinking", bool(thinking))):
|
|
if takes is None or name in takes:
|
|
tk[name] = value
|
|
if image is not None:
|
|
if takes is None or "image" in takes:
|
|
tk["image"] = image
|
|
else:
|
|
notes.append("this text encoder takes no picture, so only the text was sent")
|
|
if takes is not None and "skip_template" not in takes:
|
|
notes.append("this text encoder has no chat template, so the instructions went in as raw text")
|
|
try:
|
|
tokens = clip.tokenize(merged, **tk)
|
|
except TypeError: # a CLIP wrapper with a fixed signature
|
|
tokens = clip.tokenize(merged, **({"image": image} if image is not None else {}))
|
|
sampling = bool(do_sample) and float(temperature) > 0 # temperature 0 would divide by zero
|
|
kwargs = dict(do_sample=sampling, max_length=max_length, temperature=max(float(temperature), 0.01),
|
|
top_k=int(top_k), top_p=float(top_p), min_p=float(min_p),
|
|
repetition_penalty=float(repetition_penalty), no_repeat_ngram_size=5,
|
|
presence_penalty=float(presence_penalty),
|
|
# ComfyUI seeds a torch Generator with this when sampling: None raises
|
|
seed=int(seed) if seed else random.randrange(1, 2 ** 31))
|
|
accepted = _accepted(getattr(clip, "generate", None))
|
|
if accepted is not None:
|
|
kwargs = {k: v for k, v in kwargs.items() if k in accepted}
|
|
ids = clip.generate(tokens, **kwargs)
|
|
text = str(clip.decode(ids, skip_special_tokens=True) or "").strip()
|
|
return text, (f"generated {len(text)} chars in {time.time() - started:.1f}s on "
|
|
f"{generation_device(clip)} (cap {max_length} tokens)"
|
|
+ "".join(f"; {n}" for n in notes))
|
|
except Exception as exc: # noqa: BLE001
|
|
return "", f"generation failed: {exc}"
|
|
|
|
|
|
def enhance(clip, text, *, system="", reference_text="", chat="", style="default", **sampling):
|
|
"""-> (text, status, info). On any failure the ORIGINAL text comes back."""
|
|
original = str(text or "").strip()
|
|
info = {"before": original, "after": original, "thinking": "", "sent": "", "ok": False}
|
|
if not original:
|
|
return text, "skipped: the prompt is empty", info
|
|
system = str(system or "").strip() or STYLES.get(style, SYSTEM_PROMPT) # typed instructions win over the style
|
|
user = original + str(reference_text or "")
|
|
if chat:
|
|
# Labelled parts plus a revision rule in the SYSTEM turn: with the conversation merely
|
|
# appended, a model rewrote it as more prompt text and echoed it back.
|
|
system += "\n\n" + CHAT_SYSTEM
|
|
user = "ORIGINAL PROMPT:\n" + user + str(chat)
|
|
info["sent"] = f"SYSTEM:\n{system}\n\nMESSAGE:\n{user}"
|
|
raw, status = generate(clip, system, user, **sampling)
|
|
info["thinking"] = extract_thinking(raw)
|
|
cleaned = clean(raw)
|
|
if chat:
|
|
cleaned = cut_chat_echo(cleaned)
|
|
out, overran = trim_runaway(cleaned, max_length=sampling.get("max_length"))
|
|
if overran == "repeat":
|
|
status += ("; the model repeated itself instead of stopping and was cut at a sentence "
|
|
"(a model with a trained chat head does not do this)")
|
|
elif overran:
|
|
status += "; the answer reached the Max length setting and was cut there"
|
|
if not out:
|
|
return original, (status if "failed" in status or "unavailable" in status
|
|
else "the model returned nothing; prompt unchanged"), info
|
|
visible = "".join(out.split())
|
|
if not re.search(r"[^\W\d_]{2,}", out) or sum(c.isalpha() for c in visible) < len(visible) // 4:
|
|
# Seen on a real GPU: an encoder that is not a chat model answered with 400 commas, and that WAS sent
|
|
# on to the video model, which then made something unrelated to the prompt.
|
|
return original, ("the model wrote no usable words (" + out[:20].replace("\n", " ")
|
|
+ "...); prompt unchanged. This text encoder may not be a chat model"), info
|
|
info.update(after=out, ok=True)
|
|
return out, f"{status}; prompt {len(original)} -> {len(out)} chars", info
|