Extract route config nodes
This commit is contained in:
@@ -0,0 +1,331 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import random
|
||||
|
||||
try:
|
||||
from .prompt_builder import (
|
||||
build_cast_config_json,
|
||||
build_category_config_json,
|
||||
build_composition_pool_json,
|
||||
build_location_pool_json,
|
||||
build_thematic_location_json,
|
||||
cast_preset_choices,
|
||||
category_preset_choices,
|
||||
composition_pool_preset_choices,
|
||||
location_pool_preset_choices,
|
||||
location_theme_choices,
|
||||
subcategory_choices,
|
||||
)
|
||||
except ImportError: # Allows local smoke tests from the repository root.
|
||||
from prompt_builder import (
|
||||
build_cast_config_json,
|
||||
build_category_config_json,
|
||||
build_composition_pool_json,
|
||||
build_location_pool_json,
|
||||
build_thematic_location_json,
|
||||
cast_preset_choices,
|
||||
category_preset_choices,
|
||||
composition_pool_preset_choices,
|
||||
location_pool_preset_choices,
|
||||
location_theme_choices,
|
||||
subcategory_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",
|
||||
}
|
||||
Reference in New Issue
Block a user