Merge branch 'master' of github.com:comfyanonymous/ComfyUI

2026-07-21 23:41:28 +08:00 · 2024-09-23 12:50:31 -07:00
parent 4bdc208f29 3a0eeee320
commit fa3176f96f
51 changed files with 51256 additions and 44712 deletions
@@ -51,7 +51,46 @@ temp_directory = os.path.join(get_base_path(), "temp")
 input_directory = os.path.join(get_base_path(), "input")
 user_directory = os.path.join(get_base_path(), "user")

-_filename_list_cache = {}
+filename_list_cache = {}
+
+
+class CacheHelper:
+    """
+    Helper class for managing file list cache data.
+    """
+
+    def __init__(self):
+        self.cache: dict[str, tuple[list[str], dict[str, float], float]] = {}
+        self.active = False
+
+    def get(self, key: str, default=None) -> tuple[list[str], dict[str, float], float]:
+        if not self.active:
+            return default
+        return self.cache.get(key, default)
+
+    def set(self, key: str, value: tuple[list[str], dict[str, float], float]) -> None:
+        if self.active:
+            self.cache[key] = value
+
+    def clear(self):
+        self.cache.clear()
+
+    def __enter__(self):
+        self.active = True
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.active = False
+        self.clear()
+
+
+cache_helper = CacheHelper()
+
+
+def map_legacy(folder_name: str) -> str:
+    legacy = {"unet": "diffusion_models"}
+    return legacy.get(folder_name, folder_name)
+

 if not os.path.exists(input_directory):
    try:
@@ -150,7 +189,7 @@ def exists_annotated_filepath(name):
    return os.path.exists(filepath)


