From aac7906b2973bd082f6315c47f3f81a9a1dbbdfe Mon Sep 17 00:00:00 2001 From: limbicnation Date: Fri, 1 May 2026 02:26:19 +0200 Subject: [PATCH] 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 --- .../comfyui-prompt-generator/__init__.py | 29 ++ .../nodes/negative_prompt_node.py | 391 ++++++++++++++++++ .../nodes/prompt_combiner_node.py | 246 +++++++++++ .../nodes/prompt_refiner_node.py | 321 ++++++++++++++ .../nodes/style_applier_node.py | 93 +++++ 5 files changed, 1080 insertions(+) create mode 100644 custom_nodes/comfyui-prompt-generator/__init__.py create mode 100644 custom_nodes/comfyui-prompt-generator/nodes/negative_prompt_node.py create mode 100644 custom_nodes/comfyui-prompt-generator/nodes/prompt_combiner_node.py create mode 100644 custom_nodes/comfyui-prompt-generator/nodes/prompt_refiner_node.py create mode 100644 custom_nodes/comfyui-prompt-generator/nodes/style_applier_node.py diff --git a/custom_nodes/comfyui-prompt-generator/__init__.py b/custom_nodes/comfyui-prompt-generator/__init__.py new file mode 100644 index 0000000..18a96a9 --- /dev/null +++ b/custom_nodes/comfyui-prompt-generator/__init__.py @@ -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"] + diff --git a/custom_nodes/comfyui-prompt-generator/nodes/negative_prompt_node.py b/custom_nodes/comfyui-prompt-generator/nodes/negative_prompt_node.py new file mode 100644 index 0000000..9f8d8f9 --- /dev/null +++ b/custom_nodes/comfyui-prompt-generator/nodes/negative_prompt_node.py @@ -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) diff --git a/custom_nodes/comfyui-prompt-generator/nodes/prompt_combiner_node.py b/custom_nodes/comfyui-prompt-generator/nodes/prompt_combiner_node.py new file mode 100644 index 0000000..2859edf --- /dev/null +++ b/custom_nodes/comfyui-prompt-generator/nodes/prompt_combiner_node.py @@ -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) diff --git a/custom_nodes/comfyui-prompt-generator/nodes/prompt_refiner_node.py b/custom_nodes/comfyui-prompt-generator/nodes/prompt_refiner_node.py new file mode 100644 index 0000000..24f37d5 --- /dev/null +++ b/custom_nodes/comfyui-prompt-generator/nodes/prompt_refiner_node.py @@ -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,) diff --git a/custom_nodes/comfyui-prompt-generator/nodes/style_applier_node.py b/custom_nodes/comfyui-prompt-generator/nodes/style_applier_node.py new file mode 100644 index 0000000..356cd32 --- /dev/null +++ b/custom_nodes/comfyui-prompt-generator/nodes/style_applier_node.py @@ -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)