270 lines
9.6 KiB
Python
270 lines
9.6 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
|
|
try:
|
|
from . import krea2_eval_log
|
|
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_eval_log
|
|
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 SxCPKrea2VariantEvidence:
|
|
@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"}),
|
|
"result": (["accepted", "rejected", "inconclusive", "any"], {"default": "accepted"}),
|
|
},
|
|
"optional": {
|
|
"variant_key_in": ("STRING", {"default": ""}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "STRING", "STRING", "STRING", "INT", "STRING")
|
|
RETURN_NAMES = (
|
|
"summary",
|
|
"baseline_image_path",
|
|
"candidate_image_path",
|
|
"evidence_json",
|
|
"seed",
|
|
"decision",
|
|
)
|
|
FUNCTION = "build"
|
|
CATEGORY = "prompt_builder"
|
|
|
|
def build(self, variant_key, result="accepted", variant_key_in=""):
|
|
key = str(variant_key_in or variant_key or "").strip()
|
|
result_filter = None if result == "any" else result
|
|
entries = krea2_eval_log.entries_for_variant(key, result=result_filter)
|
|
if not entries:
|
|
empty = {
|
|
"variant_key": key,
|
|
"result": result,
|
|
"summary": "no Krea2 eval evidence found",
|
|
}
|
|
return empty["summary"], "", "", json.dumps(empty, ensure_ascii=True, sort_keys=True), -1, ""
|
|
entry = entries[0]
|
|
summary = (
|
|
f"evidence={entry.get('id')}; variant={entry.get('variant_key')}; "
|
|
f"seed={entry.get('seed')}; result={entry.get('result')}; decision={entry.get('decision')}"
|
|
)
|
|
seed = entry.get("seed")
|
|
return (
|
|
summary,
|
|
str(entry.get("baseline_image") or ""),
|
|
str(entry.get("candidate_image") or ""),
|
|
json.dumps(entry, ensure_ascii=True, sort_keys=True),
|
|
int(seed) if isinstance(seed, int) else -1,
|
|
str(entry.get("decision") or ""),
|
|
)
|
|
|
|
|
|
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,
|
|
"SxCPKrea2VariantEvidence": SxCPKrea2VariantEvidence,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"SxCPHardcorePositionPool": "SxCP Hardcore Position Pool",
|
|
"SxCPHardcoreActionFilter": "SxCP Hardcore Action Filter",
|
|
"SxCPKrea2PoseVariant": "SxCP Krea2 Pose Variant",
|
|
"SxCPKrea2VariantEvidence": "SxCP Krea2 Variant Evidence",
|
|
}
|