"""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.""" 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 FunPack 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("FunPack 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"(? 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"|"": 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: """`...` (also an unterminated one), or everything before a lone `` when the template opened the block itself.""" text = str(raw or "") parts = re.findall(r"(?is)(.*?)(?:|$)", text) if not parts and "" in text.lower(): parts = [re.split(r"(?i)", 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).*?", "", text) text = re.sub(r"(?is).*", "", 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, True 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(), True return head.rsplit(" ", 1)[0].strip(), True # --- 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="", **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 SYSTEM_PROMPT 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: 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)") 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