ComfyUI-Ethanfel-Prompt-Builder/row_category_route.py

from __future__ import annotations

from dataclasses import dataclass
from typing import Any

try:
    from . import category_library as category_policy
    from . import category_template_metadata as template_policy
    from . import hardcore_position_config as hardcore_position_policy
    from . import row_item as row_item_policy
    from . import seed_config as seed_policy
    from .hardcore_text_cleanup import (
        sanitize_hardcore_axis_values,
        sanitize_hardcore_environment_anchors,
    )
except ImportError:  # Allows local smoke tests from the repository root.
    import category_library as category_policy
    import category_template_metadata as template_policy
    import hardcore_position_config as hardcore_position_policy
    import row_item as row_item_policy
    import seed_config as seed_policy
    from hardcore_text_cleanup import (
        sanitize_hardcore_axis_values,
        sanitize_hardcore_environment_anchors,
    )


def _list_from(value: Any) -> list[Any]:
    if value is None:
        return []
    if isinstance(value, list):
        return value
    return [value]


def is_pose_content_category(category: dict[str, Any], subcategory: dict[str, Any]) -> bool:
    haystack = " ".join(
        str(value)
        for value in (
            category.get("name", ""),
            category.get("slug", ""),
            category.get("item_label", ""),
            subcategory.get("name", ""),
            subcategory.get("slug", ""),
            subcategory.get("item_label", ""),
        )
    ).lower()
    return "pose" in haystack or "sex" in haystack


def cast_count_adjustment(
    requested_women_count: int,
    requested_men_count: int,
    effective_women_count: int,
    effective_men_count: int,
) -> dict[str, int]:
    if requested_women_count == effective_women_count and requested_men_count == effective_men_count:
        return {}
    return {
        "requested_women_count": requested_women_count,
        "requested_men_count": requested_men_count,
        "effective_women_count": effective_women_count,
        "effective_men_count": effective_men_count,
    }


@dataclass(frozen=True)
class CategoryItemRoute:
    category: dict[str, Any]
    subcategory: dict[str, Any]
    women_count: int
    men_count: int
    count_adjustment: dict[str, int]
    content_axis: str
    item: Any
    item_text: str
    item_name: str
    item_axis_values: dict[str, Any]
    item_template_metadata: dict[str, Any]
    formatter_hints: dict[str, Any]
    is_pose_category: bool

    def as_dict(self) -> dict[str, Any]:
        return {
            "category": self.category,
            "subcategory": self.subcategory,
            "women_count": self.women_count,
            "men_count": self.men_count,
            "count_adjustment": dict(self.count_adjustment),
            "content_axis": self.content_axis,
            "item": self.item,
            "item_text": self.item_text,
            "item_name": self.item_name,
            "item_axis_values": dict(self.item_axis_values),
            "item_template_metadata": dict(self.item_template_metadata),
            "formatter_hints": dict(self.formatter_hints),
            "is_pose_category": self.is_pose_category,
        }


def select_category_item_route_result(
    *,
    category_choice: str,
    subcategory_choice: str,
    seed_config: dict[str, int],
    seed: int,
    row_number: int,
    women_count: int,
    men_count: int,
    hardcore_position_config: dict[str, Any] | None = None,
    categories: list[dict[str, Any]] | None = None,
) -> CategoryItemRoute:
    source_categories = category_policy.load_category_library() if categories is None else categories
    parsed_hardcore_position_config = hardcore_position_config or {}
    requested_women_count = women_count
    requested_men_count = men_count

    category_rng = seed_policy.axis_rng(seed_config, "category", seed, row_number)
    subcategory_rng = seed_policy.axis_rng(seed_config, "subcategory", seed, row_number)
    filtered_categories = hardcore_position_policy.filter_hardcore_categories_for_position(
        source_categories,
        parsed_hardcore_position_config,
        women_count,
        men_count,
        category_policy.compatible_entry,
    )
    category, subcategory, women_count, men_count = category_policy.find_subcategory(
        filtered_categories,
        category_choice,
        subcategory_choice,
        category_rng,
        subcategory_rng,
        women_count,
        men_count,
    )
    count_adjustment = cast_count_adjustment(
        requested_women_count,
        requested_men_count,
        women_count,
        men_count,
    )
    if hardcore_position_policy.is_hardcore_sexual_category(category):
        subcategory = hardcore_position_policy.apply_hardcore_position_config_to_subcategory(
            subcategory,
            parsed_hardcore_position_config,
        )

    is_pose_category = is_pose_content_category(category, subcategory)
    content_axis = "pose" if is_pose_category else "content"
    content_rng = seed_policy.axis_rng(seed_config, content_axis, seed, row_number)
    item = row_item_policy.weighted_choice(content_rng, _list_from(subcategory.get("items", [subcategory["name"]])))
    item_text, item_name, item_axis_values, item_template_metadata = row_item_policy.compose_item(
        content_rng,
        category,
        subcategory,
        item,
        women_count,
        men_count,
    )
    if is_pose_category:
        item_text = sanitize_hardcore_environment_anchors(item_text)
        item_axis_values = sanitize_hardcore_axis_values(item_axis_values)

    return CategoryItemRoute(
        category=category,
        subcategory=subcategory,
        women_count=women_count,
        men_count=men_count,
        count_adjustment=count_adjustment,
        content_axis=content_axis,
        item=item,
        item_text=item_text,
        item_name=item_name,
        item_axis_values=item_axis_values,
        item_template_metadata=item_template_metadata,
        formatter_hints=template_policy.formatter_hints(item_template_metadata),
        is_pose_category=is_pose_category,
    )


def select_category_item_route(
    *,
    category_choice: str,
    subcategory_choice: str,
    seed_config: dict[str, int],
    seed: int,
    row_number: int,
    women_count: int,
    men_count: int,
    hardcore_position_config: dict[str, Any] | None = None,
    categories: list[dict[str, Any]] | None = None,
) -> dict[str, Any]:
    return select_category_item_route_result(
        category_choice=category_choice,
        subcategory_choice=subcategory_choice,
        seed_config=seed_config,
        seed=seed,
        row_number=row_number,
        women_count=women_count,
        men_count=men_count,
        hardcore_position_config=hardcore_position_config,
        categories=categories,
    ).as_dict()