Files
ComfyUI-Ethanfel-Prompt-Bui…/node_hardcore_position.py
T

137 lines
4.6 KiB
Python

from __future__ import annotations
import json
try:
from .hardcore_position_config import (
build_hardcore_action_filter_json,
build_hardcore_position_pool_json,
hardcore_position_family_choices,
hardcore_position_focus_choices,
hardcore_position_key_choices,
)
except ImportError: # Allows local smoke tests from the repository root.
from hardcore_position_config import (
build_hardcore_action_filter_json,
build_hardcore_position_pool_json,
hardcore_position_family_choices,
hardcore_position_focus_choices,
hardcore_position_key_choices,
)
SXCP_HARDCORE_POSITION_CONFIG = "SXCP_HARDCORE_POSITION_CONFIG"
def _choice_input_key(prefix, choice):
key = "".join(char if char.isalnum() else "_" for char in str(choice).lower()).strip("_")
while "__" in key:
key = key.replace("__", "_")
return f"{prefix}_{key}"
class SxCPHardcorePositionPool:
@classmethod
def INPUT_TYPES(cls):
required = {
"combine_mode": (["replace", "add"], {"default": "replace"}),
"family": (hardcore_position_family_choices(), {"default": "any"}),
}
for choice in hardcore_position_key_choices():
required[_choice_input_key("include", choice)] = ("BOOLEAN", {"default": False})
return {
"required": required,
"optional": {
"hardcore_position_config": (SXCP_HARDCORE_POSITION_CONFIG,),
},
}
RETURN_TYPES = (SXCP_HARDCORE_POSITION_CONFIG, "STRING")
RETURN_NAMES = ("hardcore_position_config", "summary")
FUNCTION = "build"
CATEGORY = "prompt_builder"
def build(self, combine_mode="replace", family="any", hardcore_position_config="", **kwargs):
selected = [
choice
for choice in hardcore_position_key_choices()
if bool(kwargs.get(_choice_input_key("include", choice), False))
]
config = build_hardcore_position_pool_json(
hardcore_position_config=hardcore_position_config or "",
combine_mode=combine_mode,
family=family,
selected_positions=selected,
)
return config, json.loads(config).get("summary", "")
class SxCPHardcoreActionFilter:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"focus": (hardcore_position_focus_choices(), {"default": "keep_pool"}),
"allow_toys": ("BOOLEAN", {"default": False}),
"allow_double": ("BOOLEAN", {"default": False}),
"allow_penetration": ("BOOLEAN", {"default": True}),
"allow_foreplay": ("BOOLEAN", {"default": True}),
"allow_interaction": ("BOOLEAN", {"default": True}),
"allow_manual": ("BOOLEAN", {"default": True}),
"allow_oral": ("BOOLEAN", {"default": True}),
"allow_outercourse": ("BOOLEAN", {"default": True}),
"allow_anal": ("BOOLEAN", {"default": True}),
"allow_climax": ("BOOLEAN", {"default": True}),
},
"optional": {
"hardcore_position_config": (SXCP_HARDCORE_POSITION_CONFIG,),
},
}
RETURN_TYPES = (SXCP_HARDCORE_POSITION_CONFIG, "STRING")
RETURN_NAMES = ("hardcore_position_config", "summary")
FUNCTION = "build"
CATEGORY = "prompt_builder"
def build(
self,
focus,
allow_toys,
allow_double,
allow_penetration,
allow_foreplay,
allow_interaction,
allow_manual,
allow_oral,
allow_outercourse,
allow_anal,
allow_climax,
hardcore_position_config="",
):
config = build_hardcore_action_filter_json(
hardcore_position_config=hardcore_position_config or "",
focus=focus,
allow_toys=allow_toys,
allow_double=allow_double,
allow_penetration=allow_penetration,
allow_foreplay=allow_foreplay,
allow_interaction=allow_interaction,
allow_manual=allow_manual,
allow_oral=allow_oral,
allow_outercourse=allow_outercourse,
allow_anal=allow_anal,
allow_climax=allow_climax,
)
return config, json.loads(config).get("summary", "")
NODE_CLASS_MAPPINGS = {
"SxCPHardcorePositionPool": SxCPHardcorePositionPool,
"SxCPHardcoreActionFilter": SxCPHardcoreActionFilter,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"SxCPHardcorePositionPool": "SxCP Hardcore Position Pool",
"SxCPHardcoreActionFilter": "SxCP Hardcore Action Filter",
}