feat: add PromptCombiner, PromptRefiner, NegativePrompt nodes
- Fix StyleApplier relative imports (from ..style_presets) - Add PromptCombiner: blend/concat/weighted_average modes with emphasis markers - Add PromptRefiner: iterative LLM refinement (1-3 passes) via Ollama - Add NegativePrompt: auto (LLM) and preset (category-based) negative generation - Update __init__.py to register all 5 nodes
This commit is contained in:
limbicnation
@@ -0,0 +1,29 @@
|
||||
"""
|
||||
ComfyUI Prompt Generator Node
|
||||
Generate Stable Diffusion prompts using Qwen3-8B via Ollama
|
||||
"""
|
||||
|
||||
from .nodes.prompt_generator_node import PromptGeneratorNode
|
||||
from .nodes.style_applier_node import StyleApplierNode
|
||||
from .nodes.prompt_combiner_node import PromptCombinerNode
|
||||
from .nodes.prompt_refiner_node import PromptRefinerNode
|
||||
from .nodes.negative_prompt_node import NegativePromptNode
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"PromptGenerator": PromptGeneratorNode,
|
||||
"StyleApplier": StyleApplierNode,
|
||||
"PromptCombiner": PromptCombinerNode,
|
||||
"PromptRefiner": PromptRefinerNode,
|
||||
"NegativePrompt": NegativePromptNode,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"PromptGenerator": "🎨 Prompt Generator (Qwen)",
|
||||
"StyleApplier": "🎬 Style Applier (Cinematic/Photo)",
|
||||
"PromptCombiner": "🔗 Prompt Combiner",
|
||||
"PromptRefiner": "✨ Prompt Refiner",
|
||||
"NegativePrompt": "⛔ Negative Prompt Generator",
|
||||
}
|
||||
|
||||
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|
||||
|
||||
@@ -0,0 +1,391 @@
|
||||
"""
|
||||
Negative Prompt Generator Node for ComfyUI
|
||||
Generates negative prompts from a positive prompt or description.
|
||||
"""
|
||||
|
||||
import subprocess
|
||||
import time
|
||||
import threading
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
from .prompt_generator_node import extract_final_prompt
|
||||
|
||||
try:
|
||||
import ollama
|
||||
OLLAMA_API_AVAILABLE = True
|
||||
except ImportError:
|
||||
OLLAMA_API_AVAILABLE = False
|
||||
|
||||
try:
|
||||
import comfy.utils
|
||||
COMFY_PROGRESS_AVAILABLE = True
|
||||
except ImportError:
|
||||
COMFY_PROGRESS_AVAILABLE = False
|
||||
|
||||
|
||||
# Default negative prompt categories for fallback
|
||||
DEFAULT_NEGATIVE_CATEGORIES: Dict[str, List[str]] = {
|
||||
"quality": [
|
||||
"low quality",
|
||||
"worst quality",
|
||||
"bad anatomy",
|
||||
"bad proportions",
|
||||
"blurry",
|
||||
"out of focus",
|
||||
"deformed",
|
||||
"disfigured",
|
||||
"extra limbs",
|
||||
"mutated",
|
||||
"poorly drawn",
|
||||
"ugly",
|
||||
],
|
||||
"artifacts": [
|
||||
"jpeg artifacts",
|
||||
"compression artifacts",
|
||||
"noise",
|
||||
"grainy",
|
||||
"pixelated",
|
||||
"oversaturated",
|
||||
"watermark",
|
||||
"signature",
|
||||
"text",
|
||||
"logo",
|
||||
"cropped",
|
||||
"out of frame",
|
||||
],
|
||||
"people": [
|
||||
"bad face",
|
||||
"asymmetric eyes",
|
||||
"crossed eyes",
|
||||
"missing fingers",
|
||||
"extra fingers",
|
||||
"fused fingers",
|
||||
"too many fingers",
|
||||
"malformed hands",
|
||||
"bad hands",
|
||||
"missing arms",
|
||||
"missing legs",
|
||||
"extra arms",
|
||||
"extra legs",
|
||||
],
|
||||
"style": [
|
||||
"cartoon",
|
||||
"anime",
|
||||
"3d render",
|
||||
"cgi",
|
||||
"plastic",
|
||||
"doll",
|
||||
"painting",
|
||||
"sketch",
|
||||
"drawing",
|
||||
"illustration",
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
class NegativePromptNode:
|
||||
"""
|
||||
ComfyUI node for generating negative prompts.
|
||||
|
||||
Supports two modes:
|
||||
- auto: Uses an LLM (Ollama) to generate a context-aware negative prompt
|
||||
based on the positive prompt content.
|
||||
- preset: Combines predefined negative keyword categories.
|
||||
|
||||
Outputs:
|
||||
- negative_prompt: The generated negative prompt string
|
||||
- category_list: Comma-separated list of categories used (for reference)
|
||||
"""
|
||||
|
||||
SYSTEM_PROMPT = """You are an expert Stable Diffusion negative prompt engineer.
|
||||
|
||||
Given a positive prompt, generate a concise negative prompt that lists only
|
||||
what should be avoided. Focus on:
|
||||
- Quality issues (blurry, low quality, bad anatomy)
|
||||
- Unwanted style elements (if the prompt specifies photorealistic, avoid cartoon/anime)
|
||||
- Artifacts and technical problems
|
||||
- Content that contradicts the positive prompt
|
||||
|
||||
Rules:
|
||||
- Return ONLY the negative prompt text, comma-separated.
|
||||
- No explanations, no markdown, no bullet points.
|
||||
- Keep it under 200 tokens.
|
||||
- Do NOT include positive concepts.
|
||||
|
||||
Positive prompt: {prompt}
|
||||
|
||||
Negative prompt:"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls) -> dict[str, Any]:
|
||||
return {
|
||||
"required": {
|
||||
"positive_prompt": (
|
||||
"STRING",
|
||||
{
|
||||
"multiline": True,
|
||||
"default": "",
|
||||
"placeholder": "Enter the positive prompt to generate negatives for...",
|
||||
},
|
||||
),
|
||||
"mode": (
|
||||
["auto", "preset"],
|
||||
{"default": "preset"},
|
||||
),
|
||||
},
|
||||
"optional": {
|
||||
"model": (
|
||||
"STRING",
|
||||
{
|
||||
"default": "qwen3:8b",
|
||||
"placeholder": "Ollama model (auto mode only)",
|
||||
},
|
||||
),
|
||||
"categories": (
|
||||
"STRING",
|
||||
{
|
||||
"default": "quality,artifacts",
|
||||
"placeholder": "Comma-separated: quality,artifacts,people,style",
|
||||
},
|
||||
),
|
||||
"temperature": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 0.3,
|
||||
"min": 0.1,
|
||||
"max": 1.0,
|
||||
"step": 0.1,
|
||||
"display": "slider",
|
||||
},
|
||||
),
|
||||
"timeout": (
|
||||
"INT",
|
||||
{
|
||||
"default": 60,
|
||||
"min": 10,
|
||||
"max": 300,
|
||||
"step": 10,
|
||||
},
|
||||
),
|
||||
"custom_negatives": (
|
||||
"STRING",
|
||||
{
|
||||
"default": "",
|
||||
"placeholder": "Additional custom negative terms, comma-separated...",
|
||||
},
|
||||
),
|
||||
},
|
||||
"hidden": {
|
||||
"unique_id": "UNIQUE_ID",
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING", "STRING")
|
||||
RETURN_NAMES = ("negative_prompt", "category_list")
|
||||
FUNCTION = "generate"
|
||||
CATEGORY = "text/generation"
|
||||
OUTPUT_NODE = False
|
||||
|
||||
def _generate_streaming(
|
||||
self,
|
||||
model: str,
|
||||
prompt: str,
|
||||
temperature: float,
|
||||
timeout: int,
|
||||
pbar: object = None,
|
||||
) -> Optional[str]:
|
||||
"""Stream ollama.generate() with timeout enforcement."""
|
||||
chunks = []
|
||||
start = time.monotonic()
|
||||
first_chunk_timeout = min(timeout * 0.6, 45)
|
||||
chunk_timeout = 20
|
||||
got_first_chunk = False
|
||||
|
||||
try:
|
||||
stream = ollama.generate(
|
||||
model=model,
|
||||
prompt=prompt,
|
||||
stream=True,
|
||||
options={"temperature": temperature, "top_p": 0.9},
|
||||
)
|
||||
|
||||
result_holder: Dict[str, Any] = {}
|
||||
|
||||
def _iter_next(it):
|
||||
try:
|
||||
result_holder["chunk"] = next(it)
|
||||
result_holder["done"] = False
|
||||
except StopIteration:
|
||||
result_holder["done"] = True
|
||||
except Exception as exc:
|
||||
result_holder["error"] = exc
|
||||
|
||||
it = iter(stream)
|
||||
chunk_count = 0
|
||||
|
||||
while True:
|
||||
elapsed = time.monotonic() - start
|
||||
if elapsed >= timeout:
|
||||
print(f"[NegativePrompt] Total timeout ({timeout}s) reached")
|
||||
break
|
||||
|
||||
result_holder.clear()
|
||||
t = threading.Thread(target=_iter_next, args=(it,), daemon=True)
|
||||
t.start()
|
||||
|
||||
wait_time = first_chunk_timeout if not got_first_chunk else chunk_timeout
|
||||
wait_time = min(wait_time, timeout - elapsed)
|
||||
t.join(timeout=wait_time)
|
||||
|
||||
if t.is_alive():
|
||||
label = "first chunk" if not got_first_chunk else "chunk"
|
||||
print(f"[NegativePrompt] Timeout waiting for {label} ({wait_time:.0f}s)")
|
||||
return None
|
||||
|
||||
if "error" in result_holder:
|
||||
raise result_holder["error"]
|
||||
|
||||
if result_holder.get("done", False):
|
||||
break
|
||||
|
||||
chunk = result_holder.get("chunk")
|
||||
if chunk is None:
|
||||
break
|
||||
|
||||
text = chunk.get("response", "")
|
||||
if text:
|
||||
chunks.append(text)
|
||||
got_first_chunk = True
|
||||
chunk_count += 1
|
||||
if pbar is not None:
|
||||
progress = min(5 + int(chunk_count * 5), 95)
|
||||
pbar.update_absolute(progress)
|
||||
|
||||
except Exception as e:
|
||||
print(f"[NegativePrompt] Streaming error: {e}")
|
||||
return None
|
||||
|
||||
if not chunks:
|
||||
return None
|
||||
|
||||
return "".join(chunks)
|
||||
|
||||
def _build_preset_negative(
|
||||
self,
|
||||
categories_str: str,
|
||||
custom_negatives: str,
|
||||
) -> Tuple[str, str]:
|
||||
"""Build negative prompt from preset categories."""
|
||||
selected = [c.strip().lower() for c in categories_str.split(",") if c.strip()]
|
||||
terms: List[str] = []
|
||||
valid_categories: List[str] = []
|
||||
|
||||
for cat in selected:
|
||||
if cat in DEFAULT_NEGATIVE_CATEGORIES:
|
||||
terms.extend(DEFAULT_NEGATIVE_CATEGORIES[cat])
|
||||
valid_categories.append(cat)
|
||||
else:
|
||||
print(f"[NegativePrompt] Unknown category '{cat}', skipping")
|
||||
|
||||
if custom_negatives:
|
||||
custom_terms = [t.strip() for t in custom_negatives.split(",") if t.strip()]
|
||||
terms.extend(custom_terms)
|
||||
|
||||
if not terms:
|
||||
return ("", ",".join(valid_categories))
|
||||
|
||||
return (", ".join(terms), ",".join(valid_categories))
|
||||
|
||||
def generate(
|
||||
self,
|
||||
positive_prompt: str,
|
||||
mode: str,
|
||||
model: str = "qwen3:8b",
|
||||
categories: str = "quality,artifacts",
|
||||
temperature: float = 0.3,
|
||||
timeout: int = 60,
|
||||
custom_negatives: str = "",
|
||||
unique_id: str = None,
|
||||
) -> Tuple[str, str]:
|
||||
"""
|
||||
Generate a negative prompt.
|
||||
|
||||
Args:
|
||||
positive_prompt: The positive prompt to generate negatives for
|
||||
mode: "auto" (LLM-generated) or "preset" (keyword categories)
|
||||
model: Ollama model for auto mode
|
||||
categories: Comma-separated category names for preset mode
|
||||
temperature: Generation temperature for auto mode
|
||||
timeout: Max generation time for auto mode
|
||||
custom_negatives: Additional custom terms for preset mode
|
||||
unique_id: ComfyUI node ID for progress tracking
|
||||
|
||||
Returns:
|
||||
Tuple of (negative_prompt, category_list)
|
||||
"""
|
||||
if not positive_prompt.strip():
|
||||
return ("[NegativePrompt] Please provide a positive prompt.", "")
|
||||
|
||||
pbar = None
|
||||
if COMFY_PROGRESS_AVAILABLE and unique_id is not None:
|
||||
try:
|
||||
pbar = comfy.utils.ProgressBar(100, node_id=unique_id)
|
||||
pbar.update_absolute(0)
|
||||
except Exception:
|
||||
pbar = None
|
||||
|
||||
if mode == "preset":
|
||||
neg, cats = self._build_preset_negative(categories, custom_negatives)
|
||||
if pbar is not None:
|
||||
pbar.update_absolute(100)
|
||||
return (neg, cats)
|
||||
|
||||
# Auto mode: use Ollama
|
||||
system_prompt = self.SYSTEM_PROMPT.format(prompt=positive_prompt.strip())
|
||||
|
||||
if OLLAMA_API_AVAILABLE:
|
||||
output = self._generate_streaming(
|
||||
model=model,
|
||||
prompt=system_prompt,
|
||||
temperature=temperature,
|
||||
timeout=timeout,
|
||||
pbar=pbar,
|
||||
)
|
||||
|
||||
if output is not None:
|
||||
cleaned = extract_final_prompt(output.strip())
|
||||
if cleaned:
|
||||
if pbar is not None:
|
||||
pbar.update_absolute(100)
|
||||
return (cleaned, "auto")
|
||||
else:
|
||||
print("[NegativePrompt] Auto mode returned empty, falling back to preset")
|
||||
|
||||
# Fallback to subprocess or preset
|
||||
if OLLAMA_API_AVAILABLE:
|
||||
print("[NegativePrompt] Streaming failed, trying subprocess fallback")
|
||||
else:
|
||||
print("[NegativePrompt] Ollama API not available, using subprocess fallback")
|
||||
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["ollama", "run", model, system_prompt],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=timeout,
|
||||
)
|
||||
if result.returncode == 0 and result.stdout.strip():
|
||||
cleaned = extract_final_prompt(result.stdout.strip())
|
||||
if cleaned:
|
||||
if pbar is not None:
|
||||
pbar.update_absolute(100)
|
||||
return (cleaned, "auto")
|
||||
except (subprocess.TimeoutExpired, FileNotFoundError, Exception) as e:
|
||||
print(f"[NegativePrompt] Subprocess fallback failed: {e}")
|
||||
|
||||
# Final fallback: return preset negative
|
||||
print("[NegativePrompt] All auto methods failed, returning preset negative")
|
||||
neg, cats = self._build_preset_negative(categories, custom_negatives)
|
||||
if pbar is not None:
|
||||
pbar.update_absolute(100)
|
||||
return (neg, cats)
|
||||
@@ -0,0 +1,246 @@
|
||||
"""
|
||||
Prompt Combiner Node for ComfyUI
|
||||
Merges multiple prompt strings with configurable blending strategies.
|
||||
"""
|
||||
|
||||
from enum import Enum
|
||||
from typing import Any, List, Literal, Tuple, get_args
|
||||
|
||||
|
||||
class CombineMode(str, Enum):
|
||||
"""Supported strategies for combining multiple prompts."""
|
||||
|
||||
BLEND = "blend"
|
||||
CONCAT = "concat"
|
||||
WEIGHTED_AVERAGE = "weighted_average"
|
||||
|
||||
|
||||
# Single source of truth for the choices exposed in INPUT_TYPES and accepted by combine().
|
||||
ModeLiteral = Literal["blend", "concat", "weighted_average"]
|
||||
|
||||
|
||||
class PromptCombinerNode:
|
||||
"""
|
||||
ComfyUI node for combining multiple prompts into a single output.
|
||||
|
||||
Supports:
|
||||
- blend: Weighted combination with emphasis markers
|
||||
- concat: Simple concatenation with separator
|
||||
- weighted_average: Text interpolation based on weights
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls) -> dict[str, Any]:
|
||||
return {
|
||||
"required": {
|
||||
"prompt_1": (
|
||||
"STRING",
|
||||
{
|
||||
"multiline": True,
|
||||
"default": "",
|
||||
"placeholder": "First prompt...",
|
||||
},
|
||||
),
|
||||
"mode": (
|
||||
list(get_args(ModeLiteral)),
|
||||
{"default": CombineMode.BLEND.value},
|
||||
),
|
||||
},
|
||||
"optional": {
|
||||
"prompt_2": (
|
||||
"STRING",
|
||||
{
|
||||
"multiline": True,
|
||||
"default": "",
|
||||
"placeholder": "Second prompt (optional)...",
|
||||
},
|
||||
),
|
||||
"prompt_3": (
|
||||
"STRING",
|
||||
{
|
||||
"multiline": True,
|
||||
"default": "",
|
||||
"placeholder": "Third prompt (optional)...",
|
||||
},
|
||||
),
|
||||
"prompt_4": (
|
||||
"STRING",
|
||||
{
|
||||
"multiline": True,
|
||||
"default": "",
|
||||
"placeholder": "Fourth prompt (optional)...",
|
||||
},
|
||||
),
|
||||
"weight_1": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 1.0,
|
||||
"min": 0.0,
|
||||
"max": 2.0,
|
||||
"step": 0.1,
|
||||
"display": "slider",
|
||||
},
|
||||
),
|
||||
"weight_2": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 1.0,
|
||||
"min": 0.0,
|
||||
"max": 2.0,
|
||||
"step": 0.1,
|
||||
"display": "slider",
|
||||
},
|
||||
),
|
||||
"weight_3": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 1.0,
|
||||
"min": 0.0,
|
||||
"max": 2.0,
|
||||
"step": 0.1,
|
||||
"display": "slider",
|
||||
},
|
||||
),
|
||||
"weight_4": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 1.0,
|
||||
"min": 0.0,
|
||||
"max": 2.0,
|
||||
"step": 0.1,
|
||||
"display": "slider",
|
||||
},
|
||||
),
|
||||
"separator": (
|
||||
"STRING",
|
||||
{
|
||||
"default": ", ",
|
||||
"placeholder": "Separator for concat mode",
|
||||
},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING",)
|
||||
RETURN_NAMES = ("combined_prompt",)
|
||||
FUNCTION = "combine"
|
||||
CATEGORY = "text/generation"
|
||||
OUTPUT_NODE = False
|
||||
|
||||
def combine(
|
||||
self,
|
||||
prompt_1: str,
|
||||
mode: str,
|
||||
prompt_2: str = "",
|
||||
prompt_3: str = "",
|
||||
prompt_4: str = "",
|
||||
weight_1: float = 1.0,
|
||||
weight_2: float = 1.0,
|
||||
weight_3: float = 1.0,
|
||||
weight_4: float = 1.0,
|
||||
separator: str = ", ",
|
||||
) -> Tuple[str]:
|
||||
"""
|
||||
Combine multiple prompts using the selected mode.
|
||||
|
||||
Args:
|
||||
prompt_1: First prompt (required)
|
||||
mode: Combination strategy
|
||||
prompt_2-4: Additional prompts (optional)
|
||||
weight_1-4: Weights for each prompt
|
||||
separator: Separator string for concat mode
|
||||
|
||||
Returns:
|
||||
Tuple containing the combined prompt string
|
||||
"""
|
||||
prompts: List[Tuple[str, float]] = [
|
||||
(p.strip(), w)
|
||||
for p, w in (
|
||||
(prompt_1, weight_1),
|
||||
(prompt_2, weight_2),
|
||||
(prompt_3, weight_3),
|
||||
(prompt_4, weight_4),
|
||||
)
|
||||
if p and p.strip()
|
||||
]
|
||||
|
||||
if not prompts:
|
||||
return ("[PromptCombiner] At least one prompt is required.",)
|
||||
|
||||
try:
|
||||
selected = CombineMode(mode)
|
||||
except ValueError:
|
||||
valid = ", ".join(m.value for m in CombineMode)
|
||||
return (f"[PromptCombiner] Unknown mode {mode!r}. Valid: {valid}",)
|
||||
|
||||
if len(prompts) == 1:
|
||||
return (prompts[0][0],)
|
||||
|
||||
if selected == CombineMode.BLEND:
|
||||
return (self._blend(prompts),)
|
||||
elif selected == CombineMode.CONCAT:
|
||||
return (self._concat(prompts, separator),)
|
||||
elif selected == CombineMode.WEIGHTED_AVERAGE:
|
||||
return (self._weighted_average(prompts),)
|
||||
|
||||
# Fallback — should never reach here
|
||||
return ("[PromptCombiner] Internal error: unhandled mode.",)
|
||||
|
||||
def _blend(self, prompts: List[Tuple[str, float]]) -> str:
|
||||
"""
|
||||
Blend prompts using ComfyUI-style emphasis markers.
|
||||
Higher weight = more parentheses emphasis.
|
||||
"""
|
||||
parts = []
|
||||
for text, weight in prompts:
|
||||
if weight <= 0:
|
||||
continue
|
||||
# Map weight to emphasis levels
|
||||
if weight >= 1.5:
|
||||
parts.append(f"(({text}))")
|
||||
elif weight >= 1.2:
|
||||
parts.append(f"({text})")
|
||||
elif weight <= 0.5:
|
||||
parts.append(f"[{text}]")
|
||||
else:
|
||||
parts.append(text)
|
||||
return ", ".join(parts)
|
||||
|
||||
def _concat(self, prompts: List[Tuple[str, float]], separator: str) -> str:
|
||||
"""Simple concatenation with separator."""
|
||||
texts = [p[0] for p in prompts]
|
||||
return separator.join(texts)
|
||||
|
||||
def _weighted_average(self, prompts: List[Tuple[str, float]]) -> str:
|
||||
"""
|
||||
Weighted text combination using emphasis markers.
|
||||
Higher-weighted prompts receive stronger ComfyUI emphasis parentheses.
|
||||
"""
|
||||
total_weight = sum(w for _, w in prompts)
|
||||
if total_weight == 0:
|
||||
return ", ".join(p[0] for p in prompts)
|
||||
|
||||
# Normalize weights relative to average
|
||||
avg_weight = total_weight / len(prompts)
|
||||
|
||||
parts = []
|
||||
for text, weight in prompts:
|
||||
# Compute emphasis level based on weight ratio to average
|
||||
ratio = weight / avg_weight if avg_weight > 0 else 1.0
|
||||
if ratio >= 2.0:
|
||||
# Strong emphasis: triple parens
|
||||
parts.append(f"((({text})))")
|
||||
elif ratio >= 1.5:
|
||||
# High emphasis: double parens
|
||||
parts.append(f"(({text}))")
|
||||
elif ratio >= 1.2:
|
||||
# Moderate emphasis: single parens
|
||||
parts.append(f"({text})")
|
||||
elif ratio <= 0.5:
|
||||
# De-emphasis: square brackets
|
||||
parts.append(f"[{text}]")
|
||||
else:
|
||||
# Neutral: no markers
|
||||
parts.append(text)
|
||||
|
||||
return ", ".join(parts)
|
||||
@@ -0,0 +1,321 @@
|
||||
"""
|
||||
Prompt Refiner Node for ComfyUI
|
||||
Refines a raw prompt through iterative LLM passes for higher quality output.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import subprocess
|
||||
import time
|
||||
import threading
|
||||
from typing import Any, Dict, Optional, Tuple
|
||||
|
||||
from .prompt_generator_node import extract_final_prompt
|
||||
|
||||
try:
|
||||
import ollama
|
||||
OLLAMA_API_AVAILABLE = True
|
||||
except ImportError:
|
||||
OLLAMA_API_AVAILABLE = False
|
||||
|
||||
try:
|
||||
import comfy.utils
|
||||
COMFY_PROGRESS_AVAILABLE = True
|
||||
except ImportError:
|
||||
COMFY_PROGRESS_AVAILABLE = False
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class PromptRefinerNode:
|
||||
"""
|
||||
ComfyUI node for refining prompts using iterative LLM passes.
|
||||
|
||||
Takes a raw prompt string, sends it to Ollama with a refinement system prompt,
|
||||
and returns an improved version. Supports 1-3 refinement passes.
|
||||
"""
|
||||
|
||||
REFINEMENT_PROMPT = """You are an expert prompt engineer for Stable Diffusion.
|
||||
|
||||
Refine the following prompt to improve its quality, specificity, and coherence.
|
||||
Keep the core subject intact but enhance:
|
||||
- Descriptive detail (textures, lighting, atmosphere)
|
||||
- Technical quality markers (8k, highly detailed, masterpiece)
|
||||
- Composition and framing cues
|
||||
- Color palette hints
|
||||
|
||||
Return ONLY the refined prompt text. No explanations, no markdown formatting.
|
||||
|
||||
Original prompt: {prompt}
|
||||
|
||||
Refined prompt:"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls) -> dict[str, Any]:
|
||||
return {
|
||||
"required": {
|
||||
"prompt": (
|
||||
"STRING",
|
||||
{
|
||||
"multiline": True,
|
||||
"default": "",
|
||||
"placeholder": "Raw prompt to refine...",
|
||||
},
|
||||
),
|
||||
"model": (
|
||||
"STRING",
|
||||
{
|
||||
"default": "qwen3:8b",
|
||||
"placeholder": "Ollama model name",
|
||||
},
|
||||
),
|
||||
},
|
||||
"optional": {
|
||||
"passes": (
|
||||
"INT",
|
||||
{
|
||||
"default": 1,
|
||||
"min": 1,
|
||||
"max": 3,
|
||||
"step": 1,
|
||||
"display": "slider",
|
||||
},
|
||||
),
|
||||
"temperature": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 0.5,
|
||||
"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",)
|
||||
RETURN_NAMES = ("refined_prompt",)
|
||||
FUNCTION = "refine"
|
||||
CATEGORY = "text/generation"
|
||||
OUTPUT_NODE = False
|
||||
|
||||
def _generate_streaming(
|
||||
self,
|
||||
model: str,
|
||||
prompt: str,
|
||||
temperature: float,
|
||||
top_p: float,
|
||||
timeout: int,
|
||||
pbar: object = None,
|
||||
) -> Optional[str]:
|
||||
"""
|
||||
Stream ollama.generate() with per-chunk and total timeout enforcement.
|
||||
Returns the full response text, or None on failure (caller should fallback).
|
||||
"""
|
||||
chunks = []
|
||||
start = time.monotonic()
|
||||
first_chunk_timeout = min(timeout * 0.6, 90)
|
||||
chunk_timeout = 30
|
||||
got_first_chunk = False
|
||||
|
||||
try:
|
||||
stream = ollama.generate(
|
||||
model=model,
|
||||
prompt=prompt,
|
||||
stream=True,
|
||||
options={"temperature": temperature, "top_p": top_p},
|
||||
)
|
||||
|
||||
result_holder: Dict[str, Any] = {}
|
||||
|
||||
def _iter_next(it):
|
||||
"""Get next chunk from iterator in a thread."""
|
||||
try:
|
||||
result_holder["chunk"] = next(it)
|
||||
result_holder["done"] = False
|
||||
except StopIteration:
|
||||
result_holder["done"] = True
|
||||
except Exception as exc:
|
||||
result_holder["error"] = exc
|
||||
|
||||
it = iter(stream)
|
||||
chunk_count = 0
|
||||
|
||||
while True:
|
||||
elapsed = time.monotonic() - start
|
||||
if elapsed >= timeout:
|
||||
print(f"[PromptRefiner] Total timeout ({timeout}s) reached")
|
||||
break
|
||||
|
||||
result_holder.clear()
|
||||
t = threading.Thread(target=_iter_next, args=(it,), daemon=True)
|
||||
t.start()
|
||||
|
||||
wait_time = first_chunk_timeout if not got_first_chunk else chunk_timeout
|
||||
wait_time = min(wait_time, timeout - elapsed)
|
||||
t.join(timeout=wait_time)
|
||||
|
||||
if t.is_alive():
|
||||
label = "first chunk" if not got_first_chunk else "chunk"
|
||||
print(f"[PromptRefiner] Timeout waiting for {label} ({wait_time:.0f}s)")
|
||||
return None
|
||||
|
||||
if "error" in result_holder:
|
||||
raise result_holder["error"]
|
||||
|
||||
if result_holder.get("done", False):
|
||||
break
|
||||
|
||||
chunk = result_holder.get("chunk")
|
||||
if chunk is None:
|
||||
break
|
||||
|
||||
text = chunk.get("response", "")
|
||||
if text:
|
||||
chunks.append(text)
|
||||
got_first_chunk = True
|
||||
chunk_count += 1
|
||||
|
||||
if pbar is not None:
|
||||
progress = min(5 + int(chunk_count * 2), 95)
|
||||
pbar.update_absolute(progress)
|
||||
|
||||
except Exception as e:
|
||||
print(f"[PromptRefiner] Streaming error: {e}")
|
||||
return None
|
||||
|
||||
if not chunks:
|
||||
return None
|
||||
|
||||
elapsed = time.monotonic() - start
|
||||
full_text = "".join(chunks)
|
||||
print(f"[PromptRefiner] Streaming complete: {len(full_text)} chars in {elapsed:.1f}s")
|
||||
return full_text
|
||||
|
||||
def refine(
|
||||
self,
|
||||
prompt: str,
|
||||
model: str,
|
||||
passes: int = 1,
|
||||
temperature: float = 0.5,
|
||||
top_p: float = 0.9,
|
||||
seed: int = -1,
|
||||
timeout: int = 120,
|
||||
unique_id: str = None,
|
||||
) -> Tuple[str]:
|
||||
"""
|
||||
Refine a prompt through iterative LLM passes.
|
||||
|
||||
Args:
|
||||
prompt: Raw prompt string to refine
|
||||
model: Ollama model to use
|
||||
passes: Number of refinement iterations (1-3)
|
||||
temperature: Generation temperature
|
||||
seed: Seed for deterministic generation (-1 for random)
|
||||
timeout: Maximum generation time per pass
|
||||
unique_id: ComfyUI node execution ID for progress tracking
|
||||
|
||||
Returns:
|
||||
Tuple containing the refined prompt string
|
||||
"""
|
||||
if not prompt.strip():
|
||||
return ("[PromptRefiner] Please provide a prompt to refine.",)
|
||||
|
||||
# Initialize progress bar
|
||||
pbar = None
|
||||
if COMFY_PROGRESS_AVAILABLE and unique_id is not None:
|
||||
try:
|
||||
pbar = comfy.utils.ProgressBar(100, node_id=unique_id)
|
||||
pbar.update_absolute(0)
|
||||
except Exception:
|
||||
pbar = None
|
||||
|
||||
current_prompt = prompt.strip()
|
||||
effective_seed: Optional[int] = None if seed == -1 else seed
|
||||
|
||||
for i in range(passes):
|
||||
print(f"[PromptRefiner] Pass {i + 1}/{passes} with model='{model}'")
|
||||
|
||||
if pbar is not None:
|
||||
progress = int((i / passes) * 100)
|
||||
pbar.update_absolute(progress)
|
||||
|
||||
# Build refinement prompt
|
||||
refinement = self.REFINEMENT_PROMPT.format(prompt=current_prompt)
|
||||
|
||||
# Derive per-pass seed so multi-pass refinement isn't a no-op
|
||||
pass_seed = None if effective_seed is None else effective_seed + i
|
||||
|
||||
output = None
|
||||
if OLLAMA_API_AVAILABLE:
|
||||
try:
|
||||
output = self._generate_streaming(
|
||||
model=model,
|
||||
prompt=refinement,
|
||||
temperature=temperature,
|
||||
top_p=top_p,
|
||||
timeout=timeout,
|
||||
pbar=pbar,
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"[PromptRefiner] Streaming failed: {e}")
|
||||
|
||||
if output is None:
|
||||
# Fallback to subprocess
|
||||
print("[PromptRefiner] Falling back to subprocess")
|
||||
try:
|
||||
cmd = ["ollama", "run", model, refinement]
|
||||
result = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)
|
||||
if result.returncode != 0:
|
||||
return (f"[PromptRefiner] Pass {i + 1} failed: {result.stderr}",)
|
||||
output = result.stdout.strip()
|
||||
except subprocess.TimeoutExpired:
|
||||
return (f"[PromptRefiner] Pass {i + 1} timed out after {timeout}s",)
|
||||
except FileNotFoundError:
|
||||
return ("[PromptRefiner] Ollama not found. Install from: https://ollama.ai",)
|
||||
except Exception as e:
|
||||
return (f"[PromptRefiner] Pass {i + 1} error: {e}",)
|
||||
|
||||
# Clean the output
|
||||
cleaned = extract_final_prompt(output.strip())
|
||||
if cleaned:
|
||||
current_prompt = cleaned
|
||||
print(f"[PromptRefiner] Pass {i + 1} complete: {len(current_prompt)} chars")
|
||||
else:
|
||||
print(f"[PromptRefiner] Pass {i + 1} returned empty, keeping previous")
|
||||
|
||||
if pbar is not None:
|
||||
pbar.update_absolute(100)
|
||||
|
||||
return (current_prompt,)
|
||||
@@ -0,0 +1,93 @@
|
||||
"""
|
||||
Style Applier Node for ComfyUI
|
||||
Applies Cinematic or Still Image style keywords to prompts.
|
||||
"""
|
||||
|
||||
from typing import Tuple
|
||||
|
||||
|
||||
class StyleApplierNode:
|
||||
"""
|
||||
ComfyUI node for applying Cinematic or Still Image style keywords to prompts.
|
||||
|
||||
Inputs:
|
||||
- prompt: Base prompt text
|
||||
- style: "cinematic" or "still_image"
|
||||
- position: Where to add keywords ("prefix", "suffix", or "wrap")
|
||||
- emphasis: Optional emphasis level ("low", "medium", "high")
|
||||
- include_technical: Include camera/technical specs
|
||||
|
||||
Outputs:
|
||||
- styled_prompt: The prompt with style keywords added
|
||||
- style_keywords: Just the style keywords (for reference)
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
"""Define input parameters for the node."""
|
||||
from ..style_presets import StylePreset
|
||||
|
||||
return {
|
||||
"required": {
|
||||
"prompt": ("STRING", {
|
||||
"multiline": True,
|
||||
"default": "",
|
||||
"placeholder": "Enter your base prompt..."
|
||||
}),
|
||||
"style": (StylePreset.get_style_choices(), {"default": "cinematic"}),
|
||||
},
|
||||
"optional": {
|
||||
"position": (["suffix", "prefix", "wrap"], {"default": "suffix"}),
|
||||
"emphasis": (["medium", "low", "high"], {"default": "medium"}),
|
||||
"include_technical": ("BOOLEAN", {"default": True}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING", "STRING",)
|
||||
RETURN_NAMES = ("styled_prompt", "style_keywords",)
|
||||
FUNCTION = "apply_style"
|
||||
CATEGORY = "text/generation"
|
||||
|
||||
def apply_style(
|
||||
self,
|
||||
prompt: str,
|
||||
style: str,
|
||||
position: str = "suffix",
|
||||
emphasis: str = "medium",
|
||||
include_technical: bool = True
|
||||
) -> Tuple[str, str]:
|
||||
"""Apply style keywords to a prompt."""
|
||||
from ..style_presets import StylePreset
|
||||
|
||||
# Normalize inputs
|
||||
prompt = prompt.strip() if prompt else ""
|
||||
style = style.strip() if style else "cinematic"
|
||||
|
||||
# Validate style
|
||||
available_styles = StylePreset.get_style_choices()
|
||||
if style not in available_styles:
|
||||
return (f"[StyleApplier] Error: Unknown style '{style}'. Available: {available_styles}", "")
|
||||
|
||||
# Get style keywords
|
||||
try:
|
||||
style_keywords = StylePreset().get_style_prompt(
|
||||
style=style,
|
||||
emphasis=emphasis,
|
||||
include_technical=include_technical
|
||||
)
|
||||
except ValueError as e:
|
||||
return (f"[StyleApplier] Error getting style: {e}", "")
|
||||
|
||||
# Handle empty prompt
|
||||
if not prompt:
|
||||
return (style_keywords, style_keywords)
|
||||
|
||||
# Apply style based on position
|
||||
if position == "prefix":
|
||||
styled_prompt = f"{style_keywords}, {prompt}"
|
||||
elif position == "suffix":
|
||||
styled_prompt = f"{prompt}, {style_keywords}"
|
||||
else: # wrap
|
||||
styled_prompt = f"{style_keywords}, {prompt}, {style_keywords}"
|
||||
|
||||
return (styled_prompt, style_keywords)
|
||||
Reference in New Issue
Block a user