88 lines
3.0 KiB
Python
88 lines
3.0 KiB
Python
import re
|
|
|
|
class Prompt_Extender:
|
|
"""
|
|
A ComfyUI node for extending double prompts.
|
|
Takes a double prompt input and adds extensions to both positive and negative parts.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"double_prompt": ("STRING", {"forceInput": True}),
|
|
"prompt_additions": ("STRING", {
|
|
"multiline": True,
|
|
"default": "beautiful, high quality\n---\nblurry, low quality"
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("double_prompt",)
|
|
FUNCTION = "extend_prompts"
|
|
CATEGORY = "hexxacubic"
|
|
|
|
def extend_prompts(self, double_prompt, prompt_additions):
|
|
# Parse input double prompt
|
|
positive_prompt, negative_prompt = self._parse_double_prompt(double_prompt)
|
|
|
|
# Parse additions
|
|
positive_addition, negative_addition = self._parse_double_prompt(prompt_additions)
|
|
|
|
# Combine the prompts
|
|
combined_positive = self._combine_prompts(positive_prompt, positive_addition)
|
|
combined_negative = self._combine_prompts(negative_prompt, negative_addition)
|
|
|
|
# Return as double prompt format
|
|
if combined_positive and combined_negative:
|
|
return (f"{combined_positive}\n---\n{combined_negative}",)
|
|
elif combined_positive:
|
|
return (combined_positive,)
|
|
elif combined_negative:
|
|
return (f"\n---\n{combined_negative}",)
|
|
else:
|
|
return ("",)
|
|
|
|
def _parse_double_prompt(self, text):
|
|
"""Parse double prompt format into positive and negative parts"""
|
|
if not text.strip():
|
|
return "", ""
|
|
|
|
lines = text.split('\n')
|
|
|
|
# Find separator (---, ----, or -----)
|
|
separator_idx = -1
|
|
for i, line in enumerate(lines):
|
|
if line.strip() in ["---", "----", "-----"]:
|
|
separator_idx = i
|
|
break
|
|
|
|
if separator_idx != -1:
|
|
positive = "\n".join(lines[:separator_idx]).strip()
|
|
negative = "\n".join(lines[separator_idx+1:]).strip()
|
|
else:
|
|
positive = text.strip()
|
|
negative = ""
|
|
|
|
return positive, negative
|
|
|
|
def _combine_prompts(self, base_prompt, addition):
|
|
"""Intelligently combines base prompt and addition"""
|
|
base_prompt = base_prompt.strip()
|
|
addition = addition.strip()
|
|
|
|
if not base_prompt:
|
|
return addition
|
|
if not addition:
|
|
return base_prompt
|
|
|
|
# Add comma between prompts if none exists
|
|
if base_prompt and addition:
|
|
if not base_prompt.endswith(',') and not base_prompt.endswith('.'):
|
|
return f"{base_prompt}, {addition}"
|
|
else:
|
|
return f"{base_prompt} {addition}"
|
|
|
|
return base_prompt
|