"""The Enhance Prompt node: an LLM rewrites the finished prompt before it is encoded.""" import collections import json import re import time from comfy_api.latest import io from ..._core import log from . import enhance as en #: The last runs' rewrites, newest last: what the Composer's readout shows. In #: memory on purpose -- it is the last run's answer, not state to keep. RUNS = collections.deque(maxlen=20) _MARKUP = re.compile(r"\([^()]*:\s*[\d.]+\)|\[[^\]]*\|[^\]]*\]|⟦[^⟧]*⟧") _LOADED = {} # one entry: the last separate encoder, so a second Generate does not load a 9 GB file again def _own_encoder(file, family): """The encoder the person picked for the enhancer, loaded the way the pipeline's own loader loads one.""" if _LOADED.get("key") != (file, family): from ...loaders.clip import nodes as clip_nodes # here, not at the top: the enhancer must load even when the CLIP loader module did not _LOADED.clear() _LOADED.update(key=(file, family), clip=clip_nodes.FunUICLIPLoader.execute(clip_name1=file, type=family).args[0]) return _LOADED["clip"] def _prompt_id(): try: from server import PromptServer return str(PromptServer.instance.last_prompt_id or "") except Exception: # noqa: BLE001 return "" class FunUIEnhancePrompt(io.ComfyNode): @classmethod def define_schema(cls) -> io.Schema: return io.Schema( node_id="FunUIEnhancePrompt", display_name="FunUI Enhance Prompt", category="FunUI/Conditioning", description="Rewrite the finished prompt with a language model before it is encoded. " "Off, or on any failure, the prompt passes through unchanged.", inputs=[ io.Clip.Input("clip", tooltip="A text encoder that can generate text (Gemma, Qwen...)."), io.String.Input("text", multiline=True, default=""), io.Boolean.Input("enabled", default=False, tooltip="Off: the prompt passes through and the model is not run."), io.String.Input("instructions", multiline=True, default="", optional=True, tooltip="What the model is told to do. Empty uses the built-in instructions."), io.Int.Input("max_length", default=400, min=32, max=4096, step=8, optional=True, tooltip="The longest rewrite, in tokens."), io.Boolean.Input("greedy", default=False, optional=True, tooltip="Always the likeliest word. Hides the sampling dials below."), io.Float.Input("temperature", default=0.7, min=0.0, max=2.0, step=0.05, optional=True), io.Float.Input("top_p", default=0.92, min=0.0, max=1.0, step=0.01, optional=True), io.Int.Input("top_k", default=50, min=1, max=1000, optional=True, tooltip="top_p and min_p only apply while top_k is above 0."), io.Float.Input("min_p", default=0.05, min=0.0, max=1.0, step=0.01, optional=True), io.Float.Input("repetition_penalty", default=1.3, min=1.0, max=2.0, step=0.05, optional=True), io.Float.Input("presence_penalty", default=0.0, min=0.0, max=2.0, step=0.05, optional=True), io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff, optional=True, tooltip="0 = a fresh rewrite every run. Any other number repeats the same rewrite."), io.Boolean.Input("thinking", default=False, optional=True, tooltip="Let a thinking model reason first. Its reasoning eats the token budget."), # Text, not a list: a file that has gone must not stop a run where Enhance is off. The Enhance panel offers the files. io.String.Input("encoder_file", default="", optional=True, tooltip="A text encoder file just for rewriting (Gemma, Qwen...). Empty uses the model's own encoder."), io.String.Input("encoder_family", default="stable_diffusion", optional=True, tooltip="Which family that file is for. Only used with an encoder file."), io.Image.Input("image", optional=True, tooltip="Shown to the model alongside the prompt."), io.Boolean.Input("use_image", default=True, optional=True, tooltip="Send the picture (when one is connected) to the model."), io.String.Input("shortcuts", multiline=True, default="", optional=True, tooltip="Shortcut names to hand the model as reference, one per line."), io.String.Input("lorebooks", multiline=True, default="", optional=True, tooltip="Paths to lorebook or shortcuts JSON files, one per line."), io.String.Input("reference_intro", multiline=True, default="", optional=True, tooltip="The line above the reference entries. Empty uses the built-in one."), io.String.Input("chat", multiline=True, default="", optional=True, tooltip="Your comments on earlier rewrites (Composer > Chat)."), io.Combo.Input("style", options=list(en.STYLES), default="default", optional=True, tooltip="default: the built-in instructions. dynamic: adds a concrete action, a change over time and secondary motion to every shot. Typed Instructions override both."), ], outputs=[io.String.Output(display_name="enhanced"), io.String.Output(display_name="status")], ) @classmethod def fingerprint_inputs(cls, seed=0, enabled=False, **_): # Seed 0 means "a new rewrite every run", and ComfyUI would otherwise hand # back the cached one for identical inputs. return time.time() if enabled and not seed else "" @classmethod def execute(cls, clip, text, enabled=False, instructions="", max_length=400, greedy=False, temperature=0.7, top_p=0.92, top_k=50, min_p=0.05, repetition_penalty=1.3, presence_penalty=0.0, seed=0, thinking=False, encoder_file="", encoder_family="stable_diffusion", image=None, use_image=True, shortcuts="", lorebooks="", reference_intro="", chat="", style="default") -> io.NodeOutput: if not enabled: # still the final prompt, and the Composer's Enhance tab shows what was encoded either way RUNS.append({"prompt_id": _prompt_id(), "status": "off: the prompt ran as typed", "before": text, "after": text, "thinking": "", "sent": "", "ok": True}) return io.NodeOutput(text, "off") if _MARKUP.search(text or ""): log.once("prompt_enhancer:markup", log.ALERT, "FunUI Prompt enhancer", "the prompt has (word:1.5) / [a|b] markup, which the model will " "read as plain text: with the enhancer on, markup does not survive a rewrite") if encoder_file: try: clip = _own_encoder(encoder_file, encoder_family) except Exception as e: # noqa: BLE001 -- the prompt must still go through log.warning("FunUI Prompt enhancer", f"could not load {encoder_file}: {e}; the original prompt is used") RUNS.append({"prompt_id": _prompt_id(), "status": f"could not load the encoder {encoder_file}: {e}", "before": text, "after": text, "thinking": "", "sent": "", "ok": False}) return io.NodeOutput(text, "encoder failed") lines = lambda s: [x for x in str(s or "").splitlines() if x.strip()] # noqa: E731 groups = en.sources(lines(shortcuts), lines(lorebooks)) try: rounds = json.loads(chat) if str(chat or "").strip() else [] except ValueError: rounds = [] log.warning("FunUI Prompt enhancer", "the Chat comments are not readable; ignored") if not isinstance(rounds, list): rounds = [] out, status, info = en.enhance( clip, text, system=instructions, reference_text=en.reference(text, groups, reference_intro), chat=en.chat_block(rounds), style=style, seed=seed, image=image if use_image else None, thinking=thinking, max_length=max_length, temperature=temperature, top_p=top_p, top_k=top_k, min_p=min_p, repetition_penalty=repetition_penalty, presence_penalty=presence_penalty, do_sample=not greedy) if not info["ok"]: # The prompt went through UNCHANGED and the person must know why. log.warning("FunUI Prompt enhancer", f"{status}; the original prompt is used") RUNS.append({"prompt_id": _prompt_id(), "status": status, **info}) return io.NodeOutput(out, status)