ComfyUI-Ethanfel-Prompt-Builder/krea_action_positions.py

from __future__ import annotations

import re
from typing import Any

try:
    from . import outercourse_action_policy as outercourse_policy
    from .krea_action_context import (
        axis_values_text,
        is_close_foreplay_text,
        is_foreplay_text,
        is_outercourse_text,
        is_toy_assisted_double_text,
        position_context_text,
    )
except ImportError:  # Allows local smoke tests with `python -c`.
    import outercourse_action_policy as outercourse_policy
    from krea_action_context import (
        axis_values_text,
        is_close_foreplay_text,
        is_foreplay_text,
        is_outercourse_text,
        is_toy_assisted_double_text,
        position_context_text,
    )


def _clean(value: Any) -> str:
    text = "" if value is None else str(value)
    text = text.replace("\n", " ")
    text = re.sub(r"\s+", " ", text).strip()
    text = re.sub(r"\s+([,.;:])", r"\1", text)
    return text


def mentions_rear_entry(text: str) -> bool:
    return bool(
        re.search(
            r"ass[- ](?:up|raised|exposed|lifted|stretched)|penis entering ass|cum (?:on|dripping from) ass|spread cheeks|lower back and ass|pussy, ass|rear[- ]entry",
            text,
        )
    )


def hardcore_pose_anchor(role_graph: str, hard_item: str, composition: str = "", axis_values: Any = None) -> str:
    text = position_context_text(role_graph, hard_item, composition, axis_values)
    item_text = " ".join(part for part in (_clean(hard_item).lower(), axis_values_text(axis_values).lower()) if part)
    position_text = ""
    if isinstance(axis_values, dict):
        position_text = _clean(axis_values.get("position", "")).lower()
    if not text:
        return ""
    if is_foreplay_text(role_graph, hard_item, composition, axis_values_text(axis_values)):
        return ""
    if is_outercourse_text(role_graph, hard_item, composition, axis_values_text(axis_values)):
        action_kind = outercourse_policy.infer_outercourse_action_kind(position_text)
        if action_kind == outercourse_policy.OUTERCOURSE_GENERIC:
            action_kind = outercourse_policy.infer_outercourse_action_kind(text)
        if action_kind == outercourse_policy.OUTERCOURSE_BOOBJOB:
            return "breast-sex outercourse pose"
        if action_kind == outercourse_policy.OUTERCOURSE_TESTICLE:
            return "testicle-sucking outercourse pose"
        if action_kind == outercourse_policy.OUTERCOURSE_PENIS_LICKING:
            return "penis-licking outercourse pose"
        if action_kind == outercourse_policy.OUTERCOURSE_HANDJOB:
            return "handjob outercourse pose"
        if action_kind == outercourse_policy.OUTERCOURSE_FOOTJOB:
            return "footjob outercourse pose"
        return "non-penetrative outercourse pose"
    if is_toy_assisted_double_text(role_graph, hard_item, composition, axis_values_text(axis_values)):
        if "face-down ass-up" in text or "face-down" in text:
            return "toy-assisted face-down rear-entry double-penetration pose"
        if "doggy style" in text or "doggy-style" in text or "all fours" in text or "rear-entry" in text:
            return "toy-assisted rear-entry double-penetration pose"
        if "bent-over" in text or "bent forward" in text:
            return "toy-assisted bent-over double-penetration pose"
        if "spooning anal" in text or "side-lying anal" in text or "side-lying" in text:
            return "toy-assisted side-lying double-penetration pose"
        if "edge-supported" in text or "bed-edge" in text or "edge-of-bed" in text:
            return "toy-assisted edge-supported double-penetration pose"
        if "standing anal" in text or "standing supported" in text or "standing" in text:
            return "toy-assisted standing double-penetration pose"
        if "kneeling anal" in text or "kneeling" in text:
            return "toy-assisted kneeling rear-entry double-penetration pose"
        return "toy-assisted rear-entry double-penetration pose"
    if "double penetration" in text or "vaginal and anal penetration" in text or "front-and-back" in text:
        if "face-down ass-up" in text:
            return "face-down rear-entry double-penetration pose"
        if "doggy style" in text or "doggy-style" in text:
            return "doggy-style double-penetration pose"
        if "bent-over" in text:
            return "bent-over double-penetration pose"
        if "spooning anal" in text or "side-lying anal" in text:
            return "side-lying double-penetration pose"
        if "bed-edge" in text or "edge-of-bed" in text:
            return "bed-edge front-and-back double-penetration pose"
        if "standing anal" in text or "standing supported" in text:
            return "standing supported front-and-back double-penetration pose"
        if "kneeling anal" in text:
            return "kneeling rear-entry double-penetration pose"
        if "standing supported" in text:
            return "standing supported front-and-back double-penetration pose"
        if "kneeling" in text:
            return "kneeling front-and-back double-penetration pose"
        return "front-and-back double-penetration pose"
    if "sixty-nine" in position_text or ("sixty-nine" in text and not position_text):
        return "sixty-nine oral pose"
    if "face-sitting" in position_text or ("face-sitting" in text and not position_text):
        return "face-sitting oral pose"
    if "side-lying oral" in position_text or (("side-lying oral position" in item_text or "side-lying oral" in text) and not position_text):
        return "side-lying oral pose"
    if (
        "edge-of-bed oral" in position_text
        or "edge-supported oral" in position_text
        or (("edge-of-bed oral position" in item_text or "edge-of-bed oral" in text or "edge-supported oral" in text) and not position_text)
    ):
        return "edge-supported oral pose"
    if "standing oral" in position_text or (("standing oral position" in item_text or "standing oral" in text) and not position_text):
        return "standing oral pose"
    if "chair oral" in position_text or (("chair oral position" in item_text or "chair oral" in text) and not position_text):
        return "chair oral pose"
    if "kneeling oral" in position_text or (("kneeling oral position" in item_text or "kneeling oral" in text) and not position_text):
        return "kneeling oral pose"
    if "straddled oral" in position_text or (("straddled oral position" in item_text or "straddled oral" in text) and not position_text):
        return "straddled cunnilingus pose"
    if "reclining cunnilingus" in position_text or (("reclining cunnilingus position" in item_text or "reclining cunnilingus" in text) and not position_text):
        return "reclining cunnilingus pose"
    if "spread-leg oral" in position_text or (("spread-leg oral position" in item_text or "spread-leg oral" in text) and not position_text):
        return "spread-leg oral pose"
    if "cunnilingus" in text or "pussy licking" in text or "mouth on her pussy" in text:
        if "reclining" in text:
            return "reclining cunnilingus pose"
        if "straddled" in text:
            return "straddled cunnilingus pose"
        return "open-thigh cunnilingus pose"
    if "oral" in text or "blowjob" in text or "penis in her mouth" in text or "penis in mouth" in text:
        if "side-lying oral position" in item_text:
            return "side-lying oral pose"
        if "spread-leg oral position" in item_text:
            return "spread-leg oral pose"
        if "edge-of-bed oral position" in item_text:
            return "edge-supported oral pose"
        if "standing oral position" in item_text:
            return "standing oral pose"
        if "chair oral position" in item_text:
            return "chair oral pose"
        if "kneeling oral position" in item_text or "kneeling" in text:
            return "kneeling oral pose"
        if "standing" in text:
            return "standing oral pose"
        if "side-lying" in text:
            return "side-lying oral pose"
        if "edge-of-bed" in text or "bed-edge" in text:
            return "edge-supported oral pose"
        if "spread-leg" in text:
            return "spread-leg oral pose"
        if "chair oral" in text:
            return "chair oral pose"
        return "mouth-to-genitals oral pose"
    if "anal" in text or mentions_rear_entry(text) or "rear-entry" in text:
        if "face-down ass-up" in text:
            return "face-down ass-up rear-entry anal pose"
        if "doggy style" in text or "doggy-style" in text:
            return "doggy-style anal pose"
        if "bed-edge" in text or "edge-of-bed" in text:
            return "bed-edge rear-entry anal pose"
        if "bent-over" in text:
            return "bent-over rear-entry anal pose"
        if "spooning anal" in text or "side-lying anal" in text:
            return "side-lying rear-entry anal pose"
        if "kneeling anal" in text:
            return "kneeling rear-entry anal pose"
        if "standing anal" in text:
            return "standing rear-entry anal pose"
        if "doggy" in text:
            return "doggy-style anal pose"
        return "rear-entry anal pose"
    if "edge-supported" in text or "raised edge" in text or "edge-of-bed" in text or "bed-edge" in text:
        return "edge-supported penetrative sex pose"
    positions = (
        "missionary",
        "reverse cowgirl",
        "cowgirl",
        "doggy style",
        "standing sex",
        "spooning sex",
        "edge-of-bed",
        "kneeling straddle",
        "lotus",
        "bent-over",
    )
    for position in positions:
        if position in text:
            return f"{position.replace('doggy style', 'doggy-style')} pose"
    if "threesome" in text or "three-body" in text:
        return "three-body explicit sex pose"
    if "group" in text or "orgy" in text:
        return "multi-body explicit sex pose"
    if re.search(r"(?<!non-)penetrat|thrust", text):
        return "hip-aligned penetrative sex pose"
    return ""


