from __future__ import annotations from dataclasses import dataclass from typing import Any try: from . import category_library as category_policy from . import row_expression as row_expression_policy from . import row_item as row_item_policy from . import row_pools as row_pool_policy from . import pov_policy from .hardcore_text_cleanup import sanitize_hardcore_environment_anchors except ImportError: # Allows local smoke tests from the repository root. import category_library as category_policy import row_expression as row_expression_policy import row_item as row_item_policy import row_pools as row_pool_policy import pov_policy from hardcore_text_cleanup import sanitize_hardcore_environment_anchors @dataclass(frozen=True) class PromptAxesRoute: scene_slug: str scene: str scene_entry: dict[str, Any] pose: str expression: str shared_expression: str character_expressions: list[str] character_expression_text: str source_composition: str composition: str composition_entry: dict[str, Any] def as_dict(self) -> dict[str, Any]: return { "scene_slug": self.scene_slug, "scene": self.scene, "scene_entry": dict(self.scene_entry), "pose": self.pose, "expression": self.expression, "shared_expression": self.shared_expression, "character_expressions": list(self.character_expressions), "character_expression_text": self.character_expression_text, "source_composition": self.source_composition, "composition": self.composition, "composition_entry": dict(self.composition_entry), } def _metadata_entry(value: Any, *, slug: str = "", text: str = "") -> dict[str, Any]: if isinstance(value, dict): entry = dict(value) elif isinstance(value, (list, tuple)) and len(value) == 2: entry = {"slug": str(value[0]), "prompt": str(value[1])} else: entry = {"prompt": str(value or "")} if slug: entry["slug"] = slug if text: if "prompt" in entry: entry["prompt"] = text elif "text" in entry: entry["text"] = text else: entry["prompt"] = text return entry def resolve_prompt_axes_result( *, category: dict[str, Any], subcategory: dict[str, Any], item: Any, subject_type: str, context: dict[str, Any], poses: str, women_count: int, men_count: int, scene_rng: Any, pose_rng: Any, expression_rng: Any, composition_rng: Any, expression_disabled: bool, expression_intensity: float, character_slots: list[dict[str, Any]] | None = None, character_slot_map: dict[str, dict[str, Any]] | None = None, expression_phase: str = "", source_role_graph: Any = "", item_axis_values: dict[str, Any] | None = None, is_pose_category: bool = False, pov_character_labels: list[str] | None = None, location_config: dict[str, Any] | None = None, composition_config: dict[str, Any] | None = None, ) -> PromptAxesRoute: character_slots = character_slots or [] character_slot_map = character_slot_map or {} pov_character_labels = pov_character_labels or [] scene_entries = category_policy.compatible_entries( row_pool_policy.scene_pool(category, subcategory, item, subject_type, location_config), women_count, men_count, ) scene_choice = row_item_policy.weighted_choice(scene_rng, scene_entries) scene_slug, scene = row_item_policy.pair_from(scene_choice) scene_entry = _metadata_entry(scene_choice, slug=scene_slug, text=scene) pose = str( category_policy.merged_field(category, subcategory, item, "pose", "") or context.get("fallback_pose") or row_item_policy.choose_text( pose_rng, category_policy.compatible_entries( row_pool_policy.pose_pool(category, subcategory, item, subject_type, poses), women_count, men_count, ), ) ) if is_pose_category: pose = sanitize_hardcore_environment_anchors(pose) expression_pool = row_pool_policy.expression_pool(category, subcategory, item) if expression_disabled: expression = "" else: expression_entries = category_policy.compatible_entries( row_expression_policy.expression_entries_for_intensity(expression_pool, expression_intensity), women_count, men_count, ) expression = row_item_policy.choose_text(expression_rng, expression_entries) if subject_type in ("couple", "group") and ";" not in expression: secondary_expression = row_item_policy.choose_distinct_text(expression_rng, expression_entries, expression) if secondary_expression: expression = f"{expression}; {secondary_expression}" shared_expression = expression character_expressions: list[str] = [] character_expression_text = "" if not expression_disabled and subject_type == "configured_cast" and character_slots: character_expressions = row_expression_policy.character_expression_entries( expression_rng, expression_pool, expression_intensity, character_slot_map, women_count, men_count, expression_phase, ) character_expression_text = "; ".join(character_expressions) character_expression_text = row_expression_policy.sanitize_character_expression_text_for_action( character_expression_text, source_role_graph, item, item_axis_values or {}, ) character_expressions = [part.strip() for part in character_expression_text.split(";") if part.strip()] if character_expression_text: expression = character_expression_text composition_entries = category_policy.compatible_entries( row_pool_policy.composition_pool(category, subcategory, item, subject_type, composition_config), women_count, men_count, ) composition_choice = row_item_policy.weighted_choice(composition_rng, composition_entries) source_composition = row_item_policy.item_text(composition_choice) composition_entry = _metadata_entry(composition_choice, text=source_composition) if is_pose_category: source_composition = sanitize_hardcore_environment_anchors(source_composition) composition_entry["prompt"] = source_composition composition = pov_policy.pov_composition_prompt(source_composition, pov_character_labels) return PromptAxesRoute( scene_slug=scene_slug, scene=scene, scene_entry=scene_entry, pose=pose, expression=expression, shared_expression=shared_expression, character_expressions=character_expressions, character_expression_text=character_expression_text, source_composition=source_composition, composition=composition, composition_entry=composition_entry, ) def resolve_prompt_axes( *, category: dict[str, Any], subcategory: dict[str, Any], item: Any, subject_type: str, context: dict[str, Any], poses: str, women_count: int, men_count: int, scene_rng: Any, pose_rng: Any, expression_rng: Any, composition_rng: Any, expression_disabled: bool, expression_intensity: float, character_slots: list[dict[str, Any]] | None = None, character_slot_map: dict[str, dict[str, Any]] | None = None, expression_phase: str = "", source_role_graph: Any = "", item_axis_values: dict[str, Any] | None = None, is_pose_category: bool = False, pov_character_labels: list[str] | None = None, location_config: dict[str, Any] | None = None, composition_config: dict[str, Any] | None = None, ) -> dict[str, Any]: return resolve_prompt_axes_result( category=category, subcategory=subcategory, item=item, subject_type=subject_type, context=context, poses=poses, women_count=women_count, men_count=men_count, scene_rng=scene_rng, pose_rng=pose_rng, expression_rng=expression_rng, composition_rng=composition_rng, expression_disabled=expression_disabled, expression_intensity=expression_intensity, character_slots=character_slots, character_slot_map=character_slot_map, expression_phase=expression_phase, source_role_graph=source_role_graph, item_axis_values=item_axis_values, is_pose_category=is_pose_category, pov_character_labels=pov_character_labels, location_config=location_config, composition_config=composition_config, ).as_dict()