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

152 lines
4.6 KiB
Python

from __future__ import annotations
import re
from typing import Any
try:
from . import formatter_detail as detail_policy
from . import formatter_input as input_policy
from . import route_metadata as route_metadata_policy
except ImportError: # Allows local smoke tests with `python tools/prompt_smoke.py`.
import formatter_detail as detail_policy
import formatter_input as input_policy
import route_metadata as route_metadata_policy
OLD_TRIGGER = "sxcpinup_coloredpencil"
DEFAULT_TRIGGER = "sxcppnl7"
DETAIL_LEVELS = detail_policy.DETAIL_LEVELS
STYLE_POLICIES = ("drop_style_tail", "keep_style_terms")
CAPTION_PROFILE_DEFAULT = "manual_controls"
CAPTION_PROFILES = {
"manual_controls": {},
"training_concise": {
"detail_level": "concise",
"style_policy": "drop_style_tail",
"include_trigger": True,
},
"training_dense": {
"detail_level": "dense",
"style_policy": "drop_style_tail",
"include_trigger": True,
},
"browsing": {
"detail_level": "balanced",
"style_policy": "keep_style_terms",
"include_trigger": False,
},
}
STYLE_TAILS = [
", coloured pencil comic illustration, crisp linework, hatching, soft pastel palette, warm sensual lighting, textured parchment paper",
", coloured pencil comic illustration, crisp linework, hatching, soft pastel palette, warm sensual lighting, textured paper",
]
ITEM_LABELS = (
"Sexual pose",
"Erotic outfit",
"Clothing",
)
ACTION_FAMILY_CAPTION_LABELS = {
"foreplay": "foreplay action",
"outercourse": "non-penetrative action",
"oral": "oral action",
"penetration": "penetrative action",
"toy_double": "toy-assisted double-contact action",
"climax": "climax action",
}
POSITION_FAMILY_CAPTION_LABELS = {
"penetrative": "penetrative action",
"foreplay": "foreplay action",
"interaction": "interaction beat",
"manual": "manual action",
"oral": "oral action",
"outercourse": "non-penetrative action",
"anal": "anal action",
"climax": "climax action",
"threesome": "three-person action",
"group": "group action",
}
def normalize_detail_level(value: str) -> str:
return detail_policy.normalize_detail_level(value)
def _choice_key(value: Any) -> str:
return str(value or "").strip().lower().replace("-", "_").replace(" ", "_")
def normalize_style_policy(value: str) -> str:
value = _choice_key(value)
return value if value in STYLE_POLICIES else "drop_style_tail"
def style_policy_choices() -> list[str]:
return list(STYLE_POLICIES)
def caption_profile_choices() -> list[str]:
return list(CAPTION_PROFILES)
def normalize_caption_profile(value: str) -> str:
value = _choice_key(value)
return value if value in CAPTION_PROFILES else CAPTION_PROFILE_DEFAULT
def apply_caption_profile(
caption_profile: str,
*,
detail_level: str,
style_policy: str,
include_trigger: bool,
) -> tuple[str, str, bool]:
profile = CAPTION_PROFILES[normalize_caption_profile(caption_profile)]
return (
normalize_detail_level(profile.get("detail_level", detail_level)),
normalize_style_policy(profile.get("style_policy", style_policy)),
bool(profile.get("include_trigger", include_trigger)),
)
def keep_style_terms(style_policy: str) -> bool:
return normalize_style_policy(style_policy) == "keep_style_terms"
def detail_allows(level: str, dense_only: bool = False) -> bool:
return detail_policy.detail_allows(level, dense_only=dense_only)
def strip_style_tail(text: str) -> str:
text = input_policy.clean_text(text)
for tail in STYLE_TAILS:
if text.endswith(tail):
return text[: -len(tail)].strip(" ,")
return text
def metadata_action_label(row: dict[str, Any], default: str = "sexual pose") -> str:
position_family = route_metadata_policy.row_position_family(row)
if position_family in POSITION_FAMILY_CAPTION_LABELS:
return POSITION_FAMILY_CAPTION_LABELS[position_family]
action_family = route_metadata_policy.row_action_family(row)
if action_family in ACTION_FAMILY_CAPTION_LABELS:
return ACTION_FAMILY_CAPTION_LABELS[action_family]
return default
def normalize_composition(text: str) -> str:
return re.sub(r"^vertical\s+", "", input_policy.clean_text(text), flags=re.IGNORECASE)
def clean_clothing(text: str) -> str:
text = input_policy.clean_text(text)
text = re.sub(r",?\s*fashion editorial styling$", "", text, flags=re.IGNORECASE)
text = re.sub(r",?\s*resort styling$", "", text, flags=re.IGNORECASE)
return text.strip(" ,")