def hardcore_pose_arrangement(anchor: str, role_graph: str, hard_item: str, composition: str = "", axis_values: Any = None) -> str:
    text = position_context_text(anchor, f"{role_graph} {hard_item}", composition, axis_values)
    position_text = ""
    if isinstance(axis_values, dict):
        position_text = _clean(axis_values.get("position", "")).lower()
    if not text:
        return ""
    mixed_woman_man = "the woman" in text and "the man" in text
    is_double = "double-penetration" in text or "double penetration" in text

    def cast_phrase(mixed: str, generic: str) -> str:
        return mixed if mixed_woman_man else generic

    def double_tail() -> str:
        return "" if "toy" in text else ", with the second penetration point aligned"

    if "sixty-nine" in position_text or ("sixty-nine" in text and not position_text):
        return cast_phrase(
            "with the woman and man inverted head-to-hips so both mouths align with genitals",
            "with both bodies inverted head-to-hips so both mouths align with genitals",
        )
    if "face-sitting" in position_text or ("face-sitting" in text and not position_text):
        return cast_phrase(
            "with the man lying back while the woman straddles his face",
            "with one partner lying back while the other straddles the face",
        )
    if (
        "reclining cunnilingus" in position_text
        or "spread-leg oral" in position_text
        or (("reclining cunnilingus" in text or "spread-leg oral" in text) and not position_text)
    ):
        if "takes the man's penis" in text or "penis in her mouth" in text:
            return cast_phrase(
                "with the man seated with legs apart and the woman positioned at his hips",
                "with the receiver seated with legs apart and the giver positioned at the hips",
            )
        return cast_phrase(
            "with the woman lying back, thighs spread, and the man positioned between her legs",
            "with the receiving partner lying back, thighs spread, and the giver positioned between the legs",
        )
    if (
        "straddled oral" in position_text
        or (("straddled cunnilingus" in text or "straddled oral" in text) and not position_text)
    ):
        return cast_phrase(
            "with the woman straddling above the man's mouth and her thighs framing his face",
            "with the receiver straddling above the giver's mouth",
        )
    if (
        "edge-of-bed oral" in position_text
        or "edge-supported oral" in position_text
        or ("edge-of-bed oral" in text and not position_text)
        or ("edge-supported oral" in text and not position_text)
    ):
        if "takes the man's penis" in text or "penis in her mouth" in text:
            return cast_phrase(
                "with the man at a raised edge and the woman kneeling at his hips",
                "with the receiver at a raised edge and the giver positioned at hip height",
            )
        return cast_phrase(
            "with the woman lying at a raised edge and the man positioned between her open thighs",
            "with the receiver lying at a raised edge and the giver positioned between open thighs",
        )
    if "standing oral" in position_text or ("standing oral" in text and not position_text):
        if "takes the man's penis" in text or "penis in her mouth" in text:
            return cast_phrase(
                "with the man standing and the woman kneeling in front of his hips",
                "with the receiver standing and the giver kneeling at hip height",
            )
        return cast_phrase(
            "with the woman standing braced and the man kneeling between her thighs",
            "with the receiver standing braced and the giver kneeling between the thighs",
        )
    if "chair oral" in position_text or ("chair oral" in text and not position_text):
        if "takes the man's penis" in text or "penis in her mouth" in text:
            return cast_phrase(
                "with the man seated in the chair and the woman kneeling between his legs at hip level",
                "with the receiver seated in the chair and the giver kneeling between the legs at hip level",
            )
        return cast_phrase(
            "with one partner seated in a chair and the other kneeling between the open thighs",
            "with the receiver seated in a chair and the giver kneeling between the open thighs",
        )
    if "side-lying oral" in position_text or ("side-lying oral" in text and not position_text):
        return "with both bodies lying on their sides and mouth aligned to genitals"
    if "kneeling oral" in position_text or ("kneeling oral" in text and not position_text):
        if "takes the man's penis" in text or "penis in her mouth" in text:
            return cast_phrase(
                "with the woman kneeling in front of the man's hips, her mouth at penis level",
                "with the giver kneeling in front of the receiver's hips",
            )
        if "mouth on her pussy" in text or "uses his mouth on" in text:
            return cast_phrase(
                "with the man kneeling between the woman's open thighs, his mouth at her pussy",
                "with the giver kneeling between the receiver's open thighs",
            )
        return "with the giver kneeling at the receiver's hips"
    if "reverse cowgirl" in text:
        return cast_phrase(
            "with the man lying on his back under the woman while she straddles his hips facing away",
            "with the lower partner lying on their back while the upper partner straddles them facing away",
        )
    if "cowgirl" in text:
        return cast_phrase(
            "with the man lying on his back under the woman while she straddles his hips on top",
            "with the lower partner lying on their back while the upper partner straddles their hips on top",
        )
    if "missionary" in text:
        return cast_phrase(
            "with the woman lying on her back under the man, legs open around his hips",
            "with the receiving partner lying on their back under the penetrating partner, legs open around the hips",
        )
    if "lotus" in text:
        return cast_phrase(
            "with the man seated upright and the woman seated in his lap facing him, legs wrapped around his hips",
            "with one partner seated upright and the other seated in their lap facing them, legs wrapped around the hips",
        )
    if "kneeling straddle" in text:
        return cast_phrase(
            "with the woman straddling the man's kneeling lap, both torsos upright and hips pressed together",
            "with one partner straddling the other's kneeling lap, torsos upright and hips pressed together",
        )
    if "doggy-style" in text:
        return cast_phrase(
            f"with the woman on all fours and the man positioned behind her at hip level{double_tail() if is_double else ''}",
            f"with the receiving partner on all fours and the penetrating partner positioned behind at hip level{double_tail() if is_double else ''}",
        )
    if "face-down" in text:
        return cast_phrase(
            f"with the woman face-down, hips raised, and the man positioned behind her{double_tail() if is_double else ''}",
            f"with the receiving partner face-down, hips raised, and the penetrating partner positioned behind{double_tail() if is_double else ''}",
        )
    if "bent-over" in text:
        return cast_phrase(
            f"with the woman bent forward at the waist and the man positioned behind her{double_tail() if is_double else ''}",
            f"with the receiving partner bent forward at the waist and the penetrating partner positioned behind{double_tail() if is_double else ''}",
        )
    if "spooning" in text or ("side-lying" in text and "oral" not in text):
        return cast_phrase(
            f"with both lying on their sides and the man positioned behind the woman{double_tail() if is_double else ''}",
            f"with both bodies lying on their sides and the penetrating partner positioned behind{double_tail() if is_double else ''}",
        )
    if "edge-of-bed" in text or "bed-edge" in text:
        return cast_phrase(
            f"with the woman lying at the bed edge, hips at the edge, and the man kneeling between her legs{double_tail() if is_double else ''}",
            f"with the receiver lying at the bed edge, hips at the edge, and the penetrating partner kneeling between the legs{double_tail() if is_double else ''}",
        )
    if "standing" in text:
        return cast_phrase(
            f"with the woman braced standing and the man aligned at her hips{double_tail() if is_double else ''}",
            f"with both partners standing and the penetrating partner aligned at the receiver's hips{double_tail() if is_double else ''}",
        )
    if "kneeling" in text and ("anal" in text or "rear-entry" in text):
        return cast_phrase(
            f"with the woman kneeling forward and the man positioned behind her{double_tail() if is_double else ''}",
            f"with the receiving partner kneeling forward and the penetrating partner positioned behind{double_tail() if is_double else ''}",
        )
    if "double-penetration" in text or "double penetration" in text:
        if "toy" in text:
            return cast_phrase(
                "with the woman on all fours and the man positioned behind her at hip level",
                "with the receiving body on all fours and the penetrating partner positioned behind at hip level",
            )
        if "from the front" in text:
            return cast_phrase(
                "with the woman held between the man behind her and a second partner in front",
                "with the receiving body held between one partner behind and a second partner in front",
            )
        return cast_phrase(
            "with the woman held in a front-and-back position so both contact points are visible",
            "with the central body held in a front-and-back position so both contact points are visible",
        )
    if "anal" in text or mentions_rear_entry(text) or "rear-entry" in text:
        return cast_phrase(
            "with the woman's hips raised, ass exposed, and the man positioned behind her",
            "with the receiving partner's hips raised and the penetrating partner positioned behind",
        )
    if "cunnilingus" in text or "mouth on her pussy" in text or "pussy licking" in text:
        return cast_phrase(
            "with the woman's thighs open and the man's mouth pressed to her pussy",
            "with the receiver's thighs open and the giver's mouth pressed to genitals",
        )
    if "oral" in text or "blowjob" in text or "penis in her mouth" in text or "penis in mouth" in text:
        if "takes the man's penis in her mouth" in text or "penis in her mouth" in text:
            return cast_phrase(
                "with the woman's mouth at the man's hips",
                "with the giver's mouth positioned at the receiver's hips",
            )
        return "with mouth and genitals aligned clearly"
    if "threesome" in text or "three-body" in text:
        return "with all three adult bodies clearly placed around the central subject"
    if "group" in text or "orgy" in text:
        return "with each adult body readable in the shared sex act"
    if re.search(r"(?<!non-)penetrat|thrust", text):
        return "with hips aligned and legs open around the contact point"
    return ""


