527 lines
20 KiB
Python
527 lines
20 KiB
Python
# ComfyUI-RPG-Characters - RPG Character Selector with DNA Locker
|
|
# Author: Lord Lethris
|
|
|
|
import hashlib
|
|
import json
|
|
import re
|
|
import secrets
|
|
import random
|
|
|
|
from .rpg_character_data.rpg_race_data import RACE_DATA
|
|
from .rpg_character_data.rpg_ethnicity_data import ETHNICITY_DATA
|
|
from .rpg_character_data.rpg_gender_data import GENDER_DATA
|
|
from .rpg_character_data.rpg_age_data import AGE_DATA
|
|
from .rpg_character_data.rpg_class_data import CLASS_DATA
|
|
from .rpg_character_data.rpg_hair_style_data import HAIR_STYLE_DATA
|
|
from .rpg_character_data.rpg_hair_colour_data import HAIR_COLOUR_DATA
|
|
from .rpg_character_data.rpg_beard_style_data import BEARD_STYLE_DATA
|
|
from .rpg_character_data.rpg_beard_colour_data import BEARD_COLOUR_DATA
|
|
from .rpg_character_data.rpg_clothes_style_data import CLOTHES_STYLE_DATA
|
|
from .rpg_character_data.rpg_augment_data import AUGMENT_DATA
|
|
from .rpg_character_data.rpg_emotion_data import EMOTION_DATA
|
|
from .rpg_character_data.rpg_scene_data import SCENE_DATA
|
|
|
|
|
|
VARIANT_PATTERN = re.compile(r"\{([^{}]+)\}")
|
|
|
|
|
|
def _race_overrides_ethnicity_visuals(race_entry):
|
|
"""Return visual trait categories explicitly controlled by the race."""
|
|
prompt = str(race_entry.get("prompt", "")).lower()
|
|
overrides = set()
|
|
if "skin tone" in prompt or "skin tones" in prompt or "skin color" in prompt or "skin colour" in prompt:
|
|
overrides.add("skin")
|
|
if re.search(r"\beyes?\b", prompt):
|
|
overrides.add("eyes")
|
|
return overrides
|
|
|
|
|
|
def _remove_ethnicity_visual_conflicts(text, overrides):
|
|
"""Keep ethnicity identity/facial traits, but yield explicit race traits.
|
|
|
|
Ethnicity data is deliberately lightweight prose rather than a structured
|
|
schema, so conflict removal is clause-based. The explicit ethnicity label
|
|
is always retained even when its first clause also contains skin/eye data.
|
|
"""
|
|
if not text or not overrides:
|
|
return text
|
|
|
|
clauses = [clause.strip() for clause in str(text).split(",")]
|
|
kept = []
|
|
skip_continuation = False
|
|
for index, clause in enumerate(clauses):
|
|
lower = clause.lower()
|
|
has_skin = bool(re.search(
|
|
r"\bskin(?:\s+tone|\s+tones|\s+color|\s+colour)?\b", lower
|
|
))
|
|
has_eyes = bool(re.search(
|
|
r"\beye(?:s|\s+(?:color|colour))?\b", lower
|
|
))
|
|
|
|
if skip_continuation:
|
|
# Eye colours often continue across comma-separated fragments:
|
|
# "... eye colors, typically brown, hazel, or green". Suppress
|
|
# those continuation fragments while preserving later substantive
|
|
# facial descriptors if a data entry happens to contain them.
|
|
substantive = r"\b(face|features?|jaw|nose|cheek|hair|build|stature|complexion|skin)\b"
|
|
if not re.search(substantive, lower):
|
|
continue
|
|
skip_continuation = False
|
|
|
|
if ("skin" in overrides and has_skin) or ("eyes" in overrides and has_eyes):
|
|
if index == 0:
|
|
identity = re.match(r"(\(Ethnicity[^)]*\))", clause, flags=re.IGNORECASE)
|
|
if identity:
|
|
kept.append(identity.group(1))
|
|
if "eyes" in overrides and has_eyes:
|
|
skip_continuation = True
|
|
continue
|
|
kept.append(clause)
|
|
|
|
result = ", ".join(c for c in kept if c)
|
|
result = re.sub(r"(^|,\s+)(?:and|with)\s+", r"\1", result, flags=re.IGNORECASE)
|
|
result = re.sub(r"\s{2,}", " ", result).strip(" ,")
|
|
return result
|
|
|
|
|
|
def _stable_seed(master_seed, locus_id, reroll_count=0):
|
|
"""Create a stable per-locus seed, including an optional reroll counter."""
|
|
payload = f"RPGCharacterSelector|{master_seed}|{locus_id}|reroll:{int(reroll_count)}".encode("utf-8")
|
|
return int.from_bytes(hashlib.sha256(payload).digest()[:8], "big")
|
|
|
|
|
|
def _parse_state(value):
|
|
if not value:
|
|
return {}
|
|
if isinstance(value, dict):
|
|
return value
|
|
try:
|
|
parsed = json.loads(value)
|
|
return parsed if isinstance(parsed, dict) else {}
|
|
except (TypeError, ValueError, json.JSONDecodeError):
|
|
return {}
|
|
|
|
|
|
def _normalise_weights(weights, count):
|
|
values = []
|
|
for index in range(count):
|
|
try:
|
|
value = float(weights.get(str(index), weights.get(index, 0)))
|
|
except (TypeError, ValueError):
|
|
value = 0.0
|
|
values.append(max(0.0, value))
|
|
|
|
# V2.4 DNA is intentionally limited to the nearest three variants.
|
|
# This also sanitises older/stale workflow state that may contain more.
|
|
if count > 3:
|
|
keep = sorted(range(count), key=lambda i: values[i], reverse=True)[:3]
|
|
keep = set(keep)
|
|
values = [value if i in keep else 0.0 for i, value in enumerate(values)]
|
|
|
|
total = sum(values)
|
|
if total <= 0:
|
|
return None
|
|
return [value / total for value in values]
|
|
|
|
|
|
def _weighted_variant(options, weights):
|
|
"""Render an explicitly sculpted DNA locus.
|
|
|
|
The output deliberately uses the user's proposed '&' syntax.
|
|
The surrounding prompt parentheses (when present) are preserved, so:
|
|
{emerald|gold} Skin tone:1.3
|
|
can become:
|
|
((emerald:0.7) & (gold:0.3) Skin tone:1.3)
|
|
"""
|
|
normalised = _normalise_weights(weights, len(options))
|
|
if not normalised:
|
|
return None
|
|
|
|
active = [
|
|
(option, weight)
|
|
for option, weight in zip(options, normalised)
|
|
if weight > 0.0005
|
|
]
|
|
|
|
if not active:
|
|
return None
|
|
|
|
if len(active) == 1:
|
|
return active[0][0]
|
|
|
|
parts = [f"({option}:{weight:.3f}".rstrip("0").rstrip(".") + ")"
|
|
for option, weight in active]
|
|
return " & ".join(parts)
|
|
|
|
|
|
def resolve_prompt_variants_with_trace(text, master_seed, locus_prefix,
|
|
locks=None, weights=None, rerolls=None):
|
|
"""Resolve every {A|B|C} group and return a DNA trace for the UI."""
|
|
locks = locks or {}
|
|
weights = weights or {}
|
|
rerolls = rerolls or {}
|
|
selections = []
|
|
dna = []
|
|
group_index = 0
|
|
|
|
def replacer(match):
|
|
nonlocal group_index
|
|
raw = match.group(1)
|
|
options = [item.strip() for item in raw.split("|")]
|
|
locus_id = f"{locus_prefix}:{group_index}"
|
|
group_index += 1
|
|
|
|
if len(options) <= 1:
|
|
# Existing behaviour for literal braces is preserved.
|
|
choice = options[0] if options else ""
|
|
selections.append(choice)
|
|
return choice
|
|
|
|
# V2.4 no longer exposes lock state in the UI. Keep the parameter for
|
|
# backwards compatibility with older workflows, but never let stale
|
|
# lock entries prevent the new Re-roll/Sculpt workflow from working.
|
|
locked_value = None
|
|
custom_weights = weights.get(locus_id)
|
|
reroll_state = rerolls.get(locus_id, 0)
|
|
avoid_value = None
|
|
preview_value = None
|
|
if isinstance(reroll_state, dict):
|
|
try:
|
|
reroll_count = max(0, int(reroll_state.get("count", 0)))
|
|
except (TypeError, ValueError):
|
|
reroll_count = 0
|
|
avoid_value = reroll_state.get("avoid")
|
|
preview_value = reroll_state.get("preview")
|
|
else:
|
|
try:
|
|
reroll_count = max(0, int(reroll_state or 0))
|
|
except (TypeError, ValueError):
|
|
reroll_count = 0
|
|
|
|
mode = "random"
|
|
if locked_value in options:
|
|
choice = locked_value
|
|
mode = "locked"
|
|
rendered = choice
|
|
weight_map = None
|
|
elif isinstance(custom_weights, dict):
|
|
rendered = _weighted_variant(options, custom_weights)
|
|
if rendered is not None:
|
|
choice = " / ".join(
|
|
f"{option} {weight:.1%}"
|
|
for option, weight in zip(
|
|
options,
|
|
_normalise_weights(custom_weights, len(options))
|
|
)
|
|
if float(custom_weights.get(str(options.index(option)), 0) or 0) > 0
|
|
)
|
|
mode = "sculpted"
|
|
weight_map = _normalise_weights(custom_weights, len(options))
|
|
else:
|
|
rendered = None
|
|
weight_map = None
|
|
else:
|
|
# The Locker can choose a reroll immediately in the UI. Persist
|
|
# that preview in the reroll state so the next ComfyUI execution
|
|
# uses exactly what the user saw, rather than changing it again
|
|
# when Run is pressed. Older states without a preview continue
|
|
# to use the deterministic seed+reroll-counter path.
|
|
if preview_value in options:
|
|
choice = preview_value
|
|
else:
|
|
rng = random.Random(_stable_seed(master_seed, locus_id, reroll_count))
|
|
if avoid_value in options and len(options) > 1:
|
|
# A deliberate locus reroll should actually move to another
|
|
# variant when alternatives exist, rather than occasionally
|
|
# landing on the same value by chance.
|
|
alternatives = [option for option in options if option != avoid_value]
|
|
choice = rng.choice(alternatives)
|
|
else:
|
|
choice = rng.choice(options)
|
|
rendered = choice
|
|
weight_map = None
|
|
|
|
if mode != "sculpted":
|
|
weight_map = [1.0 if option == choice else 0.0 for option in options]
|
|
|
|
dna.append({
|
|
"id": locus_id,
|
|
"label": f"Variant {group_index}",
|
|
"options": options,
|
|
"selected": choice,
|
|
"weights": weight_map,
|
|
"mode": mode,
|
|
})
|
|
selections.append(choice)
|
|
return rendered
|
|
|
|
resolved = VARIANT_PATTERN.sub(replacer, text)
|
|
return resolved, selections, dna
|
|
|
|
|
|
class RPGCharacterSelector:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"race": (list(RACE_DATA.keys()),),
|
|
"ethnicity": (list(ETHNICITY_DATA.keys()),),
|
|
"gender": (list(GENDER_DATA.keys()),),
|
|
"age": (list(AGE_DATA.keys()),),
|
|
"character_class": (list(CLASS_DATA.keys()),),
|
|
"hair_style": (list(HAIR_STYLE_DATA.keys()),),
|
|
"hair_colour": (list(HAIR_COLOUR_DATA.keys()),),
|
|
"beard_style": (list(BEARD_STYLE_DATA.keys()),),
|
|
"beard_colour": (list(BEARD_COLOUR_DATA.keys()),),
|
|
"clothes_style": (list(CLOTHES_STYLE_DATA.keys()),),
|
|
"Augmentations": (list(AUGMENT_DATA.keys()),),
|
|
"emotion": (list(EMOTION_DATA.keys()),),
|
|
"scene": (list(SCENE_DATA.keys()),),
|
|
"dna_seed": ("INT", {
|
|
"default": -1,
|
|
"min": -1,
|
|
"max": 4294967295,
|
|
"step": 1,
|
|
}),
|
|
# Hidden execution revision. DNA Locker UI edits must invalidate
|
|
# ComfyUI's execution cache without changing the master DNA seed.
|
|
"dna_revision": ("INT", {
|
|
"default": 0,
|
|
"min": 0,
|
|
"max": 2147483647,
|
|
"step": 1,
|
|
}),
|
|
},
|
|
"optional": {
|
|
"dna_locks": ("STRING", {"default": "{}", "multiline": False}),
|
|
"dna_weights": ("STRING", {"default": "{}", "multiline": False}),
|
|
"dna_rerolls": ("STRING", {"default": "{}", "multiline": False}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "STRING", "STRING", "STRING")
|
|
RETURN_NAMES = (
|
|
"positive_prompt",
|
|
"negative_prompt",
|
|
"Ollama_Generate_V2_Textbox_2",
|
|
"selection_summary",
|
|
)
|
|
FUNCTION = "generate_prompt"
|
|
CATEGORY = "RPG"
|
|
OUTPUT_NODE = True
|
|
|
|
def generate_prompt(
|
|
self,
|
|
race,
|
|
ethnicity,
|
|
gender,
|
|
age,
|
|
character_class,
|
|
hair_style,
|
|
hair_colour,
|
|
beard_style,
|
|
beard_colour,
|
|
clothes_style,
|
|
Augmentations,
|
|
emotion,
|
|
scene,
|
|
dna_seed=-1,
|
|
dna_revision=0,
|
|
dna_locks="{}",
|
|
dna_weights="{}",
|
|
dna_rerolls="{}",
|
|
):
|
|
if dna_seed is None or int(dna_seed) < 0:
|
|
dna_seed = secrets.randbits(32)
|
|
else:
|
|
dna_seed = int(dna_seed)
|
|
|
|
locks = _parse_state(dna_locks)
|
|
weights = _parse_state(dna_weights)
|
|
rerolls = _parse_state(dna_rerolls)
|
|
|
|
# Some ComfyUI frontend builds can treat auxiliary hidden widgets
|
|
# differently when building the execution payload. The DNA Locker
|
|
# therefore mirrors reroll/unlock state inside dna_locks as a fallback.
|
|
embedded_rerolls = locks.pop("__dna_rerolls", None)
|
|
if not rerolls and isinstance(embedded_rerolls, dict):
|
|
rerolls = embedded_rerolls
|
|
|
|
race_entry = RACE_DATA[race]
|
|
ethnicity_entry = ETHNICITY_DATA[ethnicity]
|
|
ethnicity_entry_for_prompt = dict(ethnicity_entry)
|
|
race_overrides = _race_overrides_ethnicity_visuals(race_entry)
|
|
if race_overrides and ethnicity_entry_for_prompt.get("prompt"):
|
|
ethnicity_entry_for_prompt["prompt"] = _remove_ethnicity_visual_conflicts(
|
|
ethnicity_entry_for_prompt["prompt"], race_overrides
|
|
)
|
|
|
|
selected_data = [
|
|
("Race", race, race_entry),
|
|
("Ethnicity", ethnicity, ethnicity_entry_for_prompt),
|
|
("Gender", gender, GENDER_DATA[gender]),
|
|
("Age", age, AGE_DATA[age]),
|
|
("Class", character_class, CLASS_DATA[character_class]),
|
|
("Hair Style", hair_style, HAIR_STYLE_DATA[hair_style]),
|
|
("Hair Colour", hair_colour, HAIR_COLOUR_DATA[hair_colour]),
|
|
("Beard Style", beard_style, BEARD_STYLE_DATA[beard_style]),
|
|
("Beard Colour", beard_colour, BEARD_COLOUR_DATA[beard_colour]),
|
|
("Clothes Style", clothes_style, CLOTHES_STYLE_DATA[clothes_style]),
|
|
("Augmentations", Augmentations, AUGMENT_DATA[Augmentations]),
|
|
("Emotion", emotion, EMOTION_DATA[emotion]),
|
|
("Scene", scene, SCENE_DATA[scene]),
|
|
]
|
|
|
|
resolved_prompts = []
|
|
resolved_negatives = []
|
|
all_selections = []
|
|
dna_groups = []
|
|
|
|
for field_name, selected_name, entry in selected_data:
|
|
if entry["prompt"]:
|
|
resolved, selections, dna = resolve_prompt_variants_with_trace(
|
|
entry["prompt"],
|
|
dna_seed,
|
|
f"{field_name}:positive",
|
|
locks,
|
|
weights,
|
|
rerolls,
|
|
)
|
|
resolved_prompts.append(resolved)
|
|
all_selections.extend(selections)
|
|
for group in dna:
|
|
group["field"] = field_name
|
|
group["source"] = "positive"
|
|
group["label"] = self._infer_variant_label(
|
|
entry["prompt"], group["options"], group["label"]
|
|
)
|
|
dna_groups.extend(dna)
|
|
|
|
if entry["negative_prompt"]:
|
|
resolved, selections, dna = resolve_prompt_variants_with_trace(
|
|
entry["negative_prompt"],
|
|
dna_seed,
|
|
f"{field_name}:negative",
|
|
locks,
|
|
weights,
|
|
rerolls,
|
|
)
|
|
resolved_negatives.append(resolved)
|
|
all_selections.extend(selections)
|
|
for group in dna:
|
|
group["field"] = field_name
|
|
group["source"] = "negative"
|
|
group["label"] = self._infer_variant_label(
|
|
entry["negative_prompt"], group["options"], group["label"]
|
|
)
|
|
dna_groups.extend(dna)
|
|
|
|
combined_prompt = "Create an extreme close-up portrait of " + ", ".join(resolved_prompts)
|
|
combined_negative_prompt = ", ".join(resolved_negatives)
|
|
|
|
if scene.lower().startswith("plain") or "chroma key" in scene.lower() or "solid" in scene.lower():
|
|
scene_description = (
|
|
f"Scene: Solid {scene.replace('Plain ', '').replace('Chroma key ', '')} "
|
|
"Background Only. No objects or distractions."
|
|
)
|
|
else:
|
|
scene_description = f"Scene: {scene}"
|
|
|
|
selection_summary = "\n".join([
|
|
f"Race: {race}",
|
|
f"Ethnicity: {ethnicity}",
|
|
f"Gender: {gender}",
|
|
f"Age: {age}",
|
|
f"Class: {character_class}",
|
|
f"Hair Style: {hair_style}",
|
|
f"Hair Colour: {hair_colour}",
|
|
f"Beard Style: {beard_style}",
|
|
f"Beard Colour: {beard_colour}",
|
|
f"Clothes Style: {clothes_style}",
|
|
f"Augmentations: {Augmentations}",
|
|
f"Emotion: {emotion}",
|
|
scene_description,
|
|
f"DNA Seed: {dna_seed}",
|
|
])
|
|
|
|
if race_overrides:
|
|
selection_summary += (
|
|
"\nRace Priority: " + ", ".join(sorted(race_overrides)) +
|
|
" override(s) applied to conflicting ethnicity traits."
|
|
)
|
|
|
|
ollama_description = (
|
|
"Create an extreme close-up portrait of " +
|
|
", ".join(resolved_prompts) +
|
|
", Include Art Style."
|
|
)
|
|
|
|
if all_selections:
|
|
selection_summary += "\nRandomized Choices Used: " + ", ".join(all_selections)
|
|
|
|
# Human-readable DNA record. This is also useful when the node is
|
|
# connected to logging/debugging tools.
|
|
if dna_groups:
|
|
selection_summary += "\nDNA Loci:"
|
|
for group in dna_groups:
|
|
if group["mode"] == "sculpted" and group["weights"]:
|
|
parts = [
|
|
f"{option}={weight:.3f}".rstrip("0").rstrip(".")
|
|
for option, weight in zip(group["options"], group["weights"])
|
|
if weight > 0.0005
|
|
]
|
|
resolved = " / ".join(parts)
|
|
else:
|
|
resolved = group["selected"]
|
|
selection_summary += (
|
|
f"\n- {group['field']} [{group['source']}] "
|
|
f"{group['label']}: {resolved} ({group['mode']})"
|
|
)
|
|
|
|
# ComfyUI flattens UI values by iterating each UI field. Therefore
|
|
# custom structured UI payloads must themselves be list-valued, or a
|
|
# dict here would arrive in JavaScript as ["seed", "groups"].
|
|
# Keep the DNA record as a single list item.
|
|
ui_payload = {
|
|
"dna": [{
|
|
"seed": dna_seed,
|
|
"groups": dna_groups,
|
|
}]
|
|
}
|
|
|
|
return {
|
|
"ui": ui_payload,
|
|
"result": (
|
|
combined_prompt,
|
|
combined_negative_prompt,
|
|
ollama_description,
|
|
selection_summary,
|
|
),
|
|
}
|
|
|
|
@staticmethod
|
|
def _infer_variant_label(source_text, options, fallback):
|
|
"""Best-effort label for common '... Skin tone' / '... eyes' groups."""
|
|
marker = "}"
|
|
matches = list(VARIANT_PATTERN.finditer(source_text))
|
|
for match in matches:
|
|
if [item.strip() for item in match.group(1).split("|")] == options:
|
|
suffix = source_text[match.end():]
|
|
suffix = suffix.lstrip(" ,)")
|
|
label_match = re.match(r"([A-Za-z][A-Za-z _-]{1,40})", suffix)
|
|
if label_match:
|
|
label = label_match.group(1).strip()
|
|
if label.lower().startswith(("skin tone", "eyes", "eye colour")):
|
|
return label
|
|
break
|
|
return fallback
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"RPGCharacterSelector": RPGCharacterSelector
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"RPGCharacterSelector": "RPG Character Selector"
|
|
}
|