digital-garbageandClaude Sonnet 5.5
5718a0d7e3
Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
144 lines
8.9 KiB
Python
144 lines
8.9 KiB
Python
"""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)
|