From 7362816ad2b672266b3f4bb8d1bc35c3201f9ee5 Mon Sep 17 00:00:00 2001 From: loxotron Date: Thu, 15 May 2025 07:15:16 +0300 Subject: [PATCH] speed up caching and loading --- glob/cnr_utils.py | 307 +++++++++++++++++++++++++++++----------------- 1 file changed, 193 insertions(+), 114 deletions(-) diff --git a/glob/cnr_utils.py b/glob/cnr_utils.py index 6a61aaa3..54c37427 100644 --- a/glob/cnr_utils.py +++ b/glob/cnr_utils.py @@ -5,121 +5,158 @@ import platform import time from dataclasses import dataclass from typing import List +from functools import lru_cache +from concurrent.futures import ThreadPoolExecutor +import requests import manager_core import manager_util -import requests import toml base_url = "https://api.comfy.org" +memory_cache = {} # In-memory cache for nodes and API responses +cnr_info_cache = {} # Cache for read_cnr_info results -lock = asyncio.Lock() +@lru_cache(maxsize=1) +def get_version_and_form_factor(): + """ + Get ComfyUI version and form factor once and cache result. -is_cache_loading = False + Returns: + tuple: (comfyui_version, form_factor) + """ + # Determine form factor based on environment and platform + is_desktop = bool(os.environ.get('__COMFYUI_DESKTOP_VERSION__')) + system = platform.system().lower() + is_windows = system == 'windows' + is_mac = system == 'darwin' + is_linux = system == 'linux' + + # Get ComfyUI version tag + if is_desktop: + # extract version from pyproject.toml instead of git tag + comfyui_ver = manager_core.get_current_comfyui_ver() or 'unknown' + else: + comfyui_ver = manager_core.get_comfyui_tag() or 'unknown' + + if is_desktop: + if is_windows: + form_factor = 'desktop-win' + elif is_mac: + form_factor = 'desktop-mac' + else: + form_factor = 'other' + else: + if is_windows: + form_factor = 'git-windows' + elif is_mac: + form_factor = 'git-mac' + elif is_linux: + form_factor = 'git-linux' + else: + form_factor = 'other' + + return comfyui_ver, form_factor + +def fetch_page(page, comfyui_ver, form_factor): + """ + Fetch a single page from ComfyRegistry using requests. + + Args: + page (int): Page number. + comfyui_ver (str): ComfyUI version. + form_factor (str): Form factor (e.g., desktop-win, git-linux). + + Returns: + dict: Page data or None if failed. + """ + # Add comfyui_version and form_factor to the API request + sub_uri = f"{base_url}/nodes?page={page}&limit=30&comfyui_version={comfyui_ver}&form_factor={form_factor}" + try: + response = requests.get(sub_uri, timeout=10) + if response.status_code == 200: + return response.json() + return None + except Exception: + return None + +async def fetch_all(comfyui_ver, form_factor): + """ + Fetch all nodes from ComfyRegistry using ThreadPoolExecutor. + + Args: + comfyui_ver (str): ComfyUI version. + form_factor (str): Form factor (e.g., desktop-win, git-linux). + + Returns: + dict: Dictionary of nodes. + """ + nodes = [] + first_page = fetch_page(1, comfyui_ver, form_factor) + if not first_page: + raise Exception("Failed to fetch first page") + + total_pages = first_page['totalPages'] + nodes.extend(first_page['nodes']) + + if total_pages > 1: + with ThreadPoolExecutor(max_workers=10) as executor: + tasks = [executor.submit(fetch_page, page, comfyui_ver, form_factor) + for page in range(2, total_pages + 1)] + for i, future in enumerate(tasks, 2): + result = future.result() + if result and 'nodes' in result: + nodes.extend(result['nodes']) + if i % 50 == 0: # Adjusted logging to reduce spam + print(f"FETCH ComfyRegistry Data: {i}/{total_pages}") + + print("FETCH ComfyRegistry Data [DONE]") + + for node in nodes: + if 'latest_version' not in node: + node['latest_version'] = dict(version='nightly') + + return {'nodes': nodes} async def get_cnr_data(cache_mode=True, dont_wait=True): - try: - return await _get_cnr_data(cache_mode, dont_wait) - except asyncio.TimeoutError: - print("A timeout occurred during the fetch process from ComfyRegistry.") - return await _get_cnr_data(cache_mode=True, dont_wait=True) # timeout fallback - -async def _get_cnr_data(cache_mode=True, dont_wait=True): - global is_cache_loading + """ + Fetch or load cached ComfyRegistry data. + Args: + cache_mode (bool): Use cache if True. + dont_wait (bool): Return immediately if cache is not ready. + + Returns: + list: List of node data. + """ uri = f'{base_url}/nodes' - - async def fetch_all(): - remained = True - page = 1 - - full_nodes = {} - - - # Determine form factor based on environment and platform - is_desktop = bool(os.environ.get('__COMFYUI_DESKTOP_VERSION__')) - system = platform.system().lower() - is_windows = system == 'windows' - is_mac = system == 'darwin' - is_linux = system == 'linux' - - # Get ComfyUI version tag - if is_desktop: - # extract version from pyproject.toml instead of git tag - comfyui_ver = manager_core.get_current_comfyui_ver() or 'unknown' - else: - comfyui_ver = manager_core.get_comfyui_tag() or 'unknown' - - if is_desktop: - if is_windows: - form_factor = 'desktop-win' - elif is_mac: - form_factor = 'desktop-mac' - else: - form_factor = 'other' - else: - if is_windows: - form_factor = 'git-windows' - elif is_mac: - form_factor = 'git-mac' - elif is_linux: - form_factor = 'git-linux' - else: - form_factor = 'other' - - while remained: - # Add comfyui_version and form_factor to the API request - sub_uri = f'{base_url}/nodes?page={page}&limit=30&comfyui_version={comfyui_ver}&form_factor={form_factor}' - sub_json_obj = await asyncio.wait_for(manager_util.get_data_with_cache(sub_uri, cache_mode=False, silent=True, dont_cache=True), timeout=30) - remained = page < sub_json_obj['totalPages'] - - for x in sub_json_obj['nodes']: - full_nodes[x['id']] = x - - if page % 5 == 0: - print(f"FETCH ComfyRegistry Data: {page}/{sub_json_obj['totalPages']}") - - page += 1 - time.sleep(0.5) - - print("FETCH ComfyRegistry Data [DONE]") - - for v in full_nodes.values(): - if 'latest_version' not in v: - v['latest_version'] = dict(version='nightly') - - return {'nodes': list(full_nodes.values())} - + if cache_mode: - is_cache_loading = True - cache_state = manager_util.get_cache_state(uri) - + # Check memory cache first + if uri in memory_cache: + return memory_cache[uri]['nodes'] + + cache_path = manager_util.get_cache_path(uri) + if os.path.exists(cache_path): + with open(cache_path, 'r', encoding="UTF-8", errors="ignore") as json_file: + data = json.load(json_file) + memory_cache[uri] = data + return data['nodes'] + if dont_wait: - if cache_state == 'not-cached': - return {} - else: - print("[ComfyUI-Manager] The ComfyRegistry cache update is still in progress, so an outdated cache is being used.") - with open(manager_util.get_cache_path(uri), 'r', encoding="UTF-8", errors="ignore") as json_file: - return json.load(json_file)['nodes'] - - if cache_state == 'cached': - with open(manager_util.get_cache_path(uri), 'r', encoding="UTF-8", errors="ignore") as json_file: - return json.load(json_file)['nodes'] + print("[ComfyUI-Manager] The ComfyRegistry cache update is still in progress, so an outdated cache is being used.") + return {} # Return empty if no cache and dont_wait=True try: - json_obj = await fetch_all() + comfyui_ver, form_factor = get_version_and_form_factor() + json_obj = await fetch_all(comfyui_ver, form_factor) manager_util.save_to_cache(uri, json_obj) + memory_cache[uri] = json_obj return json_obj['nodes'] - except: - res = {} - print("Cannot connect to comfyregistry.") - finally: - if cache_mode: - is_cache_loading = False - - return res - + except Exception as e: + print(f"Cannot connect to ComfyRegistry: {e}") + return {} @dataclass class NodeVersion: @@ -130,7 +167,6 @@ class NodeVersion: version: str download_url: str - def map_node_version(api_node_version): """ Maps node version data from API response to NodeVersion dataclass. @@ -162,7 +198,6 @@ def map_node_version(api_node_version): ), # Provide a default value if 'downloadUrl' is missing ) - def install_node(node_id, version=None): """ Retrieves the node version for installation. @@ -174,30 +209,59 @@ def install_node(node_id, version=None): Returns: NodeVersion: Node version data or error message. """ + cache_key = f"install_{node_id}_{version or 'latest'}" + if cache_key in memory_cache: + return memory_cache[cache_key] + if version is None: url = f"{base_url}/nodes/{node_id}/install" else: url = f"{base_url}/nodes/{node_id}/install?version={version}" - response = requests.get(url) + response = requests.get(url, timeout=10) if response.status_code == 200: # Convert the API response to a NodeVersion object - return map_node_version(response.json()) - else: - return None - + data = response.json() + result = map_node_version(data) + memory_cache[cache_key] = result + return result + return None def all_versions_of_node(node_id): + """ + Fetch all versions of a node. + + Args: + node_id (str): Node identifier. + + Returns: + list: List of version data or None. + """ + cache_key = f"versions_{node_id}" + if cache_key in memory_cache: + return memory_cache[cache_key] + url = f"{base_url}/nodes/{node_id}/versions?statuses=NodeVersionStatusActive&statuses=NodeVersionStatusPending" - - response = requests.get(url) + response = requests.get(url, timeout=10) if response.status_code == 200: - return response.json() - else: - return None - + data = response.json() + memory_cache[cache_key] = data + return data + return None def read_cnr_info(fullpath): + """ + Read CNR info from pyproject.toml and .tracking files. + + Args: + fullpath (str): Path to node directory. + + Returns: + dict: Node info or None if invalid. + """ + if fullpath in cnr_info_cache: + return cnr_info_cache[fullpath] + try: toml_path = os.path.join(fullpath, 'pyproject.toml') tracking_path = os.path.join(fullpath, '.tracking') @@ -219,18 +283,26 @@ def read_cnr_info(fullpath): repository = urls.get('Repository') if name and version: # repository is optional - return { + result = { "id": name, "version": version, "url": repository } + cnr_info_cache[fullpath] = result + return result return None except Exception: return None # not valid CNR node pack - def generate_cnr_id(fullpath, cnr_id): + """ + Generate .cnr-id file in the node's .git directory. + + Args: + fullpath (str): Path to node directory. + cnr_id (str): CNR ID to write. + """ cnr_id_path = os.path.join(fullpath, '.git', '.cnr-id') try: if not os.path.exists(cnr_id_path): @@ -239,8 +311,16 @@ def generate_cnr_id(fullpath, cnr_id): except: print(f"[ComfyUI Manager] unable to create file: {cnr_id_path}") - def read_cnr_id(fullpath): + """ + Read CNR ID from .cnr-id file in the node's .git directory. + + Args: + fullpath (str): Path to node directory. + + Returns: + str: CNR ID or None if not found. + """ cnr_id_path = os.path.join(fullpath, '.git', '.cnr-id') try: if os.path.exists(cnr_id_path): @@ -249,5 +329,4 @@ def read_cnr_id(fullpath): except: pass - return None - + return None \ No newline at end of file