新規作成

2026-05-08 20:09:17 +09:00
parent 5b528f7a58
commit d2986823ce
4 changed files with 587 additions and 1 deletions
@@ -1,2 +1,50 @@
 # ComfyUI-CheckpointWidgetRefresh
-ComfyUIを再起動せずに、checkpoint系の選択リストを現在のファイル状態へ同期する小さなUIノードです。
+
+ComfyUIを再起動せずに、Checkpoint系の選択リストを現在のファイル状態へ同期する小さなUIノードじゃ！
+
+「新しいモデルをディレクトリに置いたのに、リストに出てこないのう……再起動は面倒じゃし……」
+そんなおぬしの“平常運転な悩み”を、フロントエンドだけでなくバックエンドまで含めた鮮やかなパッチワークで解決して進ぜよう！
+
+## 特徴
+
+- **再起動いらずの即時同期**  
+  ボタンをポチッと押すだけで、新しく追加したモデルが選択肢に並び、削除したモデルは綺麗に消えるぞ。
+
+- **標準ノードにも対応**  
+  自作ノードだけでなく、標準の `CheckpointLoaderSimple` も更新対象じゃ。
+
+- **バックエンドまで貫くパッチ**  
+  ブラウザ上の表示を変えるだけではない。Python側の型定義（`INPUT_TYPES` など）まで更新するため、バリデーションで弾かれる心配も少ないのじゃ。
+
+- **徹底した軽量設計**  
+  常時監視やポーリングは一切しない。おぬしが「今だ！」とボタンを押した時だけ動く、ワークフローの邪魔をしない奥ゆかしい実用品じゃ。
+
+## 導入方法
+
+ComfyUIの `custom_nodes` ディレクトリで、以下のコマンドを打ち込むのじゃ！
+
+```bash
+git clone https://github.com/ruminar/ComfyUI-CheckpointWidgetRefresh.git
+````
+
+## 使い方
+
+1. **Checkpointを追加・削除する**
+   いつものようにモデルフォルダを整理するのじゃ。
+
+2. **ノードを追加する**
+   `Checkpoint Widget Refresh` ノードをキャンバスのどこかに置く。
+   （配線は不要じゃ！）
+
+3. **ボタンを押す**
+   `Refresh Checkpoint Widgets` ボタンを押せば、対象ノードの選択肢が最新状態に更新されるぞ！
+
+## 対応ノード
+
+* `CheckpointLoaderSimple`
+* `CheckpointNameSelector` 系ノード
+
+## ライセンス
+
+GPL-3.0
+（ComfyUI本体の掟に従っておるぞ！）
@@ -0,0 +1,9 @@
+from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
+
+WEB_DIRECTORY = "./web"
+
+__all__ = [
+    "NODE_CLASS_MAPPINGS",
+    "NODE_DISPLAY_NAME_MAPPINGS",
+    "WEB_DIRECTORY",
+]
@@ -0,0 +1,178 @@
+import os
+from aiohttp import web
+
+import folder_paths
+from server import PromptServer
+
+
+def _get_supported_checkpoint_extensions():
+    default_extensions = {".ckpt", ".pt", ".bin", ".pth", ".safetensors"}
+    extensions = getattr(folder_paths, "supported_pt_extensions", None)
+    if not extensions:
+        return default_extensions
+    return {str(ext).lower() for ext in extensions}
+
+
+def _clear_checkpoint_filename_cache():
+    cache = getattr(folder_paths, "filename_list_cache", None)
+    if isinstance(cache, dict):
+        cache.pop("checkpoints", None)
+
+
+def _scan_registered_checkpoint_folders():
+    """
+    Scan only ComfyUI-registered checkpoint folders.
+
+    This intentionally does NOT accept arbitrary user-supplied paths.
+    It should pick up files added/removed under normal checkpoints paths,
+    including paths registered through extra_model_paths.yaml.
+    """
+    results = set()
+    extensions = _get_supported_checkpoint_extensions()
+
+    try:
+        roots = folder_paths.get_folder_paths("checkpoints")
+    except Exception:
+        roots = []
+
+    for root in roots:
+        if not root or not os.path.isdir(root):
+            continue
+
+        for dirpath, _dirnames, filenames in os.walk(root):
+            for filename in filenames:
+                _, ext = os.path.splitext(filename)
+                if ext.lower() not in extensions:
+                    continue
+
+                full_path = os.path.join(dirpath, filename)
+                rel_path = os.path.relpath(full_path, root)
+                rel_path = rel_path.replace("\\", "/")
+                results.add(rel_path)
+
+    return sorted(results, key=lambda value: value.lower())
+
+
+def _get_fresh_checkpoints():
+    """
+    Prefer ComfyUI's own list after clearing its filename cache.
+    Merge direct scanning as a safety net for stale model-library states.
+    """
+    _clear_checkpoint_filename_cache()
+
+    results = set()
+
+    try:
+        results.update(folder_paths.get_filename_list("checkpoints"))
+    except Exception:
+        pass
+
+    try:
+        results.update(_scan_registered_checkpoint_folders())
+    except Exception:
+        pass
+
+    return sorted(results, key=lambda value: value.lower())
+
+
+def _patch_backend_checkpoint_widget_classes(checkpoints):
+    """
+    Patch backend class contracts for checkpoint combo nodes.
+
+    Updating only the browser-side widget list is not enough.
+    ComfyUI backend prompt validation can still use stale INPUT_TYPES / RETURN_TYPES.
+
+    Targets:
+      - CheckpointNameSelector-style nodes
+      - CheckpointLoaderSimple
+    """
+    patched = []
+
+    try:
+        import nodes as comfy_nodes
+    except Exception:
+        return patched
+
+    mappings = getattr(comfy_nodes, "NODE_CLASS_MAPPINGS", {})
+    checkpoint_values = list(checkpoints)
+
+    for class_name, node_class in list(mappings.items()):
+        name = str(class_name)
+
+        try:
+            if "CheckpointNameSelector" in name:
+                def selector_input_types(cls, _checkpoint_values=checkpoint_values):
+                    return {
+                        "required": {
+                            "ckpt_name": (list(_checkpoint_values),),
+                        },
+                    }
+
+                node_class.INPUT_TYPES = classmethod(selector_input_types)
+                node_class.RETURN_TYPES = (list(checkpoint_values), "STRING")
+                patched.append(name)
+                continue
+
+            if name == "CheckpointLoaderSimple":
+                def loader_input_types(cls, _checkpoint_values=checkpoint_values):
+                    return {
+                        "required": {
+                            "ckpt_name": (list(_checkpoint_values),),
+                        },
+                    }
+
+                node_class.INPUT_TYPES = classmethod(loader_input_types)
+                patched.append(name)
+                continue
+
+        except Exception as exc:
+            print(f"[CheckpointWidgetRefresh] failed to patch {name}: {exc}")
+
+    return patched
+
+
+routes = PromptServer.instance.routes
+
+
+@routes.get("/checkpoint_widget_refresh/checkpoints")
+async def checkpoint_widget_refresh_checkpoints(_request):
+    checkpoints = _get_fresh_checkpoints()
+    patched = _patch_backend_checkpoint_widget_classes(checkpoints)
+
+    return web.json_response({
+        "checkpoints": checkpoints,
+        "count": len(checkpoints),
+        "patched_classes": patched,
+        "patched_class_count": len(patched),
+    })
+
+
+class CheckpointWidgetRefreshPanel:
+    """
+    UI-only panel node.
+
+    The refresh button itself is implemented in web/checkpoint_widget_refresh.js.
+    This backend node exists so the button can be placed inside a workflow as a small panel.
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {},
+        }
+
+    RETURN_TYPES = ()
+    FUNCTION = "noop"
+    CATEGORY = "utils/model"
+
+    def noop(self):
+        return ()
+
+
+NODE_CLASS_MAPPINGS = {
+    "CheckpointWidgetRefreshPanel": CheckpointWidgetRefreshPanel,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "CheckpointWidgetRefreshPanel": "Checkpoint Widget Refresh",
+}
@@ -0,0 +1,351 @@
+import { app } from "../../scripts/app.js";
+import { api } from "../../scripts/api.js";
+
+const EXTENSION_NAME = "ruminar.CheckpointWidgetRefresh";
+const REFRESHER_NODE_NAME = "CheckpointWidgetRefreshPanel";
+
+let lastCheckpointList = null;
+
+function getNodeTypeName(node) {
+    return String(
+        node?.type ??
+        node?.comfyClass ??
+        node?.constructor?.type ??
+        ""
+    );
+}
+
+function getNodeDataName(nodeData) {
+    return String(nodeData?.name ?? "");
+}
+
+function isRefresherNode(node) {
+    return getNodeTypeName(node) === REFRESHER_NODE_NAME;
+}
+
+function findCheckpointWidget(node) {
+    const widgets = node?.widgets ?? [];
+    return widgets.find((widget) =>
+        widget?.name === "ckpt_name" &&
+        Array.isArray(widget?.options?.values)
+    );
+}
+
+function outputNames(node) {
+    return (node?.outputs ?? []).map((output) =>
+        String(output?.name ?? "").toLowerCase()
+    );
+}
+
+function hasSelectorOutputs(node) {
+    const names = outputNames(node);
+    return names.includes("ckpt_name") && names.includes("ckpt_name_str");
+}
+
+function hasLoaderOutputs(node) {
+    const names = outputNames(node);
+    return names.includes("model") && names.includes("clip") && names.includes("vae");
+}
+
+function shouldTargetNode(node) {
+    if (!node || isRefresherNode(node)) {
+        return false;
+    }
+
+    const typeName = getNodeTypeName(node);
+
+    if (typeName.includes("CheckpointNameSelector")) {
+        return true;
+    }
+
+    if (typeName === "CheckpointLoaderSimple") {
+        return true;
+    }
+
+    if (hasSelectorOutputs(node) && Boolean(findCheckpointWidget(node))) {
+        return true;
+    }
+
+    if (hasLoaderOutputs(node) && Boolean(findCheckpointWidget(node))) {
+        return true;
+    }
+
+    return false;
+}
+
+function shouldPatchFutureNodeType(nodeData) {
+    const name = getNodeDataName(nodeData);
+    if (name === REFRESHER_NODE_NAME) {
+        return true;
+    }
+
+    return name.includes("CheckpointNameSelector") || name === "CheckpointLoaderSimple";
+}
+
+function chooseReplacement(oldValues, newValues, currentValue) {
+    if (!Array.isArray(newValues) || newValues.length === 0) {
+        return "";
+    }
+
+    if (newValues.includes(currentValue)) {
+        return currentValue;
+    }
+
+    const oldIndex = Array.isArray(oldValues) ? oldValues.indexOf(currentValue) : -1;
+
+    if (oldIndex >= 0) {
+        // Prefer the next surviving item from the old list.
+        for (let index = oldIndex + 1; index < oldValues.length; index++) {
+            const candidate = oldValues[index];
+            if (newValues.includes(candidate)) {
+                return candidate;
+            }
+        }
+
+        // If the deleted item was at the end, choose the previous survivor.
+        for (let index = oldIndex - 1; index >= 0; index--) {
+            const candidate = oldValues[index];
+            if (newValues.includes(candidate)) {
+                return candidate;
+            }
+        }
+
+        // Last fallback: same numeric position in the new list if possible.
+        return newValues[Math.min(oldIndex, newValues.length - 1)] ?? newValues[0];
+    }
+
+    return newValues[0];
+}
+
+function arraysEqual(left, right) {
+    if (!Array.isArray(left) || !Array.isArray(right)) {
+        return false;
+    }
+    if (left.length !== right.length) {
+        return false;
+    }
+    for (let index = 0; index < left.length; index++) {
+        if (left[index] !== right[index]) {
+            return false;
+        }
+    }
+    return true;
+}
+
+function patchCheckpointSlotTypes(node, checkpoints) {
+    const values = [...checkpoints];
+
+    for (const output of node.outputs ?? []) {
+        if (String(output?.name ?? "") !== "ckpt_name") {
+            continue;
+        }
+
+        output.type = [...values];
+
+        const links = Array.isArray(output.links) ? output.links : [];
+        for (const linkId of links) {
+            const link = app.graph?.links?.[linkId];
+            if (!link) {
+                continue;
+            }
+
+            const targetNode = app.graph?.getNodeById?.(link.target_id);
+            const targetInput = targetNode?.inputs?.[link.target_slot];
+            if (targetInput && String(targetInput.name ?? "").toLowerCase().includes("ckpt")) {
+                targetInput.type = [...values];
+            }
+        }
+    }
+
+    for (const input of node.inputs ?? []) {
+        if (String(input?.name ?? "").toLowerCase().includes("ckpt")) {
+            input.type = [...values];
+        }
+    }
+}
+
+function applyCheckpointListToNode(node, checkpoints) {
+    if (!shouldTargetNode(node)) {
+        return {
+            matched: false,
+            changed: false,
+            valueChanged: false,
+        };
+    }
+
+    const widget = findCheckpointWidget(node);
+    if (!widget) {
+        return {
+            matched: true,
+            changed: false,
+            valueChanged: false,
+            reason: "checkpoint widget not found",
+        };
+    }
+
+    if (!widget.options) {
+        widget.options = {};
+    }
+
+    const oldValues = Array.isArray(widget.options.values)
+        ? [...widget.options.values]
+        : [];
+
+    const oldValue = widget.value;
+    const newValues = [...checkpoints];
+    const newValue = chooseReplacement(oldValues, newValues, oldValue);
+
+    widget.options.values = newValues;
+    widget.value = newValue;
+    patchCheckpointSlotTypes(node, newValues);
+
+    const valuesChanged = !arraysEqual(oldValues, newValues);
+    const valueChanged = oldValue !== newValue;
+
+    if (valueChanged && typeof widget.callback === "function") {
+        try {
+            widget.callback(widget.value);
+        } catch (error) {
+            console.warn("[CheckpointWidgetRefresh] widget callback failed", error);
+        }
+    }
+
+    if (valuesChanged || valueChanged) {
+        node.setDirtyCanvas?.(true, true);
+    }
+
+    return {
+        matched: true,
+        changed: valuesChanged || valueChanged,
+        valueChanged,
+        oldValue,
+        newValue,
+    };
+}
+
+async function fetchCheckpointList() {
+    const response = await api.fetchApi("/checkpoint_widget_refresh/checkpoints", {
+        cache: "no-store",
+    });
+
+    if (!response.ok) {
+        throw new Error(`HTTP ${response.status}`);
+    }
+
+    const payload = await response.json();
+    return {
+        checkpoints: Array.isArray(payload.checkpoints) ? payload.checkpoints : [],
+        patchedClassCount: payload.patched_class_count ?? 0,
+        patchedClasses: Array.isArray(payload.patched_classes) ? payload.patched_classes : [],
+    };
+}
+
+function updatePanelTitle(panelNode, text) {
+    if (!panelNode) {
+        return;
+    }
+
+    panelNode.title = text;
+    panelNode.setDirtyCanvas?.(true, true);
+    app.graph?.setDirtyCanvas?.(true, true);
+}
+
+async function refreshCheckpointWidgets(panelNode = null) {
+    updatePanelTitle(panelNode, "Refreshing checkpoint widgets...");
+
+    const payload = await fetchCheckpointList();
+    const checkpoints = payload.checkpoints;
+    lastCheckpointList = checkpoints;
+
+    let matched = 0;
+    let changed = 0;
+    let valueChanged = 0;
+
+    for (const node of app.graph?._nodes ?? []) {
+        const result = applyCheckpointListToNode(node, checkpoints);
+        if (!result.matched) {
+            continue;
+        }
+
+        matched += 1;
+        if (result.changed) {
+            changed += 1;
+        }
+        if (result.valueChanged) {
+            valueChanged += 1;
+            console.log(
+                `[CheckpointWidgetRefresh] ${getNodeTypeName(node)} #${node.id}: ` +
+                `${result.oldValue} -> ${result.newValue}`
+            );
+        }
+    }
+
+    const title = `Updated ${changed}/${matched} ckpt widgets, ${checkpoints.length} ckpts`;
+    updatePanelTitle(panelNode, title);
+
+    console.log(
+        `[CheckpointWidgetRefresh] ${title}. ` +
+        `valueChanged=${valueChanged}, backend patched=${payload.patchedClassCount}`,
+        payload.patchedClasses
+    );
+
+    app.graph?.setDirtyCanvas?.(true, true);
+}
+
+function installRefreshPanelWidgets(node) {
+    if (!node || node._checkpointWidgetRefreshWidgetsInstalled) {
+        return;
+    }
+
+    node.addWidget(
+        "button",
+        "Refresh Checkpoint Widgets",
+        "refresh",
+        async () => {
+            try {
+                await refreshCheckpointWidgets(node);
+            } catch (error) {
+                console.error("[CheckpointWidgetRefresh] refresh failed", error);
+                updatePanelTitle(node, "Refresh failed. See console.");
+            }
+        }
+    );
+
+    node.title = "Checkpoint Widget Refresh";
+    node.size = [320, 60];
+    node._checkpointWidgetRefreshWidgetsInstalled = true;
+}
+
+function installFuturePatchHook(nodeType, nodeData) {
+    if (!shouldPatchFutureNodeType(nodeData)) {
+        return;
+    }
+
+    const originalOnNodeCreated = nodeType.prototype.onNodeCreated;
+
+    nodeType.prototype.onNodeCreated = function (...args) {
+        const result = originalOnNodeCreated?.apply(this, args);
+
+        if (getNodeDataName(nodeData) === REFRESHER_NODE_NAME) {
+            installRefreshPanelWidgets(this);
+            return result;
+        }
+
+        if (lastCheckpointList) {
+            setTimeout(() => {
+                applyCheckpointListToNode(this, lastCheckpointList);
+                app.graph?.setDirtyCanvas?.(true, true);
+            }, 0);
+        }
+
+        return result;
+    };
+}
+
+app.registerExtension({
+    name: EXTENSION_NAME,
+
+    beforeRegisterNodeDef(nodeType, nodeData) {
+        installFuturePatchHook(nodeType, nodeData);
+    },
+});