ComfyUI-Ethanfel-Prompt-Builder/node_hardcore_position.py

from __future__ import annotations

import json

try:
    from . import krea2_pose_variant_catalog
    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.
    import krea2_pose_variant_catalog
    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}"


def _variant_family(value):
    family = str(value or "any")
    if family == "penetration":
        family = "penetrative"
    return family if family in hardcore_position_family_choices() else "any"


def _variant_positions(variant):
    valid = set(hardcore_position_key_choices())
    return [str(key) for key in variant.get("position_keys", []) if str(key) in valid]


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 SxCPKrea2PoseVariant:
    @classmethod
    def INPUT_TYPES(cls):
        keys = krea2_pose_variant_catalog.variant_keys()
        return {
            "required": {
                "variant_key": (keys or ["missing_catalog_variant"], {"default": keys[0] if keys else "missing_catalog_variant"}),
                "combine_mode": (["replace", "add"], {"default": "replace"}),
            },
            "optional": {
                "hardcore_position_config": (SXCP_HARDCORE_POSITION_CONFIG,),
            },
        }

    RETURN_TYPES = (SXCP_HARDCORE_POSITION_CONFIG, "STRING", "STRING", "STRING", "STRING", "STRING")
    RETURN_NAMES = (
        "hardcore_position_config",
        "variant_key",
        "prompt_cues",
        "avoid_cues",
        "summary",
        "variant_json",
    )
    FUNCTION = "build"
    CATEGORY = "prompt_builder"

    def build(self, variant_key, combine_mode="replace", hardcore_position_config=""):
        variant = krea2_pose_variant_catalog.get_variant(variant_key)
        if not variant:
            empty = {
                "key": str(variant_key or ""),
                "status": "missing",
                "summary": "missing Krea2 pose variant",
            }
            return hardcore_position_config or "", str(variant_key or ""), "", "", empty["summary"], json.dumps(empty, sort_keys=True)

        positions = _variant_positions(variant)
        family = _variant_family(variant.get("action_family") or variant.get("family"))
        config = build_hardcore_position_pool_json(
            hardcore_position_config=hardcore_position_config or "",
            combine_mode=combine_mode,
            family=family,
            selected_positions=positions,
        )
        prompt_cues = "; ".join(str(cue) for cue in variant.get("prompt_cues", []) if str(cue).strip())
        avoid_cues = "; ".join(str(cue) for cue in variant.get("avoid_cues", []) if str(cue).strip())
        summary = (
            f"variant={variant.get('key')}; status={variant.get('status')}; "
            f"family={family}; positions={','.join(positions) or 'none'}"
        )
        return (
            config,
            str(variant.get("key") or variant_key),
            prompt_cues,
            avoid_cues,
            summary,
            json.dumps(variant, ensure_ascii=True, sort_keys=True),
        )


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,
    "SxCPKrea2PoseVariant": SxCPKrea2PoseVariant,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "SxCPHardcorePositionPool": "SxCP Hardcore Position Pool",
    "SxCPHardcoreActionFilter": "SxCP Hardcore Action Filter",
    "SxCPKrea2PoseVariant": "SxCP Krea2 Pose Variant",
}