104 lines
3.0 KiB
Python
104 lines
3.0 KiB
Python
from __future__ import annotations
|
|
|
|
import random
|
|
from typing import Any
|
|
|
|
try:
|
|
from . import cast_context as cast_context_policy
|
|
from . import character_appearance as character_appearance_policy
|
|
from . import generate_prompt_batches as g
|
|
except ImportError: # Allows local smoke tests with top-level imports.
|
|
import cast_context as cast_context_policy
|
|
import character_appearance as character_appearance_policy
|
|
import generate_prompt_batches as g
|
|
|
|
|
|
def couple_context(
|
|
rng: random.Random,
|
|
women_count: int,
|
|
men_count: int,
|
|
) -> dict[str, str]:
|
|
primary_subject, subject_phrase, pose, effective_women_count, effective_men_count = cast_context_policy.couple_type_from_counts(
|
|
rng,
|
|
women_count,
|
|
men_count,
|
|
choose=g.choose,
|
|
couple_types=g.COUPLE_TYPES,
|
|
)
|
|
return {
|
|
"subject_type": "couple",
|
|
"subject": primary_subject,
|
|
"subject_phrase": subject_phrase,
|
|
"age": g.choose(rng, g.COUPLE_AGES),
|
|
"body": g.choose(rng, ["slim and average", "curvy and broad", "stocky and curvy", "average and athletic"]),
|
|
"skin": "",
|
|
"hair": "",
|
|
"eyes": "",
|
|
"body_phrase": "",
|
|
"fallback_pose": pose,
|
|
"women_count": str(effective_women_count),
|
|
"men_count": str(effective_men_count),
|
|
"person_count": "2",
|
|
}
|
|
|
|
|
|
def group_context(rng: random.Random, ethnicity: str) -> dict[str, str]:
|
|
eth = "Asian " if ethnicity == "asian" else ""
|
|
return {
|
|
"subject_type": "group",
|
|
"subject": f"mixed {eth}adult group",
|
|
"subject_phrase": f"A mixed {eth}adult group of women and men",
|
|
"age": g.choose(rng, g.GROUP_AGES),
|
|
"body": "diverse",
|
|
"skin": "",
|
|
"hair": "",
|
|
"eyes": "",
|
|
"body_phrase": "diverse adult body types",
|
|
}
|
|
|
|
|
|
def layout_context(subject_type: str) -> dict[str, str]:
|
|
return {
|
|
"subject_type": subject_type,
|
|
"subject": "layout scene",
|
|
"subject_phrase": "Adult layout scene",
|
|
"age": "adult",
|
|
"body": "varied",
|
|
"skin": "",
|
|
"hair": "",
|
|
"eyes": "",
|
|
"body_phrase": "varied adult figures",
|
|
}
|
|
|
|
|
|
def subject_context(
|
|
rng: random.Random,
|
|
subject_type: str,
|
|
ethnicity: str,
|
|
figure: str,
|
|
no_plus_women: bool,
|
|
no_black: bool,
|
|
women_count: int = 1,
|
|
men_count: int = 1,
|
|
) -> dict[str, str]:
|
|
if subject_type in ("woman", "man", "single_any"):
|
|
return character_appearance_policy.appearance_for_subject(
|
|
rng,
|
|
subject_type,
|
|
ethnicity,
|
|
figure,
|
|
no_plus_women,
|
|
no_black,
|
|
)
|
|
|
|
if subject_type == "configured_cast":
|
|
return cast_context_policy.configured_cast_context(women_count, men_count)
|
|
|
|
if subject_type == "couple":
|
|
return couple_context(rng, women_count, men_count)
|
|
|
|
if subject_type == "group":
|
|
return group_context(rng, ethnicity)
|
|
|
|
return layout_context(subject_type)
|