Extract seed config policy

seed_config.py

@@ -0,0 +1,165 @@
+from __future__ import annotations
+
+import json
+import random
+from typing import Any
+
+
+SEED_AXIS_SALTS = {
+    "category": 31,
+    "subcategory": 37,
+    "content": 41,
+    "person": 43,
+    "scene": 47,
+    "pose": 53,
+    "role": 57,
+    "expression": 59,
+    "composition": 61,
+}
+
+SEED_AXIS_ALIASES = {
+    "category": ("category_seed", "category"),
+    "subcategory": ("subcategory_seed", "subcategory"),
+    "content": ("content_seed", "item_seed", "outfit_seed", "sexual_pose_seed", "content"),
+    "person": ("person_seed", "appearance_seed", "cast_seed", "person"),
+    "scene": ("scene_seed", "scene"),
+    "pose": ("pose_seed", "sexual_pose_seed", "pose"),
+    "role": ("role_seed", "role", "pose_seed", "sexual_pose_seed"),
+    "expression": ("expression_seed", "face_seed", "expression"),
+    "composition": ("composition_seed", "camera_seed", "composition"),
+}
+
+SEED_LOCK_AXES = (
+    "category",
+    "subcategory",
+    "content",
+    "person",
+    "scene",
+    "pose",
+    "role",
+    "expression",
+    "composition",
+)
+SEED_MODE_CHOICES = ["auto", "follow_main", "fixed", "random"]
+
+
+def seed_mode_choices() -> list[str]:
+    return list(SEED_MODE_CHOICES)
+
+
+def row_seed(seed: int, row_number: int, salt: int = 0) -> int:
+    return int(seed) + int(row_number) * 1009 + salt * 9176
+
+
+def build_seed_config_json(
+    category_seed: int = -1,
+    subcategory_seed: int = -1,
+    content_seed: int = -1,
+    person_seed: int = -1,
+    scene_seed: int = -1,
+    pose_seed: int = -1,
+    role_seed: int = -1,
+    expression_seed: int = -1,
+    composition_seed: int = -1,
+    category_seed_mode: str = "auto",
+    subcategory_seed_mode: str = "auto",
+    content_seed_mode: str = "auto",
+    person_seed_mode: str = "auto",
+    scene_seed_mode: str = "auto",
+    pose_seed_mode: str = "auto",
+    role_seed_mode: str = "auto",
+    expression_seed_mode: str = "auto",
+    composition_seed_mode: str = "auto",
+) -> str:
+    rng = random.SystemRandom()
+
+    def axis_seed(value: int, mode: str) -> int:
+        mode = mode if mode in SEED_MODE_CHOICES else "auto"
+        if mode == "auto":
+            return int(value)
+        if mode == "random":
+            return rng.randint(0, 0xFFFFFFFF)
+        if mode == "fixed":
+            return max(0, int(value))
+        return -1
+
+    return json.dumps(
+        {
+            "category_seed": axis_seed(category_seed, category_seed_mode),
+            "subcategory_seed": axis_seed(subcategory_seed, subcategory_seed_mode),
+            "content_seed": axis_seed(content_seed, content_seed_mode),
+            "person_seed": axis_seed(person_seed, person_seed_mode),
+            "scene_seed": axis_seed(scene_seed, scene_seed_mode),
+            "pose_seed": axis_seed(pose_seed, pose_seed_mode),
+            "role_seed": axis_seed(role_seed, role_seed_mode),
+            "expression_seed": axis_seed(expression_seed, expression_seed_mode),
+            "composition_seed": axis_seed(composition_seed, composition_seed_mode),
+        },
+        ensure_ascii=True,
+        sort_keys=True,
+    )
+
+
+def build_seed_lock_config_json(
+    base_seed: int = 20260614,
+    reroll_axis: str = "none",
+    reroll_seed: int = -1,
+) -> str:
+    base_seed = int(base_seed)
+    reroll_seed = int(reroll_seed)
+    reroll_groups = {
+        "none": (),
+        "category": ("category",),
+        "subcategory": ("subcategory",),
+        "content": ("content",),
+        "person": ("person",),
+        "scene": ("scene",),
+        "pose": ("pose", "role"),
+        "role": ("role",),
+        "expression": ("expression",),
+        "composition": ("composition",),
+        "content_pose": ("content", "pose", "role"),
+        "scene_pose": ("scene", "pose", "role"),
+    }
+    reroll = set(reroll_groups.get(str(reroll_axis or "none"), ()))
+    config: dict[str, int] = {}
+    for axis in SEED_LOCK_AXES:
+        config[f"{axis}_seed"] = reroll_seed if axis in reroll else base_seed
+    return json.dumps(config, ensure_ascii=True, sort_keys=True)
+
+
+def parse_seed_config(seed_config: str | dict[str, Any] | None) -> dict[str, int]:
+    if not seed_config:
+        return {}
+    if isinstance(seed_config, dict):
+        raw = seed_config
+    else:
+        try:
+            raw = json.loads(str(seed_config))
+        except json.JSONDecodeError as exc:
+            raise ValueError(f"Invalid seed_config JSON: {exc}") from exc
+    if not isinstance(raw, dict):
+        raise ValueError("seed_config must be a JSON object")
+    parsed: dict[str, int] = {}
+    for key, value in raw.items():
+        try:
+            parsed[str(key)] = int(value)
+        except (TypeError, ValueError):
+            continue
+    return parsed
+
+
+def configured_axis_seed(seed_config: dict[str, int], axis: str) -> int | None:
+    for key in SEED_AXIS_ALIASES.get(axis, (axis,)):
+        value = seed_config.get(key)
+        if value is not None and value >= 0:
+            return value
+    return None
+
+
+def axis_rng(seed_config: dict[str, int], axis: str, base_seed: int, row_number: int) -> random.Random:
+    configured = configured_axis_seed(seed_config, axis)
+    salt = SEED_AXIS_SALTS.get(axis, 0)
+    if configured is None:
+        return random.Random(row_seed(base_seed, row_number, salt))
+    return random.Random(row_seed(configured, row_number, salt))