336 lines
12 KiB
Python
336 lines
12 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import random
|
|
|
|
try:
|
|
from .prompt_builder import (
|
|
build_cast_config_json,
|
|
build_category_config_json,
|
|
cast_preset_choices,
|
|
category_preset_choices,
|
|
subcategory_choices,
|
|
)
|
|
from .location_config import (
|
|
build_composition_pool_json,
|
|
build_location_pool_json,
|
|
build_thematic_location_json,
|
|
composition_pool_preset_choices,
|
|
location_pool_preset_choices,
|
|
location_theme_choices,
|
|
)
|
|
except ImportError: # Allows local smoke tests from the repository root.
|
|
from prompt_builder import (
|
|
build_cast_config_json,
|
|
build_category_config_json,
|
|
cast_preset_choices,
|
|
category_preset_choices,
|
|
subcategory_choices,
|
|
)
|
|
from location_config import (
|
|
build_composition_pool_json,
|
|
build_location_pool_json,
|
|
build_thematic_location_json,
|
|
composition_pool_preset_choices,
|
|
location_pool_preset_choices,
|
|
location_theme_choices,
|
|
)
|
|
|
|
|
|
SXCP_CATEGORY_CONFIG = "SXCP_CATEGORY_CONFIG"
|
|
SXCP_LOCATION_CONFIG = "SXCP_LOCATION_CONFIG"
|
|
SXCP_COMPOSITION_CONFIG = "SXCP_COMPOSITION_CONFIG"
|
|
SXCP_CAST_CONFIG = "SXCP_CAST_CONFIG"
|
|
SXCP_SEED_CONFIG = "SXCP_SEED_CONFIG"
|
|
|
|
|
|
class SxCPCategoryPreset:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"preset": (category_preset_choices(), {"default": "auto_weighted"}),
|
|
"subcategory": (subcategory_choices(), {"default": "random"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = (SXCP_CATEGORY_CONFIG, "STRING", "STRING")
|
|
RETURN_NAMES = ("category_config", "category", "subcategory")
|
|
FUNCTION = "build"
|
|
CATEGORY = "prompt_builder"
|
|
|
|
def build(self, preset, subcategory):
|
|
config = build_category_config_json(preset=preset, subcategory=subcategory)
|
|
parsed = json.loads(config)
|
|
return config, parsed["category"], parsed["subcategory"]
|
|
|
|
|
|
class SxCPLocationPool:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"enabled": ("BOOLEAN", {"default": True}),
|
|
"combine_mode": (["replace", "add"], {"default": "replace"}),
|
|
"preset": (location_pool_preset_choices(), {"default": "custom_only"}),
|
|
"custom_locations": ("STRING", {"default": "", "multiline": True}),
|
|
},
|
|
"optional": {
|
|
"location_config": (SXCP_LOCATION_CONFIG,),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = (SXCP_LOCATION_CONFIG, "STRING")
|
|
RETURN_NAMES = ("location_config", "summary")
|
|
FUNCTION = "build"
|
|
CATEGORY = "prompt_builder"
|
|
|
|
def build(self, enabled, combine_mode, preset, custom_locations, location_config=""):
|
|
config = build_location_pool_json(
|
|
enabled=enabled,
|
|
combine_mode=combine_mode,
|
|
preset=preset,
|
|
custom_locations=custom_locations or "",
|
|
location_config=location_config or "",
|
|
)
|
|
parsed = json.loads(config)
|
|
return config, parsed.get("summary", "")
|
|
|
|
|
|
class SxCPCompositionPool:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"enabled": ("BOOLEAN", {"default": True}),
|
|
"combine_mode": (["replace", "add"], {"default": "replace"}),
|
|
"preset": (composition_pool_preset_choices(), {"default": "no_outfit_check"}),
|
|
"custom_compositions": ("STRING", {"default": "", "multiline": True}),
|
|
},
|
|
"optional": {
|
|
"composition_config": (SXCP_COMPOSITION_CONFIG,),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = (SXCP_COMPOSITION_CONFIG, "STRING")
|
|
RETURN_NAMES = ("composition_config", "summary")
|
|
FUNCTION = "build"
|
|
CATEGORY = "prompt_builder"
|
|
|
|
def build(self, enabled, combine_mode, preset, custom_compositions, composition_config=""):
|
|
config = build_composition_pool_json(
|
|
enabled=enabled,
|
|
combine_mode=combine_mode,
|
|
preset=preset,
|
|
custom_compositions=custom_compositions or "",
|
|
composition_config=composition_config or "",
|
|
)
|
|
parsed = json.loads(config)
|
|
return config, parsed.get("summary", "")
|
|
|
|
|
|
class SxCPLocationTheme:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"enabled": ("BOOLEAN", {"default": True}),
|
|
"combine_mode": (["replace", "add"], {"default": "replace"}),
|
|
"theme": (location_theme_choices(), {"default": "semi_public_affair"}),
|
|
"custom_locations": ("STRING", {"default": "", "multiline": True}),
|
|
"custom_compositions": ("STRING", {"default": "", "multiline": True}),
|
|
},
|
|
"optional": {
|
|
"location_config": (SXCP_LOCATION_CONFIG,),
|
|
"composition_config": (SXCP_COMPOSITION_CONFIG,),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = (SXCP_LOCATION_CONFIG, SXCP_COMPOSITION_CONFIG, "STRING")
|
|
RETURN_NAMES = ("location_config", "composition_config", "summary")
|
|
FUNCTION = "build"
|
|
CATEGORY = "prompt_builder"
|
|
|
|
def build(
|
|
self,
|
|
enabled,
|
|
combine_mode,
|
|
theme,
|
|
custom_locations,
|
|
custom_compositions,
|
|
location_config="",
|
|
composition_config="",
|
|
):
|
|
return build_thematic_location_json(
|
|
enabled=enabled,
|
|
combine_mode=combine_mode,
|
|
theme=theme,
|
|
custom_locations=custom_locations or "",
|
|
custom_compositions=custom_compositions or "",
|
|
location_config=location_config or "",
|
|
composition_config=composition_config or "",
|
|
)
|
|
|
|
|
|
class SxCPCastControl:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"cast_mode": (cast_preset_choices(), {"default": "mixed_couple"}),
|
|
"women_count": ("INT", {"default": 1, "min": 0, "max": 12, "step": 1}),
|
|
"men_count": ("INT", {"default": 1, "min": 0, "max": 12, "step": 1}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = (SXCP_CAST_CONFIG, "INT", "INT", "STRING")
|
|
RETURN_NAMES = ("cast_config", "women_count", "men_count", "cast_summary")
|
|
FUNCTION = "build"
|
|
CATEGORY = "prompt_builder"
|
|
|
|
def build(self, cast_mode, women_count, men_count):
|
|
config = build_cast_config_json(cast_mode=cast_mode, women_count=women_count, men_count=men_count)
|
|
parsed = json.loads(config)
|
|
summary = f"{parsed['women_count']} women, {parsed['men_count']} men"
|
|
return config, parsed["women_count"], parsed["men_count"], summary
|
|
|
|
|
|
class SxCPCastBias:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"seed": ("INT", {"default": -1, "min": -1, "max": 0xFFFFFFFF, "step": 1}),
|
|
"row_number": ("INT", {"default": 1, "min": 1, "max": 1000000, "step": 1}),
|
|
"women_weights": ("STRING", {"default": "0.60,0.25,0.10,0.05"}),
|
|
"women_start_count": ("INT", {"default": 1, "min": 0, "max": 12, "step": 1}),
|
|
"men_weights": ("STRING", {"default": "0.45,0.40,0.10,0.05"}),
|
|
"men_start_count": ("INT", {"default": 0, "min": 0, "max": 12, "step": 1}),
|
|
"empty_behavior": (["force_one_woman", "force_one_man", "allow_empty"], {"default": "force_one_woman"}),
|
|
},
|
|
"optional": {
|
|
"seed_config": (SXCP_SEED_CONFIG,),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = (SXCP_CAST_CONFIG, "INT", "INT", "STRING")
|
|
RETURN_NAMES = ("cast_config", "women_count", "men_count", "cast_summary")
|
|
FUNCTION = "build"
|
|
CATEGORY = "prompt_builder"
|
|
|
|
@staticmethod
|
|
def _configured_cast_seed(seed_config):
|
|
if not seed_config:
|
|
return None
|
|
if isinstance(seed_config, dict):
|
|
raw = seed_config
|
|
else:
|
|
try:
|
|
raw = json.loads(str(seed_config))
|
|
except (TypeError, ValueError, json.JSONDecodeError):
|
|
return None
|
|
if not isinstance(raw, dict):
|
|
return None
|
|
for key in ("category_seed", "content_seed", "role_seed", "seed", "global_seed"):
|
|
try:
|
|
value = int(raw.get(key))
|
|
except (TypeError, ValueError):
|
|
continue
|
|
if value >= 0:
|
|
return value
|
|
return None
|
|
|
|
@staticmethod
|
|
def _weight_pairs(weights_text, start_count):
|
|
pairs = []
|
|
start = max(0, min(12, int(start_count)))
|
|
parts = str(weights_text or "").replace("\n", ",").split(",")
|
|
for offset, raw in enumerate(parts):
|
|
count = start + offset
|
|
if count > 12:
|
|
break
|
|
try:
|
|
weight = float(raw.strip())
|
|
except (TypeError, ValueError):
|
|
continue
|
|
if weight > 0:
|
|
pairs.append((count, weight))
|
|
return pairs or [(start, 1.0)]
|
|
|
|
@staticmethod
|
|
def _weighted_count(rng, pairs):
|
|
total = sum(weight for _count, weight in pairs)
|
|
point = rng.random() * total
|
|
upto = 0.0
|
|
for count, weight in pairs:
|
|
upto += weight
|
|
if point <= upto:
|
|
return int(count)
|
|
return int(pairs[-1][0])
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, *args, **kwargs):
|
|
seed_value = kwargs.get("seed")
|
|
if seed_value is None and args:
|
|
seed_value = args[0]
|
|
seed_config = kwargs.get("seed_config", "")
|
|
if not seed_config and len(args) > 7:
|
|
seed_config = args[7]
|
|
try:
|
|
seed = int(seed_value)
|
|
except (TypeError, ValueError):
|
|
seed = -1
|
|
if seed < 0 and cls._configured_cast_seed(seed_config) is None:
|
|
return random.random()
|
|
return tuple(args), tuple(sorted(kwargs.items()))
|
|
|
|
def build(
|
|
self,
|
|
seed,
|
|
row_number,
|
|
women_weights,
|
|
women_start_count,
|
|
men_weights,
|
|
men_start_count,
|
|
empty_behavior,
|
|
seed_config="",
|
|
):
|
|
configured_seed = self._configured_cast_seed(seed_config)
|
|
if configured_seed is None and int(seed) < 0:
|
|
rng = random.Random(random.getrandbits(64))
|
|
else:
|
|
cast_seed = configured_seed if configured_seed is not None else int(seed)
|
|
rng = random.Random(f"sxcp_cast_bias:{cast_seed}:{int(row_number)}")
|
|
women_pairs = self._weight_pairs(women_weights, women_start_count)
|
|
men_pairs = self._weight_pairs(men_weights, men_start_count)
|
|
women_count = self._weighted_count(rng, women_pairs)
|
|
men_count = self._weighted_count(rng, men_pairs)
|
|
if women_count + men_count == 0:
|
|
if empty_behavior == "force_one_man":
|
|
men_count = 1
|
|
elif empty_behavior != "allow_empty":
|
|
women_count = 1
|
|
config = build_cast_config_json(cast_mode="custom_counts", women_count=women_count, men_count=men_count)
|
|
parsed = json.loads(config)
|
|
summary = f"weighted cast: {parsed['women_count']} women, {parsed['men_count']} men"
|
|
return config, parsed["women_count"], parsed["men_count"], summary
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"SxCPCategoryPreset": SxCPCategoryPreset,
|
|
"SxCPLocationPool": SxCPLocationPool,
|
|
"SxCPCompositionPool": SxCPCompositionPool,
|
|
"SxCPLocationTheme": SxCPLocationTheme,
|
|
"SxCPCastControl": SxCPCastControl,
|
|
"SxCPCastBias": SxCPCastBias,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"SxCPCategoryPreset": "SxCP Category Preset",
|
|
"SxCPLocationPool": "SxCP Location Pool",
|
|
"SxCPCompositionPool": "SxCP Composition Pool",
|
|
"SxCPLocationTheme": "SxCP Location Theme",
|
|
"SxCPCastControl": "SxCP Cast Control",
|
|
"SxCPCastBias": "SxCP Cast Bias",
|
|
}
|