612 lines
17 KiB
Python
612 lines
17 KiB
Python
import json
|
|
import os
|
|
import random
|
|
import re
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
|
|
USER_DIR = Path(os.getcwd()) / "user" / "default"
|
|
CONFIG_DIR = USER_DIR / "ComfyUI-Complex-Prompt"
|
|
CONFIG_PATH = CONFIG_DIR / "config.json"
|
|
DEFAULT_CONFIG: dict[str, Any] = {}
|
|
CONDITION_FUNCTIONS = {
|
|
"float": float,
|
|
"int": int,
|
|
"str": str,
|
|
}
|
|
|
|
|
|
def log(message: str) -> None:
|
|
print(f"[ComfyUI-Complex-Prompt]: {message}")
|
|
|
|
|
|
def read_json(path: Path) -> dict[str, Any]:
|
|
with path.open("r", encoding="utf-8") as file:
|
|
return json.load(file)
|
|
|
|
|
|
def write_json(path: Path, data: dict[str, Any]) -> None:
|
|
with path.open("w", encoding="utf-8") as file:
|
|
json.dump(data, file, ensure_ascii=False, indent=2)
|
|
file.write("\n")
|
|
|
|
|
|
def load_config() -> dict[str, Any]:
|
|
CONFIG_DIR.mkdir(parents=True, exist_ok=True)
|
|
|
|
if not CONFIG_PATH.exists():
|
|
write_json(CONFIG_PATH, DEFAULT_CONFIG)
|
|
log(f"Created config file at {CONFIG_PATH}")
|
|
return DEFAULT_CONFIG.copy()
|
|
|
|
log(f"Reading config from {CONFIG_PATH}")
|
|
return read_json(CONFIG_PATH)
|
|
|
|
|
|
CONFIG = load_config()
|
|
|
|
ARTEMKO7V_COMPLEX_PROMPT_VARS = "ArtemKo7vComplexPromptVars"
|
|
VAR_PATTERN = re.compile(r"\$([A-Za-z_][A-Za-z0-9_]*)")
|
|
DYNAMIC_VARIANT_PATTERN = re.compile(r"\{[^{}]*\}")
|
|
|
|
_PROMPT_GENERATOR = None
|
|
|
|
|
|
class ConditionValue:
|
|
def __init__(self, value: Any):
|
|
self.value = value
|
|
|
|
def __str__(self) -> str:
|
|
return str(self.value)
|
|
|
|
def __repr__(self) -> str:
|
|
return repr(str(self))
|
|
|
|
def __bool__(self) -> bool:
|
|
return bool(str(self))
|
|
|
|
def __int__(self) -> int:
|
|
return int(str(self))
|
|
|
|
def __float__(self) -> float:
|
|
return float(str(self))
|
|
|
|
def __eq__(self, other: Any) -> bool:
|
|
return str(self) == str(self.unwrap(other))
|
|
|
|
def __ne__(self, other: Any) -> bool:
|
|
return not self == other
|
|
|
|
def __lt__(self, other: Any) -> bool:
|
|
return self.compare(other, lambda left, right: left < right)
|
|
|
|
def __le__(self, other: Any) -> bool:
|
|
return self.compare(other, lambda left, right: left <= right)
|
|
|
|
def __gt__(self, other: Any) -> bool:
|
|
return self.compare(other, lambda left, right: left > right)
|
|
|
|
def __ge__(self, other: Any) -> bool:
|
|
return self.compare(other, lambda left, right: left >= right)
|
|
|
|
def unwrap(self, other: Any) -> Any:
|
|
if isinstance(other, ConditionValue):
|
|
return other.value
|
|
|
|
return other
|
|
|
|
def compare(self, other: Any, operation):
|
|
other = self.unwrap(other)
|
|
|
|
try:
|
|
return operation(float(self), float(other))
|
|
except (TypeError, ValueError):
|
|
return operation(str(self), str(other))
|
|
|
|
|
|
def get_prompt_generator():
|
|
global _PROMPT_GENERATOR
|
|
|
|
if _PROMPT_GENERATOR is None:
|
|
try:
|
|
from dynamicprompts.generators import RandomPromptGenerator
|
|
except ImportError as error:
|
|
raise ImportError(
|
|
"ComfyUI-Complex-Prompt requires the 'dynamicprompts' package. "
|
|
"Install this extension's requirements.txt dependencies."
|
|
) from error
|
|
|
|
_PROMPT_GENERATOR = RandomPromptGenerator()
|
|
|
|
return _PROMPT_GENERATOR
|
|
|
|
|
|
def mask_vars(prompt: str) -> tuple[str, dict[str, str]]:
|
|
placeholders: dict[str, str] = {}
|
|
|
|
def replace(match: re.Match[str]) -> str:
|
|
placeholder = f"AK7VCOMPLEXPROMPTVAR{len(placeholders)}TOKEN"
|
|
placeholders[placeholder] = match.group(0)
|
|
return placeholder
|
|
|
|
return VAR_PATTERN.sub(replace, prompt), placeholders
|
|
|
|
|
|
def restore_vars(prompt: str, placeholders: dict[str, str]) -> str:
|
|
for placeholder, variable in placeholders.items():
|
|
prompt = prompt.replace(placeholder, variable)
|
|
|
|
return prompt
|
|
|
|
|
|
def has_dynamic_prompt_syntax(prompt: str) -> bool:
|
|
return bool(DYNAMIC_VARIANT_PATTERN.search(prompt)) or "__" in prompt
|
|
|
|
|
|
def generate_dynamic_prompt(prompt: str | None, seed: int) -> str:
|
|
if prompt in (None, ""):
|
|
return ""
|
|
|
|
masked_prompt, placeholders = mask_vars(prompt)
|
|
if not has_dynamic_prompt_syntax(masked_prompt):
|
|
return prompt
|
|
|
|
generator = get_prompt_generator()
|
|
# dynamicprompts treats the integer 0 as an omitted seed; [0] is explicit.
|
|
prompts = generator.generate(masked_prompt, 1, seeds=[0] if seed == 0 else seed)
|
|
generated_prompt = prompts[0] if prompts else ""
|
|
return restore_vars(generated_prompt, placeholders)
|
|
|
|
|
|
|
|
def apply_vars(prompt: str, vars: dict[str, str] | None = None) -> str:
|
|
if not vars:
|
|
return prompt
|
|
|
|
def replace(match: re.Match[str]) -> str:
|
|
variable_name = match.group(1)
|
|
if variable_name not in vars:
|
|
return match.group(0)
|
|
|
|
return str(vars[variable_name])
|
|
|
|
return VAR_PATTERN.sub(replace, prompt)
|
|
|
|
|
|
def resolve_prompt(
|
|
prompt: str | None,
|
|
seed: int,
|
|
vars: dict[str, Any] | None = None,
|
|
) -> str:
|
|
"""Expand choices, then recursively resolve the selected variable values."""
|
|
def resolve(value: str | None, resolving: frozenset[str]) -> str:
|
|
generated = generate_dynamic_prompt(value, seed)
|
|
if not vars:
|
|
return generated
|
|
|
|
def replace(match: re.Match[str]) -> str:
|
|
name = match.group(1)
|
|
# Keep unknown and cyclic references visible without expanding forever.
|
|
if name not in vars or name in resolving:
|
|
return match.group(0)
|
|
|
|
return resolve(str(vars[name]), resolving | {name})
|
|
|
|
return VAR_PATTERN.sub(replace, generated)
|
|
|
|
return resolve(prompt, frozenset())
|
|
|
|
|
|
def is_json_variable_value(value: Any) -> bool:
|
|
return isinstance(value, (int, float, str)) and not isinstance(value, bool)
|
|
|
|
|
|
def normalize_json_object_text(json_text: str) -> str:
|
|
json_text = json_text.strip()
|
|
if not json_text:
|
|
return "{}"
|
|
|
|
if json_text.startswith(("{", "[")):
|
|
return json_text
|
|
|
|
return "{" + json_text + "}"
|
|
|
|
|
|
def parse_json_vars(
|
|
json_text: str,
|
|
seed: int,
|
|
vars: dict[str, Any] | None = None,
|
|
) -> dict[str, Any]:
|
|
json_text = normalize_json_object_text(json_text)
|
|
|
|
try:
|
|
parsed = json.loads(json_text)
|
|
except json.JSONDecodeError as error:
|
|
raise ValueError(f"Invalid JSON: {error.msg}") from error
|
|
|
|
if not isinstance(parsed, dict):
|
|
raise ValueError("JSON must be an object with variable names as keys.")
|
|
|
|
result = dict(vars or {})
|
|
rng = random.Random(seed)
|
|
|
|
for key, value in parsed.items():
|
|
if is_json_variable_value(value):
|
|
result[key] = value
|
|
continue
|
|
|
|
if (
|
|
isinstance(value, list)
|
|
and value
|
|
and all(is_json_variable_value(item) for item in value)
|
|
):
|
|
result[key] = rng.choice(value)
|
|
|
|
return result
|
|
|
|
|
|
def parse_choice_map(choices: str) -> dict[str, str]:
|
|
choices = choices.strip()
|
|
if not choices:
|
|
return {}
|
|
|
|
if not choices.startswith("{"):
|
|
choices = "{" + choices + "}"
|
|
|
|
try:
|
|
parsed = json.loads(choices)
|
|
except json.JSONDecodeError as error:
|
|
raise ValueError(f"Invalid choices JSON: {error.msg}") from error
|
|
|
|
if not isinstance(parsed, dict):
|
|
raise ValueError("Choices must be a JSON object.")
|
|
|
|
result: dict[str, str] = {}
|
|
for key, value in parsed.items():
|
|
if is_json_variable_value(value):
|
|
result[str(key)] = str(value)
|
|
|
|
return result
|
|
|
|
|
|
def normalize_condition(condition: str) -> str:
|
|
normalized = VAR_PATTERN.sub(r"\1", condition)
|
|
normalized = normalized.replace("&&", " and ")
|
|
normalized = normalized.replace("||", " or ")
|
|
normalized = re.sub(r"!(?!=)", " not ", normalized)
|
|
return normalized
|
|
|
|
|
|
def should_set_variable(
|
|
condition: str,
|
|
vars: dict[str, str] | None = None,
|
|
) -> bool:
|
|
condition = condition.strip()
|
|
if not condition:
|
|
return True
|
|
|
|
try:
|
|
from simpleeval import NameNotDefined, SimpleEval
|
|
except ImportError as error:
|
|
raise ImportError(
|
|
"ComfyUI-Complex-Prompt conditional variables require the "
|
|
"'simpleeval' package. Install this extension's requirements.txt "
|
|
"dependencies."
|
|
) from error
|
|
|
|
evaluator = SimpleEval(
|
|
functions=CONDITION_FUNCTIONS,
|
|
names={
|
|
key: ConditionValue(value)
|
|
for key, value in (vars or {}).items()
|
|
if re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", key)
|
|
},
|
|
)
|
|
|
|
try:
|
|
return bool(evaluator.eval(normalize_condition(condition)))
|
|
except NameNotDefined:
|
|
return False
|
|
|
|
|
|
class ArtemKo7vComplexPrompt:
|
|
CATEGORY = "ArtemKo7v"
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("string",)
|
|
FUNCTION = "generate_prompt"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"prompt": (
|
|
"STRING",
|
|
{
|
|
"multiline": True,
|
|
"dynamicPrompts": False,
|
|
},
|
|
),
|
|
"seed": (
|
|
"INT",
|
|
{
|
|
"default": 0,
|
|
"min": 0,
|
|
"max": 18446744073709551615,
|
|
"control_after_generate": "randomize",
|
|
},
|
|
),
|
|
},
|
|
"optional": {
|
|
"vars": (ARTEMKO7V_COMPLEX_PROMPT_VARS,),
|
|
},
|
|
}
|
|
|
|
def generate_prompt(
|
|
self,
|
|
prompt: str,
|
|
seed: int,
|
|
vars: dict[str, str] | None = None,
|
|
):
|
|
return (resolve_prompt(prompt, seed, vars),)
|
|
|
|
|
|
class ArtemKo7vComplexPromptSetVariable:
|
|
CATEGORY = "ArtemKo7v"
|
|
RETURN_TYPES = (ARTEMKO7V_COMPLEX_PROMPT_VARS, "BOOLEAN")
|
|
RETURN_NAMES = ("vars", "wasSet")
|
|
FUNCTION = "set_variable"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"variable_name": (
|
|
"STRING",
|
|
{
|
|
"default": "",
|
|
},
|
|
),
|
|
"value": (
|
|
"STRING",
|
|
{
|
|
"multiline": True,
|
|
"dynamicPrompts": False,
|
|
},
|
|
),
|
|
"condition": (
|
|
"STRING",
|
|
{
|
|
"default": "",
|
|
},
|
|
),
|
|
"seed": (
|
|
"INT",
|
|
{
|
|
"default": 0,
|
|
"min": 0,
|
|
"max": 18446744073709551615,
|
|
"control_after_generate": "randomize",
|
|
},
|
|
),
|
|
"trim": (
|
|
"BOOLEAN",
|
|
{
|
|
"default": False,
|
|
},
|
|
),
|
|
},
|
|
"optional": {
|
|
"vars": (ARTEMKO7V_COMPLEX_PROMPT_VARS,),
|
|
},
|
|
}
|
|
|
|
def set_variable(
|
|
self,
|
|
variable_name: str,
|
|
value: str | None,
|
|
condition: str,
|
|
seed: int,
|
|
trim: bool = False,
|
|
vars: dict[str, str] | None = None,
|
|
):
|
|
result = dict(vars or {})
|
|
if not should_set_variable(condition, vars):
|
|
return (result, False)
|
|
|
|
value = "" if value is None else value
|
|
stored_value = resolve_prompt(value, seed, vars)
|
|
if trim:
|
|
stored_value = stored_value.strip()
|
|
|
|
result[variable_name] = stored_value
|
|
return (result, True)
|
|
|
|
|
|
class ArtemKo7vComplexPromptSetVariableByChoice:
|
|
CATEGORY = "ArtemKo7v"
|
|
RETURN_TYPES = (ARTEMKO7V_COMPLEX_PROMPT_VARS, "BOOLEAN")
|
|
RETURN_NAMES = ("vars", "wasSet")
|
|
FUNCTION = "set_variable_by_choice"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"variable_name": (
|
|
"STRING",
|
|
{
|
|
"default": "",
|
|
},
|
|
),
|
|
"source_variable": (
|
|
"STRING",
|
|
{
|
|
"default": "",
|
|
},
|
|
),
|
|
"choices": (
|
|
"STRING",
|
|
{
|
|
"multiline": True,
|
|
"dynamicPrompts": False,
|
|
"default": "",
|
|
},
|
|
),
|
|
"seed": (
|
|
"INT",
|
|
{
|
|
"default": 0,
|
|
"min": 0,
|
|
"max": 18446744073709551615,
|
|
"control_after_generate": "randomize",
|
|
},
|
|
),
|
|
"trim": (
|
|
"BOOLEAN",
|
|
{
|
|
"default": False,
|
|
},
|
|
),
|
|
},
|
|
"optional": {
|
|
"vars": (ARTEMKO7V_COMPLEX_PROMPT_VARS,),
|
|
},
|
|
}
|
|
|
|
def set_variable_by_choice(
|
|
self,
|
|
variable_name: str,
|
|
source_variable: str,
|
|
choices: str,
|
|
seed: int,
|
|
trim: bool = False,
|
|
vars: dict[str, str] | None = None,
|
|
):
|
|
result = dict(vars or {})
|
|
source_variable = source_variable.strip()
|
|
if source_variable.startswith("$"):
|
|
source_variable = source_variable[1:]
|
|
|
|
if not source_variable or source_variable not in result:
|
|
return (result, False)
|
|
|
|
choice_map = parse_choice_map(choices)
|
|
source_value = str(result[source_variable])
|
|
value = choice_map.get(source_value, choice_map.get("__default__"))
|
|
if value is None:
|
|
return (result, False)
|
|
|
|
stored_value = resolve_prompt(value, seed, vars)
|
|
if trim:
|
|
stored_value = stored_value.strip()
|
|
|
|
result[variable_name] = stored_value
|
|
return (result, True)
|
|
|
|
|
|
class ArtemKo7vComplexPromptParseJSON:
|
|
CATEGORY = "ArtemKo7v"
|
|
RETURN_TYPES = (ARTEMKO7V_COMPLEX_PROMPT_VARS,)
|
|
RETURN_NAMES = ("vars",)
|
|
FUNCTION = "parse_json"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"json_text": (
|
|
"STRING",
|
|
{
|
|
"multiline": True,
|
|
"dynamicPrompts": False,
|
|
"default": "",
|
|
},
|
|
),
|
|
"seed": (
|
|
"INT",
|
|
{
|
|
"default": 0,
|
|
"min": 0,
|
|
"max": 18446744073709551615,
|
|
"control_after_generate": "randomize",
|
|
},
|
|
),
|
|
},
|
|
"optional": {
|
|
"vars": (ARTEMKO7V_COMPLEX_PROMPT_VARS,),
|
|
},
|
|
}
|
|
|
|
def parse_json(
|
|
self,
|
|
json_text: str,
|
|
seed: int,
|
|
vars: dict[str, Any] | None = None,
|
|
):
|
|
return (parse_json_vars(json_text, seed, vars),)
|
|
|
|
|
|
class ArtemKo7vComplexPromptCombineVars:
|
|
"""Combines vars inputs in order, with later inputs taking precedence."""
|
|
|
|
CATEGORY = "ArtemKo7v"
|
|
MAX_VAR_INDEX = 64
|
|
RETURN_TYPES = (ARTEMKO7V_COMPLEX_PROMPT_VARS,)
|
|
RETURN_NAMES = ("vars",)
|
|
FUNCTION = "combine_vars"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"vars_1": (ARTEMKO7V_COMPLEX_PROMPT_VARS,),
|
|
},
|
|
}
|
|
|
|
def combine_vars(self, vars_1: dict[str, Any], **kwargs):
|
|
result = dict(vars_1 or {})
|
|
for index in range(2, self.MAX_VAR_INDEX + 1):
|
|
vars = kwargs.get(f"vars_{index}")
|
|
if vars:
|
|
result.update(vars)
|
|
|
|
return (result,)
|
|
|
|
|
|
class ArtemKo7vComplexPromptEmptyString:
|
|
CATEGORY = "ArtemKo7v"
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("string",)
|
|
FUNCTION = "get_empty_string"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {"required": {}}
|
|
|
|
def get_empty_string(self):
|
|
return ("",)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"ArtemKo7vComplexPrompt": ArtemKo7vComplexPrompt,
|
|
"ArtemKo7vComplexPromptSetVariable": ArtemKo7vComplexPromptSetVariable,
|
|
"ArtemKo7vComplexPromptSetVariableByChoice": (
|
|
ArtemKo7vComplexPromptSetVariableByChoice
|
|
),
|
|
"ArtemKo7vComplexPromptParseJSON": ArtemKo7vComplexPromptParseJSON,
|
|
"ArtemKo7vComplexPromptCombineVars": ArtemKo7vComplexPromptCombineVars,
|
|
"ArtemKo7vComplexPromptEmptyString": ArtemKo7vComplexPromptEmptyString,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"ArtemKo7vComplexPrompt": "Complex Prompt",
|
|
"ArtemKo7vComplexPromptSetVariable": "Complex Prompt Set Variable",
|
|
"ArtemKo7vComplexPromptSetVariableByChoice": (
|
|
"Complex Prompt Set Variable By Choice"
|
|
),
|
|
"ArtemKo7vComplexPromptParseJSON": "Complex Prompt Parse JSON",
|
|
"ArtemKo7vComplexPromptCombineVars": "Complex Prompt Combine Vars",
|
|
"ArtemKo7vComplexPromptEmptyString": "Complex Prompt Empty String",
|
|
}
|