251 lines
7.8 KiB
Python
251 lines
7.8 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import random
|
|
from typing import Any, Iterable
|
|
|
|
|
|
SEED_AXIS_SALTS = {
|
|
"category": 31,
|
|
"subcategory": 37,
|
|
"content": 41,
|
|
"clothing": 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", "sexual_pose_seed", "content"),
|
|
"clothing": ("clothing_seed", "outfit_seed", "wardrobe_seed", "content_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",
|
|
"clothing",
|
|
"person",
|
|
"scene",
|
|
"pose",
|
|
"role",
|
|
"expression",
|
|
"composition",
|
|
)
|
|
SEED_MODE_CHOICES = ["auto", "follow_main", "fixed", "random"]
|
|
SEED_REROLL_GROUPS = {
|
|
"none": (),
|
|
"category": ("category",),
|
|
"subcategory": ("subcategory",),
|
|
"content": ("content",),
|
|
"clothing": ("clothing",),
|
|
"person": ("person",),
|
|
"scene": ("scene",),
|
|
"pose": ("pose", "role"),
|
|
"role": ("role",),
|
|
"expression": ("expression",),
|
|
"composition": ("composition",),
|
|
"content_pose": ("content", "pose", "role"),
|
|
"content_clothing": ("content", "clothing"),
|
|
"clothing_pose": ("clothing", "pose", "role"),
|
|
"scene_pose": ("scene", "pose", "role"),
|
|
}
|
|
SEED_REROLL_AXIS_CHOICES = list(SEED_REROLL_GROUPS.keys())
|
|
|
|
|
|
def _normal_key(value: Any) -> str:
|
|
return str(value or "").strip().lower().replace("-", "_").replace(" ", "_")
|
|
|
|
|
|
def seed_mode_choices() -> list[str]:
|
|
return list(SEED_MODE_CHOICES)
|
|
|
|
|
|
def normalize_seed_mode(value: Any) -> str:
|
|
normalized = _normal_key(value)
|
|
aliases = {
|
|
"follow": "follow_main",
|
|
"followmain": "follow_main",
|
|
"follow_main_seed": "follow_main",
|
|
"main": "follow_main",
|
|
"main_seed": "follow_main",
|
|
}
|
|
normalized = aliases.get(normalized, normalized)
|
|
return normalized if normalized in SEED_MODE_CHOICES else "auto"
|
|
|
|
|
|
def seed_reroll_axis_choices() -> list[str]:
|
|
return list(SEED_REROLL_AXIS_CHOICES)
|
|
|
|
|
|
def normalize_reroll_axis(value: Any) -> str:
|
|
normalized = _normal_key(value)
|
|
aliases = {
|
|
"contentpose": "content_pose",
|
|
"contentclothing": "content_clothing",
|
|
"clothingpose": "clothing_pose",
|
|
"scenepose": "scene_pose",
|
|
}
|
|
normalized = aliases.get(normalized, normalized)
|
|
return normalized if normalized in SEED_REROLL_GROUPS else "none"
|
|
|
|
|
|
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,
|
|
clothing_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",
|
|
clothing_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 = normalize_seed_mode(mode)
|
|
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),
|
|
"clothing_seed": axis_seed(clothing_seed, clothing_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 = set(SEED_REROLL_GROUPS[normalize_reroll_axis(reroll_axis)])
|
|
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 configured_seed_from_axes(
|
|
seed_config: str | dict[str, Any] | None,
|
|
axes: Iterable[str],
|
|
*,
|
|
extra_keys: Iterable[str] = (),
|
|
) -> int | None:
|
|
try:
|
|
parsed = parse_seed_config(seed_config)
|
|
except ValueError:
|
|
return None
|
|
for axis in axes:
|
|
value = configured_axis_seed(parsed, axis)
|
|
if value is not None:
|
|
return value
|
|
for key in extra_keys:
|
|
value = parsed.get(str(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))
|
|
|
|
|
|
def axis_seed_trace(
|
|
seed_config: str | dict[str, Any] | None,
|
|
base_seed: int,
|
|
row_number: int,
|
|
axes: Iterable[str] = SEED_LOCK_AXES,
|
|
) -> dict[str, dict[str, int | str]]:
|
|
parsed = parse_seed_config(seed_config)
|
|
trace: dict[str, dict[str, int | str]] = {}
|
|
for axis in axes:
|
|
configured = configured_axis_seed(parsed, axis)
|
|
seed_value = int(configured) if configured is not None else int(base_seed)
|
|
source = "configured" if configured is not None else "main"
|
|
salt = SEED_AXIS_SALTS.get(axis, 0)
|
|
trace[axis] = {
|
|
"source": source,
|
|
"seed": seed_value,
|
|
"rng_seed": row_seed(seed_value, row_number, salt),
|
|
}
|
|
return trace
|