Files
ComfyUI-Ethanfel-Prompt-Bui…/hardcore_role_graphs.py

177 lines
7.1 KiB
Python

from __future__ import annotations
import random
from typing import Any
try:
from . import item_axis_policy
from .hardcore_role_anal import build_anal_or_double_role_graph
from .hardcore_role_climax import build_climax_role_graph
from .hardcore_role_fallback import (
build_men_only_role_graph,
build_mixed_group_fallback_role_graph,
build_solo_role_graph,
build_support_sentence,
build_women_only_role_graph,
)
from .hardcore_role_interaction import (
build_foreplay_role_graph,
build_group_coordination_role_graph,
build_interaction_role_graph,
build_manual_role_graph,
)
from .hardcore_role_oral import build_oral_role_graph
from .hardcore_role_outercourse import build_outercourse_role_graph
from .hardcore_role_penetration import build_penetration_role_graph
except ImportError: # Allows local smoke tests with `python -c`.
import item_axis_policy
from hardcore_role_anal import build_anal_or_double_role_graph
from hardcore_role_climax import build_climax_role_graph
from hardcore_role_fallback import (
build_men_only_role_graph,
build_mixed_group_fallback_role_graph,
build_solo_role_graph,
build_support_sentence,
build_women_only_role_graph,
)
from hardcore_role_interaction import (
build_foreplay_role_graph,
build_group_coordination_role_graph,
build_interaction_role_graph,
build_manual_role_graph,
)
from hardcore_role_oral import build_oral_role_graph
from hardcore_role_outercourse import build_outercourse_role_graph
from hardcore_role_penetration import build_penetration_role_graph
def _lettered(prefix: str, count: int) -> list[str]:
letters = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
return [f"{prefix.capitalize()} {letters[index]}" for index in range(max(0, count))]
def _pick_distinct(rng: random.Random, items: list[str], count: int) -> list[str]:
if not items:
return []
if len(items) >= count:
return rng.sample(items, count)
picked = list(items)
while len(picked) < count:
picked.append(items[rng.randrange(len(items))])
return picked
def _participant_context(women_count: int, men_count: int) -> dict[str, list[str]]:
women = _lettered("woman", women_count)
men = _lettered("man", men_count)
return {"women": women, "men": men, "people": women + men}
def build_hardcore_role_graph(
rng: random.Random,
subcategory: dict[str, Any],
context: dict[str, Any],
item_axis_values: dict[str, Any] | None = None,
pov_labels: list[str] | None = None,
) -> str:
if context.get("subject_type") != "configured_cast":
return ""
women_count = int(context.get("women_count") or 0)
men_count = int(context.get("men_count") or 0)
people_count = women_count + men_count
if people_count <= 0:
return ""
participants = _participant_context(women_count, men_count)
women = participants["women"]
men = participants["men"]
people = participants["people"]
slug = str(subcategory.get("slug") or subcategory.get("name") or "").lower()
item_text = item_axis_policy.context_text(axis_values=item_axis_values)
def any_person(exclude: set[str] | None = None) -> str:
exclude = exclude or set()
pool = [person for person in people if person not in exclude] or people
return rng.choice(pool)
def any_woman(exclude: set[str] | None = None) -> str:
exclude = exclude or set()
pool = [person for person in women if person not in exclude] or [person for person in people if person not in exclude] or people
return rng.choice(pool)
def any_man(exclude: set[str] | None = None) -> str:
exclude = exclude or set()
pool = [person for person in men if person not in exclude] or [person for person in people if person not in exclude] or people
return rng.choice(pool)
if people_count == 1:
return build_solo_role_graph(people[0], women_count, slug, item_text, item_axis_values)
if women_count > 0 and men_count == 0:
a, b = _pick_distinct(rng, women, 2)
c = any_woman({a, b}) if len(women) >= 3 else ""
used = {a, b}
if any(token in slug for token in ("foreplay", "body_worship", "clothing_position", "dominant_guidance", "camera_performance", "aftercare")):
graph = build_interaction_role_graph(a, b, c, slug, item_text, item_axis_values)
if c and "camera_performance" in slug:
used.add(c)
elif "foreplay" in slug:
graph = build_foreplay_role_graph(a, b, item_text, item_axis_values)
else:
graph, used = build_women_only_role_graph(
slug,
a,
b,
c,
c or any_woman({a}),
item_text,
item_axis_values,
)
return graph + build_support_sentence(rng, people, used)
if men_count > 0 and women_count == 0:
a, b = _pick_distinct(rng, men, 2)
c = any_man({a, b}) if len(men) >= 3 else ""
graph, used = build_men_only_role_graph(
slug,
a,
b,
c,
c or any_man({a}),
item_text,
item_axis_values,
)
return graph + build_support_sentence(rng, people, used)
woman = any_woman()
man = any_man()
third = any_person({woman, man}) if people_count >= 3 else ""
if "manual_stimulation" in slug:
graph = build_manual_role_graph(woman, man, item_text, item_axis_values)
elif "group_coordination" in slug:
graph = build_group_coordination_role_graph(
woman,
man,
third,
any_person({woman, man}) if not third else "",
item_text,
item_axis_values,
)
elif any(token in slug for token in ("foreplay", "body_worship", "clothing_position", "dominant_guidance", "camera_performance", "aftercare")):
graph = build_interaction_role_graph(woman, man, third, slug, item_text, item_axis_values)
elif "foreplay" in slug:
graph = build_foreplay_role_graph(woman, man, item_text, item_axis_values)
elif "outercourse" in slug:
graph = build_outercourse_role_graph(woman, man, item_text, item_axis_values, pov_labels)
elif "oral" in slug:
graph = build_oral_role_graph(woman, man, item_text, item_axis_values, pov_labels)
elif "anal" in slug or "double" in slug:
graph = build_anal_or_double_role_graph(woman, man, third, people_count, item_text, item_axis_values)
elif "threesome" in slug or "group" in slug or "orgy" in slug:
graph = build_mixed_group_fallback_role_graph(woman, man, third, any_person({woman, man}), slug)
elif "cumshot" in slug or "climax" in slug:
graph = build_climax_role_graph(woman, man, third, item_text, item_axis_values)
else:
graph = build_penetration_role_graph(woman, man, item_text, item_axis_values)
return graph + build_support_sentence(rng, people, {woman, man, third} if third else {woman, man})