def arrangement_duplicates_role(arrangement: str, role_graph: str) -> bool:
    arrangement_lower = _clean(arrangement).lower()
    role_lower = _clean(role_graph).lower()
    if not arrangement_lower or not role_lower:
        return False
    markers = (
        "bed edge",
        "on all fours",
        "face-down",
        "hips raised",
        "bent forward",
        "straddl",
        "on her back",
        "on their sides",
        "on her side",
        "seated in",
        "sits in",
        "lap",
        "kneeling between",
        "kneels between",
        "kneeling in front",
        "kneels in front",
        "positioned behind",
        "standing",
    )
    return any(marker in arrangement_lower and marker in role_lower for marker in markers)


def action_position_phrase(action: str) -> str:
    action = _clean(action).lower()
    if is_close_foreplay_text(action):
        return "single-frame close-body first-person position"
    if "pov reverse cowgirl" in action:
        return "reverse-cowgirl first-person position"
    if "pov cowgirl" in action:
        return "cowgirl first-person position"
    if "pov missionary" in action:
        return "missionary first-person position"
    if "pov raised-edge" in action or "raised edge" in action:
        return "raised-edge open-thigh position"
    if "pov doggy" in action or "on all fours" in action:
        return "all-fours rear-entry position"
    if "pov bent-over" in action or "bent forward" in action:
        return "bent-over rear-entry position"
    if "pov face-down" in action:
        return "face-down rear-entry position"
    if "pov standing" in action:
        return "standing rear-entry position"
    if "pov side-lying" in action:
        return "side-lying position"
    if "pov lotus" in action:
        return "lap-straddling position"
    if "face-down" in action and "ass raised" in action:
        return "face-down raised-hip position"
    if "on all fours" in action:
        return "all-fours raised-hip position"
    if "bends forward" in action or "bent forward" in action:
        return "bent-over raised-hip position"
    if "lies on her back" in action and ("thighs open" in action or "legs open" in action):
        return "open-thigh reclined position"
    if "lies at the bed edge" in action or "bed edge" in action:
        return "bed-edge position"
    if "lies on her side" in action:
        return "side-lying position"
    if "kneels in front" in action:
        return "kneeling-at-hip-height position"
    if "straddles" in action or "squats over" in action:
        return "straddling position"
    if "sits in the man's lap" in action:
        return "lap-straddling position"
    if "stands braced" in action:
        return "standing braced position"
    if "held between" in action or "front-and-back" in action:
        return "front-and-back position"
    if "lies between" in action:
        return "between-partners position"
    return ""