| import re |
| import os |
| import json |
| from git import Repo |
| from torchvision.datasets.utils import download_url |
| import concurrent |
|
|
| builtin_nodes = set() |
|
|
|
|
| def scan_in_file(filename, is_builtin=False): |
| global builtin_nodes |
|
|
| try: |
| with open(filename, encoding='utf-8') as file: |
| code = file.read() |
| except UnicodeDecodeError: |
| with open(filename, encoding='cp949') as file: |
| code = file.read() |
|
|
| pattern = r"_CLASS_MAPPINGS\s*=\s*{([^}]*)}" |
| regex = re.compile(pattern, re.MULTILINE | re.DOTALL) |
|
|
| nodes = set() |
| class_dict = {} |
|
|
| pattern2 = r'^[^=]*_CLASS_MAPPINGS\["(.*?)"\]' |
| keys = re.findall(pattern2, code) |
| for key in keys: |
| nodes.add(key.strip()) |
|
|
| pattern3 = r'^[^=]*_CLASS_MAPPINGS\[\'(.*?)\'\]' |
| keys = re.findall(pattern3, code) |
| for key in keys: |
| nodes.add(key.strip()) |
|
|
| matches = regex.findall(code) |
| for match in matches: |
| dict_text = match |
|
|
| key_value_pairs = re.findall(r"\"([^\"]*)\"\s*:\s*([^,\n]*)", dict_text) |
| for key, value in key_value_pairs: |
| class_dict[key.strip()] = value.strip() |
|
|
| key_value_pairs = re.findall(r"'([^']*)'\s*:\s*([^,\n]*)", dict_text) |
| for key, value in key_value_pairs: |
| class_dict[key.strip()] = value.strip() |
|
|
| for key, value in class_dict.items(): |
| nodes.add(key.strip()) |
|
|
| update_pattern = r"_CLASS_MAPPINGS.update\s*\({([^}]*)}\)" |
| update_match = re.search(update_pattern, code) |
| if update_match: |
| update_dict_text = update_match.group(1) |
| update_key_value_pairs = re.findall(r"\"([^\"]*)\"\s*:\s*([^,\n]*)", update_dict_text) |
| for key, value in update_key_value_pairs: |
| class_dict[key.strip()] = value.strip() |
| nodes.add(key.strip()) |
|
|
| metadata = {} |
| lines = code.strip().split('\n') |
| for line in lines: |
| if line.startswith('@'): |
| if line.startswith("@author:") or line.startswith("@title:") or line.startswith("@nickname:") or line.startswith("@description:"): |
| key, value = line[1:].strip().split(':') |
| metadata[key.strip()] = value.strip() |
|
|
| if is_builtin: |
| builtin_nodes += set(nodes) |
| else: |
| for x in builtin_nodes: |
| if x in nodes: |
| nodes.remove(x) |
|
|
| return nodes, metadata |
|
|
|
|
| def get_py_file_paths(dirname): |
| file_paths = [] |
| |
| for root, dirs, files in os.walk(dirname): |
| if ".git" in root or "__pycache__" in root: |
| continue |
|
|
| for file in files: |
| if file.endswith(".py"): |
| file_path = os.path.join(root, file) |
| file_paths.append(file_path) |
| |
| return file_paths |
|
|
|
|
| def get_nodes(target_dir): |
| py_files = [] |
| directories = [] |
| |
| for item in os.listdir(target_dir): |
| if ".git" in item or "__pycache__" in item: |
| continue |
|
|
| path = os.path.abspath(os.path.join(target_dir, item)) |
| |
| if os.path.isfile(path) and item.endswith(".py"): |
| py_files.append(path) |
| elif os.path.isdir(path): |
| directories.append(path) |
| |
| return py_files, directories |
|
|
|
|
| def get_git_urls_from_json(json_file): |
| with open(json_file, encoding='utf-8') as file: |
| data = json.load(file) |
|
|
| custom_nodes = data.get('custom_nodes', []) |
| git_clone_files = [] |
| for node in custom_nodes: |
| if node.get('install_type') == 'git-clone': |
| files = node.get('files', []) |
| if files: |
| git_clone_files.append((files[0], node.get('title'), node.get('nodename_pattern'))) |
|
|
| git_clone_files.append(("https://github.com/comfyanonymous/ComfyUI", "ComfyUI", None)) |
|
|
| return git_clone_files |
|
|
|
|
| def get_py_urls_from_json(json_file): |
| with open(json_file, encoding='utf-8') as file: |
| data = json.load(file) |
|
|
| custom_nodes = data.get('custom_nodes', []) |
| py_files = [] |
| for node in custom_nodes: |
| if node.get('install_type') == 'copy': |
| files = node.get('files', []) |
| if files: |
| py_files.append((files[0], node.get('title'), node.get('nodename_pattern'))) |
|
|
| return py_files |
|
|
|
|
| def clone_or_pull_git_repository(git_url): |
| repo_name = git_url.split("/")[-1].split(".")[0] |
| repo_dir = os.path.join(os.getcwd(), ".tmp", repo_name) |
|
|
| if os.path.exists(repo_dir): |
| try: |
| repo = Repo(repo_dir) |
| origin = repo.remote(name="origin") |
| origin.pull(rebase=True) |
| repo.git.submodule('update', '--init', '--recursive') |
| print(f"Pulling {repo_name}...") |
| except Exception as e: |
| print(f"Pulling {repo_name} failed: {e}") |
| else: |
| try: |
| Repo.clone_from(git_url, repo_dir, recursive=True) |
| print(f"Cloning {repo_name}...") |
| except Exception as e: |
| print(f"Cloning {repo_name} failed: {e}") |
|
|
|
|
| def update_custom_nodes(): |
| tmp_dir = os.path.join(os.getcwd(), ".tmp") |
| if not os.path.exists(tmp_dir): |
| os.makedirs(tmp_dir) |
|
|
| node_info = {} |
|
|
| git_url_titles = get_git_urls_from_json('custom-node-list.json') |
|
|
| def process_git_url_title(url, title, node_pattern): |
| name = os.path.basename(url) |
| if name.endswith(".git"): |
| name = name[:-4] |
| |
| node_info[name] = (url, title, node_pattern) |
| clone_or_pull_git_repository(url) |
|
|
| with concurrent.futures.ThreadPoolExecutor(10) as executor: |
| for url, title, node_pattern in git_url_titles: |
| executor.submit(process_git_url_title, url, title, node_pattern) |
|
|
| py_url_titles_and_pattern = get_py_urls_from_json('custom-node-list.json') |
|
|
| def download_and_store_info(url_title_and_pattern): |
| url, title, node_pattern = url_title_and_pattern |
| name = os.path.basename(url) |
| if name.endswith(".py"): |
| node_info[name] = (url, title, node_pattern) |
|
|
| try: |
| download_url(url, ".tmp") |
| except: |
| print(f"[ERROR] Cannot download '{url}'") |
|
|
| with concurrent.futures.ThreadPoolExecutor(10) as executor: |
| executor.map(download_and_store_info, py_url_titles_and_pattern) |
|
|
| return node_info |
|
|
|
|
| def gen_json(node_info): |
| |
| node_files, node_dirs = get_nodes(".tmp") |
|
|
| comfyui_path = os.path.abspath(os.path.join('.tmp', "ComfyUI")) |
| node_dirs.remove(comfyui_path) |
| node_dirs = [comfyui_path] + node_dirs |
|
|
| data = {} |
| for dirname in node_dirs: |
| py_files = get_py_file_paths(dirname) |
| metadata = {} |
| |
| nodes = set() |
| for py in py_files: |
| nodes_in_file, metadata_in_file = scan_in_file(py, dirname == "ComfyUI") |
| nodes.update(nodes_in_file) |
| metadata.update(metadata_in_file) |
| |
| dirname = os.path.basename(dirname) |
|
|
| if len(nodes) > 0 or (dirname in node_info and node_info[dirname][2] is not None): |
| nodes = list(nodes) |
| nodes.sort() |
|
|
| if dirname in node_info: |
| git_url, title, node_pattern = node_info[dirname] |
| metadata['title_aux'] = title |
| if node_pattern is not None: |
| metadata['nodename_pattern'] = node_pattern |
| data[git_url] = (nodes, metadata) |
| else: |
| print(f"WARN: {dirname} is removed from custom-node-list.json") |
|
|
| for file in node_files: |
| nodes, metadata = scan_in_file(file) |
|
|
| if len(nodes) > 0 or (dirname in node_info and node_info[dirname][2] is not None): |
| nodes = list(nodes) |
| nodes.sort() |
|
|
| file = os.path.basename(file) |
|
|
| if file in node_info: |
| url, title, node_pattern = node_info[file] |
| metadata['title_aux'] = title |
| if node_pattern is not None: |
| metadata['nodename_pattern'] = node_pattern |
| data[url] = (nodes, metadata) |
| else: |
| print(f"Missing info: {url}") |
|
|
| |
| extensions = [name for name in os.listdir('.tmp') if os.path.isdir(os.path.join('.tmp', name))] |
|
|
| for extension in extensions: |
| node_list_json_path = os.path.join('.tmp', extension, 'node_list.json') |
| if os.path.exists(node_list_json_path): |
| git_url, title, node_pattern = node_info[extension] |
|
|
| with open(node_list_json_path, 'r', encoding='utf-8') as f: |
| node_list_json = json.load(f) |
|
|
| metadata_in_url = {} |
| if git_url not in data: |
| nodes = set() |
| else: |
| nodes_in_url, metadata_in_url = data[git_url] |
| nodes = set(nodes_in_url) |
|
|
| for x, desc in node_list_json.items(): |
| nodes.add(x.strip()) |
|
|
| metadata_in_url['title_aux'] = title |
| if node_pattern is not None: |
| metadata['nodename_pattern'] = node_pattern |
| nodes = list(nodes) |
| nodes.sort() |
| data[git_url] = (nodes, metadata_in_url) |
|
|
| json_path = f"extension-node-map.json" |
| with open(json_path, "w", encoding='utf-8') as file: |
| json.dump(data, file, indent=4, sort_keys=True) |
|
|
|
|
| print("### ComfyUI Manager Node Scanner ###") |
|
|
| print("\n# Updating extensions\n") |
| updated_node_info = update_custom_nodes() |
|
|
| print("\n# 'extension-node-map.json' file is generated.\n") |
| gen_json(updated_node_info) |
|
|
