mirror of
https://github.com/Comfy-Org/ComfyUI-Manager.git
synced 2026-02-15 23:52:32 +08:00
speed up caching and loading
This commit is contained in:
parent
41006c3a33
commit
7362816ad2
@ -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
|
||||
Loading…
Reference in New Issue
Block a user