-def add_model_folder_path(folder_name, full_folder_path: Optional[str] = None, extensions: Optional[set[str]] = None) -> str:
+def add_model_folder_path(folder_name, full_folder_path: Optional[str] = None, extensions: Optional[set[str]] = None, is_default: bool = False) -> str:
    """
    Registers a model path for the given canonical name.
    :param folder_name: the folder name
@@ -165,7 +204,10 @@ def add_model_folder_path(folder_name, full_folder_path: Optional[str] = None, e

    folder_path = folder_names_and_paths[folder_name]
    if full_folder_path not in folder_path.paths:
-        folder_path.paths.append(full_folder_path)
+        if is_default:
+            folder_path.paths.insert(0, full_folder_path)
+        else:
+            folder_path.paths.append(full_folder_path)

    if extensions is not None:
        folder_path.supported_extensions |= extensions
@@ -244,7 +286,15 @@ def get_full_path(folder_name, filename) -> Optional[str | bytes | os.PathLike]:
    return None


-def get_filename_list_(folder_name):
+def get_full_path_or_raise(folder_name: str, filename: str) -> str:
+    full_path = get_full_path(folder_name, filename)
+    if full_path is None:
+        raise FileNotFoundError(f"Model in folder '{folder_name}' with filename '{filename}' not found.")
+    return full_path
+
+
+def get_filename_list_(folder_name: str) -> tuple[list[str], dict[str, float], float]:
+    folder_name = map_legacy(folder_name)
    global folder_names_and_paths
    output_list = set()
    folders = folder_names_and_paths[folder_name]
@@ -257,12 +307,17 @@ def get_filename_list_(folder_name):
    return sorted(list(output_list)), output_folders, time.perf_counter()


-def cached_filename_list_(folder_name):
-    global _filename_list_cache
+def cached_filename_list_(folder_name: str) -> tuple[list[str], dict[str, float], float] | None:
+    strong_cache = cache_helper.get(folder_name)
+    if strong_cache is not None:
+        return strong_cache
+
+    global filename_list_cache
    global folder_names_and_paths
-    if folder_name not in _filename_list_cache:
+    folder_name = map_legacy(folder_name)
+    if folder_name not in filename_list_cache:
        return None
-    out = _filename_list_cache[folder_name]
+    out = filename_list_cache[folder_name]

    for x in out[1]:
        time_modified = out[1][x]
@@ -279,12 +334,14 @@ def cached_filename_list_(folder_name):
    return out


-def get_filename_list(folder_name):
+def get_filename_list(folder_name: str) -> list[str]:
+    folder_name = map_legacy(folder_name)
    out = cached_filename_list_(folder_name)
    if out is None:
        out = get_filename_list_(folder_name)
-        global _filename_list_cache
-        _filename_list_cache[folder_name] = out
+        global filename_list_cache
+        filename_list_cache[folder_name] = out
+    cache_helper.set(folder_name, out)
    return list(out[0])


@@ -345,8 +402,8 @@ def create_directories():


 def invalidate_cache(folder_name):
-    global _filename_list_cache
-    _filename_list_cache.pop(folder_name, None)
+    global filename_list_cache
+    filename_list_cache.pop(folder_name, None)


 def filter_files_content_types(files: list[str], content_types: list[Literal["image", "video", "audio"]]) -> list[str]:
@@ -82,7 +82,10 @@ def prompt_worker(q: AbstractPromptQueue, _server: server_module.PromptServer):


 async def run(server, address='', port=8188, verbose=True, call_on_start=None):
-    await asyncio.gather(server.start(address, port, verbose, call_on_start), server.publish_loop())
+    addresses = []
+    for addr in address.split(","):
+        addresses.append((addr, port))
+    await asyncio.gather(server.start_multi_address(addresses, call_on_start), server.publish_loop())


 def cleanup_temp():
@@ -223,12 +226,15 @@ async def main(from_script_dir: Optional[Path] = None):
            import webbrowser
            if os.name == 'nt' and address == '0.0.0.0' or address == '':
                address = '127.0.0.1'
+            if ':' in address:
+                address = "[{}]".format(address)
            webbrowser.open(f"http://{address}:{port}")

        call_on_start = startup_server

    server.address = args.listen
    server.port = args.port
+
    try:
        await server.setup()
        await run(server, address=args.listen, port=args.port, verbose=not args.dont_print_server,
@@ -29,7 +29,7 @@ from can_ada import URL, parse as urlparse  # pylint: disable=no-name-in-module
 from typing_extensions import NamedTuple

 from .latent_preview_image_encoding import encode_preview_image
-from .. import interruption
+from .. import interruption, model_management
 from .. import node_helpers
 from .. import utils
 from ..api_server.routes.internal.internal_routes import InternalRoutes
@@ -256,6 +256,12 @@ class PromptServer(ExecutorToClientProgress):
            embeddings = folder_paths.get_filename_list("embeddings")
            return web.json_response(list(map(lambda a: os.path.splitext(a)[0], embeddings)))

+        @routes.get("/models")
+        def list_model_types(request):
+            model_types = list(folder_paths.folder_names_and_paths.keys())
+
+            return web.json_response(model_types)
+
        @routes.get("/models/{folder}")
        async def get_models(request):
            folder = request.match_info.get("folder", None)
@@ -505,12 +511,17 @@ class PromptServer(ExecutorToClientProgress):
        async def system_stats(request):
            device = get_torch_device()
            device_name = get_torch_device_name(device)
+            cpu_device = model_management.torch.device("cpu")
+            ram_total = model_management.get_total_memory(cpu_device)
+            ram_free = model_management.get_free_memory(cpu_device)
            vram_total, torch_vram_total = get_total_memory(device, torch_total_too=True)
            vram_free, torch_vram_free = get_free_memory(device, torch_free_too=True)

            system_stats = {
                "system": {
                    "os": os.name,
+                    "ram_total": ram_total,
+                    "ram_free": ram_free,
                    "comfyui_version": get_comfyui_version(),
                    "python_version": sys.version,
                    "pytorch_version": torch_version,
@@ -568,14 +579,15 @@ class PromptServer(ExecutorToClientProgress):

        @routes.get("/object_info")
        async def get_object_info(request):
-            out = {}
-            for x in self.nodes.NODE_CLASS_MAPPINGS:
-                try:
-                    out[x] = node_info(x)
-                except Exception as e:
-                    logging.error(f"[ERROR] An error occurred while retrieving information for the '{x}' node.")
-                    logging.error(traceback.format_exc())
-            return web.json_response(out)
+            with folder_paths.cache_helper:
+                out = {}
+                for x in self.nodes.NODE_CLASS_MAPPINGS:
+                    try:
+                        out[x] = node_info(x)
+                    except Exception as e:
+                        logging.error(f"[ERROR] An error occurred while retrieving information for the '{x}' node.")
+                        logging.error(traceback.format_exc())
+                return web.json_response(out)

        @routes.get("/object_info/{node_class}")
        async def get_object_info_node(request):
@@ -969,17 +981,29 @@ class PromptServer(ExecutorToClientProgress):
            await self.send(*msg)

    async def start(self, address: str | None, port: int | None, verbose=True, call_on_start=None):
+        await self.start_multi_address([(address, port)], call_on_start=call_on_start, verbose=verbose)
+
+    async def start_multi_address(self, addresses, call_on_start=None, verbose=True):
        runner = web.AppRunner(self.app, access_log=None)
        await runner.setup()
-        site = web.TCPSite(runner, host=address, port=port)
-        await site.start()
+        for addr in addresses:
+            address = addr[0]
+            port = addr[1]
+            site = web.TCPSite(runner, address, port)
+            await site.start()

-        self.address = address
-        self.port = port
+            if not hasattr(self, 'address'):
+                self.address = address #TODO: remove this
+                self.port = port
+
+            if ':' in address:
+                address_print = "[{}]".format(address)
+            else:
+                address_print = address

        if verbose:
            logging.info("Starting server")
-            logging.info("To see the GUI go to: http://{}:{}".format("localhost" if address == "0.0.0.0" else address, port))
+            logging.info("To see the GUI go to: http://{}:{}".format("localhost" if address_print == "0.0.0.0" else address, port))
        if call_on_start is not None:
            call_on_start(address, port)