nodes/Limbicnation-ComfyUI-PromptGenerator
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Merge pull request #12 from Limbicnation/feature/dual-stream-prompt-refiner

Browse Source 
Add dual-stream prompt refiner node
This commit is contained in:
Gero Doll
2026-06-10 22:32:08 +02:00
committed by GitHub
parent 70ba1bf8b7 eab2efc461
commit 8a56ed45a6
6 changed files with 404 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+4
View File
@@ -6,6 +6,7 @@ Generate Stable Diffusion prompts using Qwen3-8B via Ollama
try:
from .nodes.negative_prompt_node import NegativePromptNode
from .nodes.prompt_combiner_node import PromptCombinerNode
from .nodes.prompt_dual_stream_refiner_node import PromptDualStreamRefinerNode
from .nodes.prompt_generator_node import PromptGeneratorNode
from .nodes.prompt_refiner_node import PromptRefinerNode
from .nodes.style_applier_node import StyleApplierNode
@@ -13,6 +14,7 @@ except ImportError:
# Fallback for test environments where ComfyUI isn't present
from nodes.negative_prompt_node import NegativePromptNode
from nodes.prompt_combiner_node import PromptCombinerNode
from nodes.prompt_dual_stream_refiner_node import PromptDualStreamRefinerNode
from nodes.prompt_generator_node import PromptGeneratorNode
from nodes.prompt_refiner_node import PromptRefinerNode
from nodes.style_applier_node import StyleApplierNode
@@ -21,6 +23,7 @@ NODE_CLASS_MAPPINGS = {
"Limbicnation_PromptGenerator": PromptGeneratorNode,
"Limbicnation_StyleApplier": StyleApplierNode,
"Limbicnation_PromptRefiner": PromptRefinerNode,
"Limbicnation_PromptDualStreamRefiner": PromptDualStreamRefinerNode,
"Limbicnation_NegativePrompt": NegativePromptNode,
"Limbicnation_PromptCombiner": PromptCombinerNode,
}
@@ -29,6 +32,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
"Limbicnation_PromptGenerator": "Prompt Generator (Qwen)",
"Limbicnation_StyleApplier": "Style Applier",
"Limbicnation_PromptRefiner": "Prompt Refiner",
"Limbicnation_PromptDualStreamRefiner": "Prompt Dual-Stream Refiner",
"Limbicnation_NegativePrompt": "Negative Prompt Generator",
"Limbicnation_PromptCombiner": "Prompt Combiner",
}
config/Modelfile.dualstream
+39
View File
@@ -0,0 +1,39 @@
# Dual-Stream Prompt Refiner Modelfile
#
# Registers the pre-merged Q8 GGUF from
# Limbicnation/qwen2-5-7b-dual-stream-prompt-lora with Ollama so the
# "Prompt Dual-Stream Refiner" ComfyUI node can use it.
#
# Usage:
# 1. Download qwen2-5-7b-dual-stream-q8.gguf from the HF repo into this dir:
# huggingface-cli download Limbicnation/qwen2-5-7b-dual-stream-prompt-lora \
# qwen2-5-7b-dual-stream-q8.gguf --local-dir .
# 2. Create the Ollama model (name contains "prompt" so it sorts to the top
# of the node's model dropdown — see OllamaClient.LORA_KEYWORDS):
# ollama create limbicnation-dualstream-prompt -f config/Modelfile.dualstream
# 3. Verify the output format before relying on the node's parser:
# ollama run limbicnation-dualstream-prompt "a mystical forest at twilight"
# Confirm it emits "Positive: ...\nNegative: ...". If the trained format
# differs, adjust INSTRUCTION_PROMPT / parse_dual_stream in
# nodes/prompt_dual_stream_refiner_node.py to match.
# 4. Restart ComfyUI — the model appears in the node's dropdown.
# Pre-merged Q8 GGUF (already contains the LoRA — no ADAPTER needed).
FROM ./qwen2-5-7b-dual-stream-q8.gguf
# Base is Qwen2.5-7B-Instruct, which uses the standard ChatML template.
# (Ollama applies a sensible Qwen2.5 default; uncomment to pin explicitly.)
# TEMPLATE """{{ if .System }}<|im_start|>system
# {{ .System }}<|im_end|>
# {{ end }}<|im_start|>user
# {{ .Prompt }}<|im_end|>
# <|im_start|>assistant
# """
SYSTEM """You are an expert Stable Diffusion prompt engineer. Given a description, you produce a detailed positive prompt and a matching negative prompt, formatted as:
Positive: <positive prompt>
Negative: <negative prompt>"""
PARAMETER temperature 0.7
PARAMETER top_p 0.9
PARAMETER num_ctx 4096
docs/plans/2026-06-10-qwen2-5-dual-stream-refiner.md
+61
View File
@@ -0,0 +1,61 @@
# Qwen2.5 Dual-Stream Prompt Refiner — Implementation (Ollama/GGUF)
> **Superseded note:** An earlier draft of this plan loaded the model locally via
> `transformers` + PEFT. That approach was **invalid** — the HF repo
> `Limbicnation/qwen2-5-7b-dual-stream-prompt-lora` is a LoRA *adapter* with **no base
> weights**, so `AutoModelForCausalLM.from_pretrained(...)` on it crashes. It also added a
> heavy ML stack to an Ollama-only project and loaded a 7B model into ComfyUI's diffusion
> VRAM. This document replaces it with the implemented Ollama/GGUF approach.
**Goal:** A ComfyUI node that turns a description into a positive + negative prompt pair,
powered by the pre-merged Q8 GGUF the HF repo already ships, served through Ollama.
**Repo facts:** PEFT adapter (`adapter_config.json` + `adapter_model.safetensors`) **plus**
`qwen2-5-7b-dual-stream-q8.gguf` (8.1 GB). Base: `Qwen/Qwen2.5-7B-Instruct`. "Dual-stream" =
a single causal LM trained to emit positive + negative blocks (no special architecture). The
shipped `chat_template.jinja` is the stock Qwen2.5 ChatML template and does **not** define the
positive/negative format — so the node imposes the format via its instruction prompt and parses
defensively.
## What was implemented
1. **`nodes/prompt_dual_stream_refiner_node.py`** — `PromptDualStreamRefinerNode`
(`Limbicnation_PromptDualStreamRefiner`, display "Prompt Dual-Stream Refiner",
category `text/generation`). Inputs: `prompt`, `model` (dropdown from
`OllamaClient.discover_models()`), optional `temperature`/`top_p`/`seed`/`timeout`, hidden
`unique_id`. Outputs: `(positive_prompt, negative_prompt)`. Reuses `OllamaClient` for
streaming, progress bar, timeout, llama-runner crash categorization, and subprocess fallback.
Module-level `parse_dual_stream()` strips thinking/markdown via `extract_final_prompt` then
splits on a delimiter-tolerant `Negative:` label.
2. **`config/Modelfile.dualstream`** — registers the GGUF with Ollama. Suggested model name
`limbicnation-dualstream-prompt` (contains `prompt`, a `LORA_KEYWORDS` token, so it sorts to
the top of the dropdown).
3. **`__init__.py`** — node registered in both `NODE_CLASS_MAPPINGS` and
`NODE_DISPLAY_NAME_MAPPINGS`.
4. **Tests** — `tests/unit/test_prompt_dual_stream_refiner.py` (parser variants + refine flow);
`tests/unit/test_node_registration.py` `EXPECTED_NODES` updated.
## No new runtime dependencies
`requirements.txt` stays `ollama` / `jinja2` / `pyyaml`.
## Setup & verification
```bash
# 1. Get the GGUF into Ollama
huggingface-cli download Limbicnation/qwen2-5-7b-dual-stream-prompt-lora \
qwen2-5-7b-dual-stream-q8.gguf --local-dir config
ollama create limbicnation-dualstream-prompt -f config/Modelfile.dualstream
# 2. Confirm the trained output format BEFORE trusting the parser
ollama run limbicnation-dualstream-prompt "a mystical forest at twilight"
# -> expect "Positive: ...\nNegative: ...". If different, tune
# INSTRUCTION_PROMPT / parse_dual_stream to match.
# 3. Code checks
ruff check . && ruff format --check .
python -m pytest tests/unit/test_prompt_dual_stream_refiner.py tests/unit/test_node_registration.py -q
# 4. Manual: restart ComfyUI, add "Prompt Dual-Stream Refiner" (text/generation),
# select the model, run a description, confirm two correct outputs + progress bar.
```
nodes/prompt_dual_stream_refiner_node.py
+209
View File
@@ -0,0 +1,209 @@
"""
Dual-Stream Prompt Refiner Node for ComfyUI.
Refines a raw description into a *pair* of prompts — a positive prompt and a
negative prompt — in a single generation pass, served via Ollama.
Designed for the `Limbicnation/qwen2-5-7b-dual-stream-prompt-lora` model after
its shipped Q8 GGUF has been registered with Ollama (see
config/Modelfile.dualstream), but works with any chat-capable Ollama model: the
"Positive:" / "Negative:" output format is imposed by the instruction prompt
rather than relying on the model's training, and parsing is delimiter-tolerant.
"""
import logging
import re
from typing import Any
from .adapters.ollama_client import OllamaClient
from .prompt_generator_node import extract_final_prompt
logger = logging.getLogger(__name__)
# Strip a leading "Positive[ prompt]:" label, tolerating markdown bold (**) both
# before the label and after the colon. e.g. "**Positive:** text".
# A colon is required (not a hyphen) so art terms like "negative-space" in the
# prompt body are not mistaken for a label.
_POS_LABEL = re.compile(r"^\s*\**\s*positive(?:\s+prompt)?\s*\**\s*:\s*\**\s*", re.IGNORECASE)
# Locate the negative section. Anchored to a line start (the imposed output
# format puts "Negative:" on its own line) so the word "negative" appearing
# mid-sentence in the positive prompt (e.g. "negative space") never triggers a split.
_NEG_LABEL = re.compile(r"(?:^|\n)\s*\**\s*negative(?:\s+prompt)?\s*\**\s*:\s*\**\s*", re.IGNORECASE)
def parse_dual_stream(text: str) -> tuple[str, str]:
"""Split a model response into (positive_prompt, negative_prompt).
Thinking/markdown noise is stripped first via ``extract_final_prompt`` so a
"Negative:" mention inside a reasoning block can't cause a mis-split. The
text is then split at the first "Negative:" label; everything before it is
the positive prompt (with any leading "Positive:" label removed). If no
negative label is found, the whole response is treated as the positive
prompt and the negative is empty.
"""
cleaned = (extract_final_prompt(text) or text).strip()
neg_match = _NEG_LABEL.search(cleaned)
if neg_match:
positive = cleaned[: neg_match.start()]
negative = cleaned[neg_match.end() :]
else:
positive = cleaned
negative = ""
positive = _POS_LABEL.sub("", positive, count=1).strip().strip('"').strip()
negative = negative.strip().strip('"').strip()
return positive, negative
class PromptDualStreamRefinerNode:
"""ComfyUI node that turns a description into a positive + negative prompt pair."""
INSTRUCTION_PROMPT = """You are an expert Stable Diffusion prompt engineer.
From the description below, produce TWO prompts:
1. A detailed positive prompt — subject, style, lighting, composition, color, and quality tags.
2. A negative prompt — artifacts and qualities to avoid (e.g. blurry, lowres, deformed, watermark).
Respond in EXACTLY this format, with no explanations or markdown:
Positive: <positive prompt>
Negative: <negative prompt>
Description: {prompt}"""
@classmethod
def _get_available_models(cls) -> list[str]:
"""Fetch available Ollama models via OllamaClient (LoRA/prompt models first)."""
client = OllamaClient(logger_prefix="PromptDualStreamRefiner")
return client.discover_models()
@classmethod
def INPUT_TYPES(cls) -> dict[str, Any]:
available_models = cls._get_available_models()
return {
"required": {
"prompt": (
"STRING",
{
"multiline": True,
"default": "",
"placeholder": "Description to expand into positive + negative prompts...",
},
),
"model": (
available_models,
{
"default": available_models[0] if available_models else "qwen3:8b",
"tooltip": "Select Ollama model. Dual-stream/LoRA models appear first.",
},
),
},
"optional": {
"temperature": (
"FLOAT",
{
"default": 0.7,
"min": 0.1,
"max": 1.0,
"step": 0.1,
"display": "slider",
},
),
"top_p": (
"FLOAT",
{
"default": 0.9,
"min": 0.1,
"max": 1.0,
"step": 0.1,
"display": "slider",
},
),
"seed": (
"INT",
{
"default": -1,
"min": -1,
"max": 2**31 - 1,
"step": 1,
},
),
"timeout": (
"INT",
{
"default": 120,
"min": 30,
"max": 600,
"step": 10,
},
),
},
"hidden": {"unique_id": "UNIQUE_ID"},
}
RETURN_TYPES = ("STRING", "STRING")
RETURN_NAMES = ("positive_prompt", "negative_prompt")
FUNCTION = "refine"
CATEGORY = "text/generation"
OUTPUT_NODE = False
def refine(
self,
prompt: str,
model: str,
temperature: float = 0.7,
top_p: float = 0.9,
seed: int = -1,
timeout: int = 120,
unique_id: str | None = None,
) -> tuple[str, str]:
"""Generate a positive/negative prompt pair from a raw description.
Returns:
(positive_prompt, negative_prompt). On failure the error message is
returned as the positive output and the negative is empty, matching
the error-surfacing convention of PromptRefinerNode.
"""
if not prompt.strip():
return ("[PromptDualStreamRefiner] Please provide a description.", "")
client = OllamaClient(logger_prefix="PromptDualStreamRefiner")
pbar = client.create_progress_bar(unique_id)
effective_seed: int | None = None if seed == -1 else seed
instruction = self.INSTRUCTION_PROMPT.format(prompt=prompt.strip())
logger.info("Dual-stream refine with model='%s'", model)
result = client.generate_streaming(
model=model,
prompt=instruction,
temperature=temperature,
top_p=top_p,
timeout=timeout,
pbar=pbar,
seed=effective_seed,
)
if result.kind == "ok" and result.text is not None:
output = result.text
elif result.kind in ("model_crash", "server_error", "unavailable"):
# Subprocess fallback won't help for these; surface directly.
return (f"[PromptDualStreamRefiner] {result.message}", "")
else:
# timeout / transient — try the CLI subprocess fallback.
success, output = client.generate_subprocess(model, instruction, timeout)
if not success:
return (f"[PromptDualStreamRefiner] {output}", "")
positive, negative = parse_dual_stream(output)
if pbar is not None:
pbar.update_absolute(100)
if not positive and not negative:
logger.warning("Dual-stream parse produced empty output")
return ("[PromptDualStreamRefiner] Model returned no usable prompt.", "")
logger.info("Dual-stream complete: +%d / -%d chars", len(positive), len(negative))
return (positive, negative)
tests/unit/test_node_registration.py
+1
View File
@@ -31,6 +31,7 @@ EXPECTED_NODES: frozenset[str] = frozenset(
"Limbicnation_PromptGenerator",
"Limbicnation_StyleApplier",
"Limbicnation_PromptRefiner",
"Limbicnation_PromptDualStreamRefiner",
"Limbicnation_NegativePrompt",
"Limbicnation_PromptCombiner",
}
tests/unit/test_prompt_dual_stream_refiner.py
+90
View File
@@ -0,0 +1,90 @@
"""Unit tests for the dual-stream refiner: output parsing and refine() flow."""
from unittest.mock import patch
from nodes.adapters.ollama_client import StreamResult
from nodes.prompt_dual_stream_refiner_node import (
PromptDualStreamRefinerNode,
parse_dual_stream,
)
class TestParseDualStream:
"""parse_dual_stream must tolerate label/delimiter variants."""
def test_plain_positive_negative(self):
pos, neg = parse_dual_stream("Positive: a cat\nNegative: blurry, lowres")
assert pos == "a cat"
assert neg == "blurry, lowres"
def test_prompt_suffix_labels(self):
pos, neg = parse_dual_stream("Positive prompt: a dog\nNegative prompt: deformed")
assert pos == "a dog"
assert neg == "deformed"
def test_markdown_bold_labels(self):
pos, neg = parse_dual_stream("**Positive:** a fox\n**Negative:** ugly")
assert pos == "a fox"
assert neg == "ugly"
def test_case_insensitive(self):
pos, neg = parse_dual_stream("POSITIVE: bright\nNEGATIVE: dark")
assert pos == "bright"
assert neg == "dark"
def test_no_negative_section(self):
pos, neg = parse_dual_stream("Positive: just a positive prompt")
assert pos == "just a positive prompt"
assert neg == ""
def test_unlabeled_text_is_positive(self):
pos, neg = parse_dual_stream("a forest at twilight")
assert pos == "a forest at twilight"
assert neg == ""
def test_multiline_blocks_preserved(self):
text = "Positive: line one,\nline two\nNegative: bad, worse"
pos, neg = parse_dual_stream(text)
assert "line one" in pos and "line two" in pos
assert neg == "bad, worse"
def test_negative_space_in_body_is_not_a_split(self):
"""The art term 'negative space' must not be mistaken for a label."""
pos, neg = parse_dual_stream("a portrait with strong negative space, studio lighting")
assert pos == "a portrait with strong negative space, studio lighting"
assert neg == ""
def test_hyphenated_negative_space_in_body_is_not_a_split(self):
pos, neg = parse_dual_stream("Positive: negative-space composition\nNegative: blurry")
assert pos == "negative-space composition"
assert neg == "blurry"
class TestRefineFlow:
"""refine() should parse a successful stream into two outputs."""
def test_successful_generation_splits_streams(self):
node = PromptDualStreamRefinerNode()
with patch(
"nodes.prompt_dual_stream_refiner_node.OllamaClient.generate_streaming",
return_value=StreamResult(text="Positive: a knight\nNegative: blurry", kind="ok"),
):
pos, neg = node.refine(prompt="a knight", model="qwen3:8b")
assert pos == "a knight"
assert neg == "blurry"
def test_empty_prompt_short_circuits(self):
node = PromptDualStreamRefinerNode()
pos, neg = node.refine(prompt=" ", model="qwen3:8b")
assert pos.startswith("[PromptDualStreamRefiner]")
assert neg == ""
def test_model_crash_surfaces_error(self):
node = PromptDualStreamRefinerNode()
with patch(
"nodes.prompt_dual_stream_refiner_node.OllamaClient.generate_streaming",
return_value=StreamResult(kind="model_crash", message="runner crashed"),
):
pos, neg = node.refine(prompt="a knight", model="broken")
assert "runner crashed" in pos
assert neg == ""