269 lines
9.9 KiB
Python
269 lines
9.9 KiB
Python
from __future__ import annotations
|
|
|
|
import random
|
|
from typing import Any
|
|
|
|
try:
|
|
from . import character_config as character_policy
|
|
from . import character_profile as character_profile_policy
|
|
from . import character_slot as character_slot_policy
|
|
from . import generate_prompt_batches as g
|
|
from . import seed_config as seed_policy
|
|
except ImportError: # Allows local smoke tests with top-level imports.
|
|
import character_config as character_policy
|
|
import character_profile as character_profile_policy
|
|
import character_slot as character_slot_policy
|
|
import generate_prompt_batches as g
|
|
import seed_config as seed_policy
|
|
|
|
|
|
def _choose(rng: random.Random, items: list[Any]) -> Any:
|
|
return items[rng.randrange(len(items))]
|
|
|
|
|
|
def slot_softcore_outfit(slot: dict[str, Any] | None, rng: random.Random | None = None) -> str:
|
|
if not slot:
|
|
return ""
|
|
outfit = character_policy.slot_value(slot.get("softcore_outfit"))
|
|
if outfit:
|
|
return outfit
|
|
if rng is None:
|
|
return ""
|
|
return character_policy.characteristic_choice(
|
|
character_policy.parse_characteristics_config(slot.get("characteristics")),
|
|
"softcore_outfits",
|
|
rng,
|
|
)
|
|
|
|
|
|
def slot_hardcore_clothing(slot: dict[str, Any] | None, rng: random.Random | None = None) -> str:
|
|
if not slot:
|
|
return ""
|
|
clothing = character_policy.slot_value(slot.get("hardcore_clothing"))
|
|
if clothing:
|
|
return clothing
|
|
if rng is None:
|
|
return ""
|
|
return character_policy.characteristic_choice(
|
|
character_policy.parse_characteristics_config(slot.get("characteristics")),
|
|
"hardcore_clothing",
|
|
rng,
|
|
)
|
|
|
|
|
|
def hair_descriptor_from_slot(base_hair: Any, slot: dict[str, Any], rng: random.Random) -> str:
|
|
hair_config = character_policy.parse_hair_config(slot.get("hair_config"))
|
|
color_choice = character_policy.normalize_hair_choice(slot.get("hair_color"), character_policy.CHARACTER_HAIR_COLOR_CHOICES)
|
|
length_choice = character_policy.normalize_hair_choice(slot.get("hair_length"), character_policy.CHARACTER_HAIR_LENGTH_CHOICES)
|
|
style_choice = character_policy.normalize_hair_choice(slot.get("hair_style"), character_policy.CHARACTER_HAIR_STYLE_CHOICES)
|
|
color_options = hair_config.get("colors") or []
|
|
length_options = hair_config.get("lengths") or []
|
|
style_options = hair_config.get("styles") or []
|
|
if (
|
|
color_choice == "random"
|
|
and length_choice == "random"
|
|
and style_choice == "random"
|
|
and not color_options
|
|
and not length_options
|
|
and not style_options
|
|
):
|
|
return ""
|
|
if color_choice != "random":
|
|
color_key = color_choice
|
|
elif color_options:
|
|
color_key = _choose(rng, color_options)
|
|
else:
|
|
color_key = character_policy.infer_hair_color_key(base_hair)
|
|
|
|
if length_choice != "random":
|
|
length_key = length_choice
|
|
elif length_options:
|
|
length_key = _choose(rng, length_options)
|
|
else:
|
|
length_key = character_policy.infer_hair_length_key(base_hair)
|
|
|
|
if style_choice != "random":
|
|
style_key = style_choice
|
|
elif style_options:
|
|
style_key = _choose(rng, style_options)
|
|
else:
|
|
style_key = character_policy.infer_hair_style_key(base_hair)
|
|
if color_key == "random":
|
|
color_key = character_policy.choose_hair_key(rng, character_policy.CHARACTER_HAIR_COLOR_CHOICES)
|
|
if length_key == "random":
|
|
length_key = character_policy.choose_hair_key(rng, character_policy.CHARACTER_HAIR_LENGTH_CHOICES)
|
|
if style_key == "random":
|
|
style_key = character_policy.choose_hair_key(rng, character_policy.CHARACTER_HAIR_STYLE_CHOICES)
|
|
if length_key == "updo" and style_key not in ("ponytail", "braid", "braids", "bun", "messy_bun", "locs", "twists"):
|
|
style_key = _choose(rng, ["ponytail", "braid", "bun", "messy_bun"])
|
|
return character_policy.hair_phrase_from_parts(color_key, length_key, style_key)
|
|
|
|
|
|
def appearance_for_subject(
|
|
rng: random.Random,
|
|
subject_type: str,
|
|
ethnicity: str,
|
|
figure: str,
|
|
no_plus_women: bool,
|
|
no_black: bool,
|
|
) -> dict[str, str]:
|
|
if subject_type == "single_any":
|
|
subject_type = "woman" if rng.random() < 0.82 else "man"
|
|
|
|
if subject_type == "man":
|
|
men_ethnicity = ethnicity if ethnicity else "any"
|
|
subject, age, body, skin, hair, eyes = g.choose(rng, g.by_ethnicity(g.MEN, men_ethnicity))
|
|
return {
|
|
"subject_type": "man",
|
|
"subject": subject,
|
|
"subject_phrase": subject,
|
|
"age": age,
|
|
"body": body,
|
|
"skin": skin,
|
|
"hair": hair,
|
|
"eyes": eyes,
|
|
"body_phrase": f"{body} figure",
|
|
}
|
|
|
|
subject, age, body, skin, hair, eyes = g.choose_woman(rng, ethnicity, no_plus_women, no_black)
|
|
figure_note = g.choose(rng, g.figure_pool(figure))
|
|
return {
|
|
"subject_type": "woman",
|
|
"subject": subject,
|
|
"subject_phrase": subject,
|
|
"age": age,
|
|
"body": body,
|
|
"skin": skin,
|
|
"hair": hair,
|
|
"eyes": eyes,
|
|
"body_phrase": character_profile_policy.body_phrase(body, figure_note),
|
|
"figure": figure_note,
|
|
}
|
|
|
|
|
|
def context_from_character_slot(
|
|
rng: random.Random,
|
|
slot: dict[str, Any],
|
|
subject_type: str,
|
|
ethnicity: str,
|
|
figure: str,
|
|
no_plus_women: bool,
|
|
no_black: bool,
|
|
) -> dict[str, Any]:
|
|
slot_ethnicity = character_policy.slot_value(slot.get("ethnicity"))
|
|
slot_body = character_policy.slot_value(slot.get("body"))
|
|
effective_ethnicity = slot_ethnicity or ethnicity
|
|
effective_figure = character_slot_policy.slot_effective_figure(slot, subject_type, figure)
|
|
effective_no_plus = bool(no_plus_women) and not slot_body
|
|
effective_no_black = bool(no_black) and not slot_ethnicity
|
|
appearance_rng = character_slot_policy.slot_context_rng(slot, rng)
|
|
context = appearance_for_subject(
|
|
appearance_rng,
|
|
subject_type,
|
|
effective_ethnicity,
|
|
effective_figure,
|
|
effective_no_plus,
|
|
effective_no_black,
|
|
)
|
|
|
|
characteristics = character_policy.parse_characteristics_config(slot.get("characteristics"))
|
|
age = character_policy.slot_value(slot.get("age")) or character_policy.characteristic_choice(characteristics, "ages", appearance_rng)
|
|
body_phrase = character_policy.slot_value(slot.get("body_phrase"))
|
|
if not slot_body:
|
|
slot_body = character_policy.characteristic_choice(characteristics, "bodies", appearance_rng)
|
|
if age:
|
|
context["age"] = age
|
|
if slot_body:
|
|
context["body"] = slot_body
|
|
if subject_type == "woman":
|
|
context["body_phrase"] = character_profile_policy.body_phrase(slot_body, context.get("figure", ""))
|
|
else:
|
|
context["body_phrase"] = f"{slot_body} figure"
|
|
if body_phrase:
|
|
context["body_phrase"] = body_phrase
|
|
skin_value = character_policy.slot_value(slot.get("skin"))
|
|
if skin_value:
|
|
context["skin"] = skin_value
|
|
eyes_value = character_policy.slot_value(slot.get("eyes"))
|
|
if not eyes_value:
|
|
eyes_value = character_policy.eye_phrase_from_key(character_policy.characteristic_choice(characteristics, "eyes", appearance_rng))
|
|
if eyes_value:
|
|
context["eyes"] = eyes_value
|
|
hair_value = character_policy.slot_value(slot.get("hair"))
|
|
if hair_value:
|
|
context["hair"] = hair_value
|
|
else:
|
|
hair_descriptor = hair_descriptor_from_slot(context.get("hair"), slot, appearance_rng)
|
|
if hair_descriptor:
|
|
context["hair"] = hair_descriptor
|
|
context["descriptor_detail"] = character_policy.normalize_descriptor_detail(slot.get("descriptor_detail"))
|
|
context["presence_mode"] = character_policy.normalize_presence_mode(slot.get("presence_mode"), subject_type)
|
|
context["expression_enabled"] = character_slot_policy.slot_expression_enabled(slot)
|
|
expression_intensity = character_slot_policy.slot_expression_intensity(slot)
|
|
if expression_intensity is not None:
|
|
context["expression_intensity"] = expression_intensity
|
|
context["subject_type"] = subject_type
|
|
context["subject"] = subject_type
|
|
context["subject_phrase"] = subject_type
|
|
return context
|
|
|
|
|
|
def character_context_for_label(
|
|
label: str,
|
|
label_map: dict[str, dict[str, Any]],
|
|
rng: random.Random,
|
|
ethnicity: str,
|
|
figure: str,
|
|
no_plus_women: bool,
|
|
no_black: bool,
|
|
) -> tuple[dict[str, Any], dict[str, Any] | None]:
|
|
subject_type = "man" if label.startswith("Man ") else "woman"
|
|
slot = label_map.get(label)
|
|
if slot:
|
|
return context_from_character_slot(rng, slot, subject_type, ethnicity, figure, no_plus_women, no_black), slot
|
|
return appearance_for_subject(rng, subject_type, ethnicity, figure, no_plus_women, no_black), None
|
|
|
|
|
|
def apply_character_context_to_row(row: dict[str, Any], context: dict[str, Any]) -> dict[str, Any]:
|
|
for key in (
|
|
"subject_type",
|
|
"subject",
|
|
"subject_phrase",
|
|
"age",
|
|
"body",
|
|
"body_phrase",
|
|
"skin",
|
|
"hair",
|
|
"eyes",
|
|
"figure",
|
|
"descriptor_detail",
|
|
"presence_mode",
|
|
"expression_enabled",
|
|
"expression_intensity",
|
|
):
|
|
value = context.get(key)
|
|
if value is not None and value != "":
|
|
row[key] = value
|
|
if context.get("age"):
|
|
row["age_band"] = context["age"]
|
|
return row
|
|
|
|
|
|
def row_from_character_slot(character_slot: str | dict[str, Any] | None) -> dict[str, Any]:
|
|
slots = character_slot_policy.parse_character_cast(character_slot)
|
|
if not slots:
|
|
return {}
|
|
slot = slots[-1]
|
|
if character_slot_policy.slot_seed(slot) >= 0:
|
|
subject_type = str(slot.get("subject_type") or "woman")
|
|
return context_from_character_slot(
|
|
random.Random(seed_policy.row_seed(character_slot_policy.slot_seed(slot), 1, 719)),
|
|
slot,
|
|
subject_type,
|
|
"any",
|
|
"curvy",
|
|
False,
|
|
False,
|
|
)
|
|
return slot
|