reload embedded client

Author: ad-astra-video

server/api/api.py

@@ -1,15 +1,17 @@
-from pathlib import Path
-import os
-import json
-from git import Repo
-import logging
-import sys
 from aiohttp import web
 
 from api.nodes.nodes import list_nodes, install_node, delete_node
 from api.models.models import list_models, add_model, delete_model
 from api.settings.settings import set_twilio_account_info
 
+from comfy.nodes.package_typing import ExportedNodes
+from comfy.nodes.package import _comfy_nodes, import_all_nodes_in_workspace
+from comfy.cmd.execution import nodes
+
+from api.nodes.nodes import force_import_all_nodes_in_workspace
+#use a different node import
+import_all_nodes_in_workspace = force_import_all_nodes_in_workspace
+
 def add_routes(app):
     app.router.add_get("/env/list_nodes", nodes)
     app.router.add_post("/env/install_nodes", install_nodes)
@@ -19,9 +21,38 @@ def add_routes(app):
     app.router.add_post("/env/add_models", add_models)
     app.router.add_post("/env/delete_models", delete_models)
 
+    app.router.add_post("/env/reload", reload)
     app.router.add_post("/env/set_account_info", set_account_info)
 
+async def reload(request):
+    '''
+    Reload ComfyUI environment
 
+    '''
+    
+    #reset embedded client
+    from comfy.client.embedded_comfy_client import EmbeddedComfyClient
+    from comfy.cli_args_types import Configuration
+    await request.app["pipeline"].client.comfy_client.__aexit__()
+    
+    #reset imports to clear imported nodes
+    global _comfy_nodes
+    if len(_comfy_nodes) > 0:
+        import sys
+        import importlib
+        del sys.modules['comfy.nodes.package']
+        del sys.modules['comfy.cmd.execution']
+        globals()['comfy.nodes.package'] = importlib.import_module('comfy.nodes.package')
+        globals()['comfy.cmd.execution'] = importlib.import_module('comfy.cmd.execution')
+        #use a different node import
+        import_all_nodes_in_workspace = force_import_all_nodes_in_workspace
+        _comfy_nodes = import_all_nodes_in_workspace()
+    
+    #load new embedded client
+    request.app["pipeline"].client.comfy_client = EmbeddedComfyClient(Configuration(cwd=request.app["workspace"], disable_cuda_malloc=True, gpu_only=True))
+        
+    return web.json_response({"success": True, "error": None})
+    
 async def nodes(request):
     '''
     List all custom nodes in the workspace
@@ -87,6 +118,24 @@ async def install_nodes(request):
         for node in nodes:
             await install_node(node, workspace_dir)
             installed_nodes.append(node['url'])
+        
+        #restart embedded client        
+        #request.app["pipeline"].client.set_prompt(request.app["pipeline"].client.prompt)
+        
+        # reimport nodes to workspace
+        #from comfy.cmd.execution import nodes
+        #from comfy.nodes.package import import_all_nodes_in_workspace
+        #update_nodes = import_all_nodes_in_workspace()
+        #nodes.update(update_nodes)
+
+        #config = request.app["pipeline"].client.comfy_client._configuration
+        #from comfy.client.embedded_comfy_client import EmbeddedComfyClient
+        #request.app["pipeline"].client.comfy_client = EmbeddedComfyClient(config)
+        #await request.app["pipeline"].warm()
+
+        #from pipeline import Pipeline
+        #request.app["pipeline"] = Pipeline(cwd=request.app["workspace"], disable_cuda_malloc=True, gpu_only=True)
+
         return web.json_response({"success": True, "error": None, "installed_nodes": installed_nodes})
     except Exception as e:
         return web.json_response({"success": False, "error": str(e), "installed_nodes": installed_nodes}, status=500)

server/api/nodes/nodes.py

@@ -127,3 +127,59 @@ async def delete_node(node, workspace_dir):
     except Exception as e:
         logger.error(f"error deleting node {node['name']}")
         raise Exception(f"error deleting node: {e}")
+
+
+from comfy.nodes.package import ExportedNodes
+from comfy.nodes.package import _comfy_nodes, _import_and_enumerate_nodes_in_module
+from functools import reduce
+from importlib.metadata import entry_points
+import types
+
+def force_import_all_nodes_in_workspace(vanilla_custom_nodes=True, raise_on_failure=False) -> ExportedNodes:
+    # now actually import the nodes, to improve control of node loading order
+    from comfy_extras import nodes as comfy_extras_nodes  # pylint: disable=absolute-import-used
+    from comfy.cli_args import args
+    from comfy.nodes import base_nodes
+    from comfy.nodes.vanilla_node_importing import mitigated_import_of_vanilla_custom_nodes
+
+    # only load these nodes once
+
+    base_and_extra = reduce(lambda x, y: x.update(y),
+                            map(lambda module_inner: _import_and_enumerate_nodes_in_module(module_inner, raise_on_failure=raise_on_failure), [
+                                # this is the list of default nodes to import
+                                base_nodes,
+                                comfy_extras_nodes
+                            ]),
+                            ExportedNodes())
+    custom_nodes_mappings = ExportedNodes()
+
+    if args.disable_all_custom_nodes:
+        logging.info("Loading custom nodes was disabled, only base and extra nodes were loaded")
+        _comfy_nodes.update(base_and_extra)
+        return _comfy_nodes
+
+    # load from entrypoints
+    for entry_point in entry_points().select(group='comfyui.custom_nodes'):
+        # Load the module associated with the current entry point
+        try:
+            module = entry_point.load()
+        except ModuleNotFoundError as module_not_found_error:
+            logging.error(f"A module was not found while importing nodes via an entry point: {entry_point}. Please ensure the entry point in setup.py is named correctly", exc_info=module_not_found_error)
+            continue
+
+        # Ensure that what we've loaded is indeed a module
+        if isinstance(module, types.ModuleType):
+            custom_nodes_mappings.update(
+                _import_and_enumerate_nodes_in_module(module, print_import_times=True))
+
+    # load the vanilla custom nodes last
+    if vanilla_custom_nodes:
+        custom_nodes_mappings += mitigated_import_of_vanilla_custom_nodes()
+
+    # don't allow custom nodes to overwrite base nodes
+    custom_nodes_mappings -= base_and_extra
+
+    _comfy_nodes.update(base_and_extra + custom_nodes_mappings)
+    
+    return _comfy_nodes
+