diff --git a/app/ui/features/composer/enhance.js b/app/ui/features/composer/enhance.js index 6bec88d4..9c4be5c9 100644 --- a/app/ui/features/composer/enhance.js +++ b/app/ui/features/composer/enhance.js @@ -9,7 +9,7 @@ const DIALS = [["temperature", 0.7, 0, 2, 0.05, "Higher gives more varied wordin export function enhance(app, own) { const ps = app.pipeline, page = c.region.stack({ gap: "sm" }); - let defaults = null, last = null, alive = true; + let defaults = null, last = null, alive = true, spec = null; own(() => { alive = false; }); const slot = () => (ps.slots() || []).find((s) => s.node === NODE) || null; const set = async (name, value) => { const sl = slot(); if (sl) { await ps.save({ inputs: { [sl.id]: { [name]: value } } }); draw(); } }; @@ -22,6 +22,12 @@ export function enhance(app, own) { return [c.label.section({ text: "Last run" }), c.hint.default({ text: last.status || "" }), part("Before", last.before), c.field.default({ label: "After", control: c.text.sm({ text: last.after || "" }) }), part("Model's thinking", last.thinking), part("Sent to the model (chat run)", last.sent)].filter(Boolean); } + const choices = (name) => (((spec && spec.widgets) || []).find((w) => w.name === name) || {}).choices || []; + const pick = (name, label, first, hint) => { + const cur = (slot().inputs || {})[name] || "", list = choices(name); + const opts = [...(first ? [{ value: "", label: first }] : []), ...[...list, ...(cur && !list.includes(cur) ? [cur] : [])].map((v) => ({ value: v, label: String(v) }))]; + return field(label, c.select.md({ label, value: cur, options: opts, onChange: (v) => set(name, v) }), hint); + }; function draw() { if (!alive) return; const sl = slot(); @@ -33,6 +39,8 @@ export function enhance(app, own) { c.hint.default({ text: "A language model rewrites the finished prompt before it is encoded. Off, or if it fails, your prompt is used as typed. (word:1.5) and [a@2-4] markup does not survive a rewrite." }), c.toggle.default({ label: "Enhance prompt", checked: on, onChange: (x) => set("enabled", x) }), ...(!on ? readout() : [ + pick("encoder_file", "Encoder", "The model's own encoder", "The text encoder that writes the rewrite. Models like H3's own encoder cannot write prose: pick a chat model such as Gemma 4 here."), + ...((val("encoder_file", "") && choices("encoder_family").length) ? [pick("encoder_family", "Encoder family", "", "Which family that file is for.")] : []), field("Instructions", c.textarea.md({ label: "Instructions", rows: 6, value: val("instructions", ""), placeholder: defaults ? defaults.instructions : "(built-in instructions)", onCommit: (x) => set("instructions", x) }), "Empty uses the built-in instructions."), num("Max tokens", val("max_length", 400), 32, 4096, 8, "max_length", "The longest the rewrite may be."), c.toggle.default({ label: "Use the picture", hint: "Sends the scene's picture to the model along with the prompt.", checked: val("use_image", true), onChange: (x) => set("use_image", x) }), @@ -49,6 +57,7 @@ export function enhance(app, own) { } draw(); ps.ensureLoaded().then(draw); + app.api.describeNodes([NODE]).then((r) => { spec = (r.nodes || {})[NODE] || null; draw(); }).catch(() => {}); // Another project's pipeline, or an edit made elsewhere: redraw, but never under a field being typed in. const typing = () => page.node.contains(document.activeElement); own(ps.subscribe(() => { if (!typing()) draw(); })); diff --git a/app/ui/features/composer/enhance.test.js b/app/ui/features/composer/enhance.test.js index b98ae319..033aa093 100644 --- a/app/ui/features/composer/enhance.test.js +++ b/app/ui/features/composer/enhance.test.js @@ -9,7 +9,7 @@ test.after(() => teardownDom()); test("the Enhance tab follows the open project's pipeline: no enhancer there, no toggle shown", async () => { let live = [{ id: "e", node: "FunPackEnhancePrompt", inputs: { enabled: true } }], heard; const ps = { slots: () => live, ensureLoaded: async () => {}, save: async () => {}, subscribe: (fn) => { heard = fn; return () => {}; } }; - const api = { enhancerDefaults: async () => ({}), enhancerRuns: async () => ({ runs: [] }) }; + const api = { describeNodes: async () => ({ nodes: {} }), enhancerDefaults: async () => ({}), enhancerRuns: async () => ({ runs: [] }) }; const page = enhance({ pipeline: ps, api }, () => {}); document.body.append(page.node); assert.match(page.node.textContent, /Enhance prompt/); @@ -21,7 +21,7 @@ test("the Enhance tab follows the open project's pipeline: no enhancer there, no test("a run that ends while Enhance is open shows its final prompt without reopening the tab", async () => { const ps = { slots: () => [{ id: "e", node: "FunPackEnhancePrompt", inputs: {} }], ensureLoaded: async () => {}, save: async () => {}, subscribe: () => () => {} }; let runs = [], tell; - const api = { enhancerDefaults: async () => ({}), enhancerRuns: async () => ({ runs }) }; + const api = { describeNodes: async () => ({ nodes: {} }), enhancerDefaults: async () => ({}), enhancerRuns: async () => ({ runs }) }; const generate = { subscribe: (fn) => { tell = fn; fn({ phase: "idle" }); return () => {}; } }; const page = enhance({ pipeline: ps, api, generate }, () => {}); await new Promise((r) => setTimeout(r, 0)); @@ -44,3 +44,23 @@ test("Chat on a pipeline with no enhancer says so instead of taking comments tha heard(slots, "edit"); assert.ok(page.node.querySelector("textarea")); }); + +test("the Encoder pick lists the encoder files, and its family shows only once a file is picked", async () => { + let inputs = { enabled: true }, saved; + const ps = { slots: () => [{ id: "e", node: "FunPackEnhancePrompt", inputs }], ensureLoaded: async () => {}, save: async (b) => { saved = b; inputs = { ...inputs, ...b.inputs.e }; }, subscribe: () => () => {} }; + const widgets = [{ name: "encoder_file", type: "COMBO", choices: ["", "gemma4.safetensors"] }, { name: "encoder_family", type: "COMBO", choices: ["stable_diffusion", "ltxv"] }]; + const api = { describeNodes: async () => ({ nodes: { FunPackEnhancePrompt: { widgets } } }), enhancerDefaults: async () => ({}), enhancerRuns: async () => ({ runs: [] }) }; + const page = enhance({ pipeline: ps, api }, () => {}); + document.body.append(page.node); + await new Promise((r) => setTimeout(r, 0)); + assert.match(page.node.textContent, /Encoder/); + assert.match(page.node.textContent, /gemma4\.safetensors/); + assert.doesNotMatch(page.node.textContent, /Encoder family/); + inputs = { enabled: true, encoder_file: "gemma4.safetensors" }; + page.node.dispatchEvent(new Event("noop")); + const ps2 = { ...ps, slots: () => [{ id: "e", node: "FunPackEnhancePrompt", inputs }] }; + const again = enhance({ pipeline: ps2, api }, () => {}); + document.body.append(again.node); + await new Promise((r) => setTimeout(r, 0)); + assert.match(again.node.textContent, /Encoder family/); +}); diff --git a/core/tests/test_shipped_nodes.py b/core/tests/test_shipped_nodes.py index a4f1a057..c0683a81 100644 --- a/core/tests/test_shipped_nodes.py +++ b/core/tests/test_shipped_nodes.py @@ -60,7 +60,7 @@ FROZEN = { "FunPackShotCamera": ["text", "settings", "length", "frame_rate", "project"], "FunPackEnhancePrompt": [ "clip", "text", "enabled", "instructions", "max_length", "greedy", "temperature", "top_p", - "top_k", "min_p", "repetition_penalty", "presence_penalty", "seed", "thinking", "image", "use_image", + "top_k", "min_p", "repetition_penalty", "presence_penalty", "seed", "thinking", "encoder_file", "encoder_family", "image", "use_image", "shortcuts", "lorebooks", "reference_intro", "chat", ], "FunPackApplyPromptMarkup": ["model", "clip", "positive", "latent", "markup"], diff --git a/modules/conditioning/prompt_enhancer/__init__.py b/modules/conditioning/prompt_enhancer/__init__.py index eb3661d7..f22d32b7 100644 --- a/modules/conditioning/prompt_enhancer/__init__.py +++ b/modules/conditioning/prompt_enhancer/__init__.py @@ -12,7 +12,7 @@ project is never rewritten as one blob and re-split. Whether a given text encoder writes sensible prose is an open question per model (H3's conditioning checkpoint has no trained decoder head; Gemma3/Qwen3 do). v4's -result: confirmed working on Gemma3-12B. v5 on a rental: works with Gemma 4 E4B (a real rewrite in 2.5 s); H3's own encoder writes no prose, so there the run keeps your prompt and says why. On H3, add a second FunPack CLIP Loader (Models & Pipeline), pick the Gemma file, and wire it to the enhance node's clip: verified end to end. +result: confirmed working on Gemma3-12B. v5 on a rental: works with Gemma 4 E4B (a real rewrite in 2.5 s); H3's own encoder writes no prose, so there the run keeps your prompt and says why. On H3, pick a chat encoder (Gemma 4) in Composer > Enhance > Encoder: verified end to end. """ from .nodes import RUNS, FunPackEnhancePrompt diff --git a/modules/conditioning/prompt_enhancer/nodes.py b/modules/conditioning/prompt_enhancer/nodes.py index 5075c7f9..3a1e0aad 100644 --- a/modules/conditioning/prompt_enhancer/nodes.py +++ b/modules/conditioning/prompt_enhancer/nodes.py @@ -8,6 +8,7 @@ import time from comfy_api.latest import io from ..._core import log +from ...loaders.clip import nodes as clip_nodes from . import enhance as en #: The last runs' rewrites, newest last: what the Composer's readout shows. In @@ -18,6 +19,17 @@ 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): + _LOADED.clear() + _LOADED.update(key=(file, family), clip=clip_nodes.FunPackCLIPLoader.execute(clip_name1=file, type=family).args[0]) + return _LOADED["clip"] + + def _prompt_id(): try: from server import PromptServer @@ -57,6 +69,10 @@ class FunPackEnhancePrompt(io.ComfyNode): 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."), + io.Combo.Input("encoder_file", options=[""] + clip_nodes.encoder_files(), default="", optional=True, + tooltip="A text encoder file just for rewriting (Gemma, Qwen...). Empty uses the model's own encoder."), + io.Combo.Input("encoder_family", options=clip_nodes.clip_types(), 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."), @@ -82,7 +98,7 @@ class FunPackEnhancePrompt(io.ComfyNode): @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, image=None, use_image=True, shortcuts="", + presence_penalty=0.0, seed=0, thinking=False, encoder_file="", encoder_family="stable_diffusion", image=None, use_image=True, shortcuts="", lorebooks="", reference_intro="", chat="") -> io.NodeOutput: if not enabled: # still the final prompt, and the Composer's Enhance tab shows what was encoded either way @@ -93,6 +109,14 @@ class FunPackEnhancePrompt(io.ComfyNode): log.once("prompt_enhancer:markup", log.ALERT, "FunPack Prompt enhancer", "the prompt has (word:1.5) / [a@2-4] / [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("FunPack 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: diff --git a/modules/conditioning/prompt_enhancer/tests/test_prompt_enhancer.py b/modules/conditioning/prompt_enhancer/tests/test_prompt_enhancer.py index c59f3802..3cd88ba7 100644 --- a/modules/conditioning/prompt_enhancer/tests/test_prompt_enhancer.py +++ b/modules/conditioning/prompt_enhancer/tests/test_prompt_enhancer.py @@ -260,3 +260,13 @@ def test_a_tokenizer_that_reads_the_picture_out_of_kwargs_gets_it(): clip = Clip("A fox runs.") _, status, _ = en().enhance(clip, "fox", image=object()) assert "takes no picture" not in status and "image" in clip.seen + + +def test_a_picked_encoder_that_will_not_load_keeps_the_prompt_and_says_why(monkeypatch): + from modules.conditioning.prompt_enhancer import nodes + def boom(file, family): + raise RuntimeError("no such family") + monkeypatch.setattr(nodes, "_own_encoder", boom) + out = nodes.FunPackEnhancePrompt.execute(clip=object(), text="a fox", enabled=True, encoder_file="x.gguf") + assert out.args[0] == "a fox" and out.args[1] == "encoder failed" + assert "x.gguf" in nodes.RUNS[-1]["status"] and not nodes.RUNS[-1]["ok"]