Upload BAN_Expression_To_Default.py

BAN_Expression_To_Default.py

@@ -0,0 +1,164 @@
+from __future__ import annotations
+
+import re
+
+
+DEFAULT_EXPRESSION_1 = "calm expression"
+DEFAULT_EXPRESSION_2 = "neutral expression"
+DEFAULT_EXPRESSION_3 = "slight smug smile"
+
+
+# Matches:
+# expression.1=...
+# expression.2: ...
+# expression_tag1=...
+# expression1=...
+_EXPR_RE = re.compile(
+    r"^expression(?:\.|_tag)?([1-5])\s*[:=]\s*(.*)$",
+    flags=re.IGNORECASE,
+)
+
+_SKIN_RE = re.compile(
+    r"^skin(?:\.|_tag)?([1-5])\s*[:=]",
+    flags=re.IGNORECASE,
+)
+
+_AESTHETIC_RE = re.compile(
+    r"^aesthetic(?:\.|_tag)?([1-5])\s*[:=]",
+    flags=re.IGNORECASE,
+)
+
+_EQUIP_RE = re.compile(
+    r"^(?:equip|equipment)\.",
+    flags=re.IGNORECASE,
+)
+
+_GPT_BAM_BLOCK_RE = re.compile(
+    r"GPT_BAM_START###(.*?)###GPT_BAM_END",
+    flags=re.IGNORECASE | re.DOTALL,
+)
+
+
+def _split_segments(payload: str) -> list[str]:
+    """Split a GPT_BAM payload into ### segments, removing empty/newline-only parts."""
+    payload = (payload or "").replace("\r", "\n")
+    return [seg.strip() for seg in payload.split("###") if seg.strip()]
+
+
+def _is_expression_segment(seg: str) -> bool:
+    return _EXPR_RE.match(seg) is not None
+
+
+def _find_insertion_index(segments: list[str]) -> int:
+    """
+    Where to insert expression.1/2/3 if no expression exists yet:
+
+    1) first existing expression position
+    2) after last skin tag
+    3) after last aesthetic tag
+    4) after last equip.* tag
+    5) otherwise at end
+    """
+    first_expr_idx = None
+    last_skin_idx = None
+    last_aesthetic_idx = None
+    last_equip_idx = None
+
+    for i, seg in enumerate(segments):
+        if _is_expression_segment(seg) and first_expr_idx is None:
+            first_expr_idx = i
+        if _SKIN_RE.match(seg):
+            last_skin_idx = i
+        if _AESTHETIC_RE.match(seg):
+            last_aesthetic_idx = i
+        if _EQUIP_RE.match(seg):
+            last_equip_idx = i
+
+    if first_expr_idx is not None:
+        return first_expr_idx
+    if last_skin_idx is not None:
+        return last_skin_idx + 1
+    if last_aesthetic_idx is not None:
+        return last_aesthetic_idx + 1
+    if last_equip_idx is not None:
+        return last_equip_idx + 1
+    return len(segments)
+
+
+def _rewrite_payload(payload: str) -> str:
+    segments = _split_segments(payload)
+    insertion_index = _find_insertion_index(segments)
+
+    new_expression_segments = [
+        f"expression.1={DEFAULT_EXPRESSION_1}",
+        f"expression.2={DEFAULT_EXPRESSION_2}",
+        f"expression.3={DEFAULT_EXPRESSION_3}",
+    ]
+
+    out_segments: list[str] = []
+    inserted = False
+
+    for i, seg in enumerate(segments):
+        if not inserted and i == insertion_index:
+            out_segments.extend(new_expression_segments)
+            inserted = True
+
+        # Remove ALL existing expression.1..5 / expression_tag1..5
+        if _is_expression_segment(seg):
+            continue
+
+        out_segments.append(seg)
+
+    if not inserted:
+        out_segments.extend(new_expression_segments)
+
+    return "###".join(out_segments)
+
+
+def _rewrite_gpt_bam_text(text: str) -> str:
+    """
+    If a GPT_BAM block is found, only rewrite that block and preserve any text outside it.
+    If no GPT_BAM block is found, treat the whole input as a raw payload and rewrite it.
+    """
+    text = text or ""
+    match = _GPT_BAM_BLOCK_RE.search(text)
+
+    if not match:
+        # Fallback: treat entire text as payload and wrap it back into GPT_BAM
+        rewritten_payload = _rewrite_payload(text)
+        return f"GPT_BAM_START###{rewritten_payload}###GPT_BAM_END"
+
+    original_payload = match.group(1)
+    rewritten_payload = _rewrite_payload(original_payload)
+    new_block = f"GPT_BAM_START###{rewritten_payload}###GPT_BAM_END"
+
+    return text[:match.start()] + new_block + text[match.end():]
+
+
+class BAM_expression_default:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "BAM-format_In": ("STRING", {"multiline": True, "default": ""}),
+            }
+        }
+
+    RETURN_TYPES = ("STRING",)
+    RETURN_NAMES = ("BAM-format_OUT",)
+    FUNCTION = "apply"
+    CATEGORY = "BAM"
+
+    def apply(self, **kwargs):
+        bam_in = kwargs.get("BAM-format_In", "")
+        bam_out = _rewrite_gpt_bam_text(bam_in)
+        return (bam_out,)
+
+
+NODE_CLASS_MAPPINGS = {
+    "BAM_expression_default": BAM_expression_default,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "BAM_expression_default": "BAM_expression_default",
+}