113 lines
3.9 KiB
Python
113 lines
3.9 KiB
Python
# pyright: reportArgumentType=false
|
|
import os
|
|
import json
|
|
from server import PromptServer
|
|
from aiohttp import web
|
|
|
|
# --- Constants ---
|
|
NODE_DIR = os.path.dirname(os.path.abspath(__file__))
|
|
LABELS_FILE = os.path.join(NODE_DIR, "switch_any_labels.json")
|
|
MAX_INPUTS = 10 # 允许的标签最大绝对数量
|
|
|
|
# --- Helper Functions for File I/O ---
|
|
|
|
def read_labels_from_file():
|
|
"""从JSON文件中读取标签列表。"""
|
|
if not os.path.exists(LABELS_FILE):
|
|
return []
|
|
try:
|
|
with open(LABELS_FILE, "r", encoding="utf-8") as f:
|
|
labels = json.load(f)
|
|
return labels if isinstance(labels, list) else []
|
|
except (IOError, json.JSONDecodeError):
|
|
return []
|
|
|
|
def write_labels_to_file(labels):
|
|
"""将给定的标签列表保存到JSON文件中。"""
|
|
try:
|
|
# 过滤空字符串和空白字符串，保持原始顺序和重复项
|
|
filtered_labels = [label for label in labels if label and str(label).strip()]
|
|
existing_labels = read_labels_from_file()
|
|
if filtered_labels != existing_labels:
|
|
with open(LABELS_FILE, "w", encoding="utf-8") as f:
|
|
json.dump(filtered_labels, f, indent=2)
|
|
except IOError as e:
|
|
print(f"[Switch Any] ERROR: Could not write labels to file: {e}")
|
|
|
|
# --- Custom Node Definition ---
|
|
|
|
class SwitchAny:
|
|
"""
|
|
一个用于定义和选择标签以进行路由的节点。
|
|
输入字段的数量会随着您添加更多标签而增加。
|
|
使用 'Save & Refresh UI' 按钮来保存更改并更新界面。
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
saved_labels = read_labels_from_file()
|
|
dropdown_labels = saved_labels if saved_labels else [""]
|
|
|
|
inputs = {
|
|
"required": {
|
|
"select_label": (dropdown_labels, ),
|
|
},
|
|
"optional": {}
|
|
}
|
|
|
|
# 总是显示MAX_INPUTS个输入框，确保工作流迁移时状态完整
|
|
for i in range(1, MAX_INPUTS + 1):
|
|
default_value = saved_labels[i-1] if i <= len(saved_labels) else ""
|
|
inputs["optional"][f"label_{i}"] = ("STRING", {
|
|
"default": default_value,
|
|
"multiline": False,
|
|
"placeholder": f"Label for index {i}"
|
|
})
|
|
|
|
return inputs
|
|
|
|
RETURN_TYPES = ("INT", "STRING")
|
|
RETURN_NAMES = ("index", "label")
|
|
FUNCTION = "execute"
|
|
CATEGORY = "ComfyUI_Switch"
|
|
|
|
def execute(self, select_label, **kwargs):
|
|
# 此方法在工作流运行时触发，确保标签也被保存
|
|
current_labels_from_inputs = []
|
|
for i in range(1, MAX_INPUTS + 1):
|
|
label = kwargs.get(f"label_{i}", "").strip()
|
|
if label:
|
|
current_labels_from_inputs.append(label)
|
|
write_labels_to_file(current_labels_from_inputs)
|
|
|
|
# 确定输出
|
|
final_labels = read_labels_from_file()
|
|
output_index = 0
|
|
output_label = ""
|
|
if select_label in final_labels:
|
|
output_index = final_labels.index(select_label) + 1
|
|
output_label = select_label
|
|
|
|
return (output_index, output_label)
|
|
|
|
# --- API Route for Saving Labels ---
|
|
@PromptServer.instance.routes.post("/switch_any/save_labels")
|
|
async def save_labels_route(request):
|
|
data = await request.json()
|
|
labels = data.get("labels", [])
|
|
|
|
if isinstance(labels, list):
|
|
try:
|
|
write_labels_to_file(labels)
|
|
return web.json_response({"status": "ok", "message": "Labels saved."})
|
|
except Exception as e:
|
|
return web.json_response({"status": "error", "message": str(e)}, status=500)
|
|
|
|
return web.json_response({"status": "error", "message": "Invalid data format, 'labels' must be a list."}, status=400)
|
|
|
|
# --- Node Registration ---
|
|
NODE_CLASS_MAPPINGS = { "SwitchAny": SwitchAny }
|
|
NODE_DISPLAY_NAME_MAPPINGS = { "SwitchAny": "Switch Any" }
|
|
|
|
# --- Web Directory Setup for JS file ---
|
|
WEB_DIRECTORY = "js" |