serve workflow templates from custom_nodes (#6193)

* add GET /workflow_templates

* serve workflow templates from custom_nodes

* refactor into custom_node_manager, add test

* remove unused import

* revert changes in folder_paths

* Remove trailing whitespace.

* account for multiple custom_nodes paths

app/custom_node_manager.py

@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+import os
+import folder_paths
+import glob
+from aiohttp import web
+
+class CustomNodeManager:
+    """
+    Placeholder to refactor the custom node management features from ComfyUI-Manager.
+    Currently it only contains the custom workflow templates feature.
+    """
+    def add_routes(self, routes, webapp, loadedModules):     
+
+        @routes.get("/workflow_templates")
+        async def get_workflow_templates(request):
+            """Returns a web response that contains the map of custom_nodes names and their associated workflow templates. The ones without templates are omitted."""
+            files = [
+                file
+                for folder in folder_paths.get_folder_paths("custom_nodes")
+                for file in glob.glob(os.path.join(folder, '*/example_workflows/*.json'))
+            ]
+            workflow_templates_dict = {} # custom_nodes folder name -> example workflow names
+            for file in files:
+                custom_nodes_name = os.path.basename(os.path.dirname(os.path.dirname(file)))
+                workflow_name = os.path.splitext(os.path.basename(file))[0]
+                workflow_templates_dict.setdefault(custom_nodes_name, []).append(workflow_name)
+            return web.json_response(workflow_templates_dict)
+        
+        # Serve workflow templates from custom nodes.
+        for module_name, module_dir in loadedModules:
+            workflows_dir = os.path.join(module_dir, 'example_workflows')
+            if os.path.exists(workflows_dir):
+                webapp.add_routes([web.static('/api/workflow_templates/' + module_name, workflows_dir)])

nodes.py

@@ -2047,6 +2047,9 @@ NODE_DISPLAY_NAME_MAPPINGS = {
 
 EXTENSION_WEB_DIRS = {}
 
+# Dictionary of successfully loaded module names and associated directories.
+LOADED_MODULE_DIRS = {}
+
 
 def get_module_name(module_path: str) -> str:
     """
@@ -2088,6 +2091,8 @@ def load_custom_node(module_path: str, ignore=set(), module_parent="custom_nodes
         sys.modules[module_name] = module
         module_spec.loader.exec_module(module)
 
+        LOADED_MODULE_DIRS[module_name] = os.path.abspath(module_dir)
+
         if hasattr(module, "WEB_DIRECTORY") and getattr(module, "WEB_DIRECTORY") is not None:
             web_dir = os.path.abspath(os.path.join(module_dir, getattr(module, "WEB_DIRECTORY")))
             if os.path.isdir(web_dir):

server.py

@@ -30,6 +30,7 @@ import node_helpers
 from app.frontend_management import FrontendManager
 from app.user_manager import UserManager
 from app.model_manager import ModelFileManager
+from app.custom_node_manager import CustomNodeManager
 from typing import Optional
 from api_server.routes.internal.internal_routes import InternalRoutes
 
@@ -153,6 +154,7 @@ class PromptServer():
 
         self.user_manager = UserManager()
         self.model_file_manager = ModelFileManager()
+        self.custom_node_manager = CustomNodeManager()
         self.internal_routes = InternalRoutes(self)
         self.supports = ["custom_nodes_from_web"]
         self.prompt_queue = None
@@ -697,6 +699,7 @@ class PromptServer():
     def add_routes(self):
         self.user_manager.add_routes(self.routes)
         self.model_file_manager.add_routes(self.routes)
+        self.custom_node_manager.add_routes(self.routes, self.app, nodes.LOADED_MODULE_DIRS.items())
         self.app.add_subapp('/internal', self.internal_routes.get_app())
 
         # Prefix every route with /api for easier matching for delegation.
@@ -713,6 +716,7 @@ class PromptServer():
         self.app.add_routes(api_routes)
         self.app.add_routes(self.routes)
 
+        # Add routes from web extensions.
         for name, dir in nodes.EXTENSION_WEB_DIRS.items():
             self.app.add_routes([web.static('/extensions/' + name, dir)])
 

tests-unit/app_test/custom_node_manager_test.py

@@ -0,0 +1,40 @@
+import pytest
+from aiohttp import web
+from unittest.mock import patch
+from app.custom_node_manager import CustomNodeManager
+
+pytestmark = (
+    pytest.mark.asyncio
+)  # This applies the asyncio mark to all test functions in the module
+
+@pytest.fixture
+def custom_node_manager():
+    return CustomNodeManager()
+
+@pytest.fixture
+def app(custom_node_manager):
+    app = web.Application()
+    routes = web.RouteTableDef()
+    custom_node_manager.add_routes(routes, app, [("ComfyUI-TestExtension1", "ComfyUI-TestExtension1")])
+    app.add_routes(routes)
+    return app
+
+async def test_get_workflow_templates(aiohttp_client, app, tmp_path):
+    client = await aiohttp_client(app)
+    # Setup temporary custom nodes file structure with 1 workflow file
+    custom_nodes_dir = tmp_path / "custom_nodes"
+    example_workflows_dir = custom_nodes_dir / "ComfyUI-TestExtension1" / "example_workflows"
+    example_workflows_dir.mkdir(parents=True)
+    template_file = example_workflows_dir / "workflow1.json"
+    template_file.write_text('')
+
+    with patch('folder_paths.folder_names_and_paths', {
+        'custom_nodes': ([str(custom_nodes_dir)], None)
+    }):
+        response = await client.get('/workflow_templates')
+        assert response.status == 200
+        workflows_dict = await response.json()
+        assert isinstance(workflows_dict, dict)
+        assert "ComfyUI-TestExtension1" in workflows_dict
+        assert isinstance(workflows_dict["ComfyUI-TestExtension1"], list)
+        assert workflows_dict["ComfyUI-TestExtension1"][0] == "workflow1"