新規作成

nodes.py

@@ -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",
+}