Extract subject context policy

2026-06-27 08:41:13 +02:00
parent 70a8698cbe
commit d9275f5f0c
5 changed files with 140 additions and 56 deletions
@@ -0,0 +1,103 @@
+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)