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 pose: str expression: str shared_expression: str character_expressions: list[str] character_expression_text: str source_composition: str composition: str def as_dict(self) -> dict[str, Any]: return { "scene_slug": self.scene_slug, "scene": self.scene, "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, } 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_slug, scene = row_item_policy.choose_pair( scene_rng, category_policy.compatible_entries( row_pool_policy.scene_pool(category, subcategory, item, subject_type, location_config), women_count, men_count, ), ) 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 source_composition = row_item_policy.choose_text( composition_rng, category_policy.compatible_entries( row_pool_policy.composition_pool(category, subcategory, item, subject_type, composition_config), women_count, men_count, ), ) if is_pose_category: source_composition = sanitize_hardcore_environment_anchors(source_composition) composition = pov_policy.pov_composition_prompt(source_composition, pov_character_labels) return PromptAxesRoute( scene_slug=scene_slug, scene=scene, pose=pose, expression=expression, shared_expression=shared_expression, character_expressions=character_expressions, character_expression_text=character_expression_text, source_composition=source_composition, composition=composition, ) 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()