fred-dev/comfy_ui_ali
0
1import ast2import re3import os4import json5from git import Repo6import concurrent7import datetime8import concurrent.futures9import requests10 11builtin_nodes = set()12 13import sys14 15from urllib.parse import urlparse16from github import Github17 18 19def download_url(url, dest_folder, filename=None):20 # Ensure the destination folder exists21 if not os.path.exists(dest_folder):22 os.makedirs(dest_folder)23 24 # Extract filename from URL if not provided25 if filename is None:26 filename = os.path.basename(url)27 28 # Full path to save the file29 dest_path = os.path.join(dest_folder, filename)30 31 # Download the file32 response = requests.get(url, stream=True)33 if response.status_code == 200:34 with open(dest_path, 'wb') as file:35 for chunk in response.iter_content(chunk_size=1024):36 if chunk:37 file.write(chunk)38 else:39 raise Exception(f"Failed to download file from {url}")40 41 42# prepare temp dir43if len(sys.argv) > 1:44 temp_dir = sys.argv[1]45else:46 temp_dir = os.path.join(os.getcwd(), ".tmp")47 48if not os.path.exists(temp_dir):49 os.makedirs(temp_dir)50 51 52skip_update = '--skip-update' in sys.argv or '--skip-all' in sys.argv53skip_stat_update = '--skip-stat-update' in sys.argv or '--skip-all' in sys.argv54 55if not skip_stat_update:56 g = Github(os.environ.get('GITHUB_TOKEN'))57else:58 g = None59 60 61print(f"TEMP DIR: {temp_dir}")62 63 64parse_cnt = 065 66 67def extract_nodes(code_text):68 global parse_cnt69 70 try:71 if parse_cnt % 100 == 0:72 print(".", end="", flush=True)73 parse_cnt += 174 75 code_text = re.sub(r'\\[^"\']', '', code_text)76 parsed_code = ast.parse(code_text)77 78 assignments = (node for node in parsed_code.body if isinstance(node, ast.Assign))79 80 for assignment in assignments:81 if isinstance(assignment.targets[0], ast.Name) and assignment.targets[0].id in ['NODE_CONFIG', 'NODE_CLASS_MAPPINGS']:82 node_class_mappings = assignment.value83 break84 else:85 node_class_mappings = None86 87 if node_class_mappings:88 s = set()89 90 for key in node_class_mappings.keys:91 if key is not None and isinstance(key.value, str):92 s.add(key.value.strip())93 94 return s95 else:96 return set()97 except:98 return set()99 100 101# scan102def scan_in_file(filename, is_builtin=False):103 global builtin_nodes104 105 try:106 with open(filename, encoding='utf-8') as file:107 code = file.read()108 except UnicodeDecodeError:109 with open(filename, encoding='cp949') as file:110 code = file.read()111 112 pattern = r"_CLASS_MAPPINGS\s*=\s*{([^}]*)}"113 regex = re.compile(pattern, re.MULTILINE | re.DOTALL)114 115 nodes = set()116 class_dict = {}117 118 nodes |= extract_nodes(code)119 code = re.sub(r'^#.*?$', '', code, flags=re.MULTILINE)120 121 def extract_keys(pattern, code):122 keys = re.findall(pattern, code)123 return {key.strip() for key in keys}124 125 def update_nodes(nodes, new_keys):126 nodes |= new_keys127 128 patterns = [129 r'^[^=]*_CLASS_MAPPINGS\["(.*?)"\]',130 r'^[^=]*_CLASS_MAPPINGS\[\'(.*?)\'\]',131 r'@register_node\("(.+)",\s*\".+"\)',132 r'"(\w+)"\s*:\s*{"class":\s*\w+\s*'133 ]134 135 with concurrent.futures.ThreadPoolExecutor() as executor:136 futures = {executor.submit(extract_keys, pattern, code): pattern for pattern in patterns}137 for future in concurrent.futures.as_completed(futures):138 update_nodes(nodes, future.result())139 140 matches = regex.findall(code)141 for match in matches:142 dict_text = match143 144 key_value_pairs = re.findall(r"\"([^\"]*)\"\s*:\s*([^,\n]*)", dict_text)145 for key, value in key_value_pairs:146 class_dict[key.strip()] = value.strip()147 148 key_value_pairs = re.findall(r"'([^']*)'\s*:\s*([^,\n]*)", dict_text)149 for key, value in key_value_pairs:150 class_dict[key.strip()] = value.strip()151 152 for key, value in class_dict.items():153 nodes.add(key.strip())154 155 update_pattern = r"_CLASS_MAPPINGS.update\s*\({([^}]*)}\)"156 update_match = re.search(update_pattern, code)157 if update_match:158 update_dict_text = update_match.group(1)159 update_key_value_pairs = re.findall(r"\"([^\"]*)\"\s*:\s*([^,\n]*)", update_dict_text)160 for key, value in update_key_value_pairs:161 class_dict[key.strip()] = value.strip()162 nodes.add(key.strip())163 164 metadata = {}165 lines = code.strip().split('\n')166 for line in lines:167 if line.startswith('@'):168 if line.startswith("@author:") or line.startswith("@title:") or line.startswith("@nickname:") or line.startswith("@description:"):169 key, value = line[1:].strip().split(':', 1)170 metadata[key.strip()] = value.strip()171 172 if is_builtin:173 builtin_nodes += set(nodes)174 else:175 for x in builtin_nodes:176 if x in nodes:177 nodes.remove(x)178 179 return nodes, metadata180 181 182def get_py_file_paths(dirname):183 file_paths = []184 185 for root, dirs, files in os.walk(dirname):186 if ".git" in root or "__pycache__" in root:187 continue188 189 for file in files:190 if file.endswith(".py"):191 file_path = os.path.join(root, file)192 file_paths.append(file_path)193 194 return file_paths195 196 197def get_nodes(target_dir):198 py_files = []199 directories = []200 201 for item in os.listdir(target_dir):202 if ".git" in item or "__pycache__" in item:203 continue204 205 path = os.path.abspath(os.path.join(target_dir, item))206 207 if os.path.isfile(path) and item.endswith(".py"):208 py_files.append(path)209 elif os.path.isdir(path):210 directories.append(path)211 212 return py_files, directories213 214 215def get_git_urls_from_json(json_file):216 with open(json_file, encoding='utf-8') as file:217 data = json.load(file)218 219 custom_nodes = data.get('custom_nodes', [])220 git_clone_files = []221 for node in custom_nodes:222 if node.get('install_type') == 'git-clone':223 files = node.get('files', [])224 if files:225 git_clone_files.append((files[0], node.get('title'), node.get('preemptions'), node.get('nodename_pattern')))226 227 git_clone_files.append(("https://github.com/comfyanonymous/ComfyUI", "ComfyUI", None, None))228 229 return git_clone_files230 231 232def get_py_urls_from_json(json_file):233 with open(json_file, encoding='utf-8') as file:234 data = json.load(file)235 236 custom_nodes = data.get('custom_nodes', [])237 py_files = []238 for node in custom_nodes:239 if node.get('install_type') == 'copy':240 files = node.get('files', [])241 if files:242 py_files.append((files[0], node.get('title'), node.get('preemptions'), node.get('nodename_pattern')))243 244 return py_files245 246 247def clone_or_pull_git_repository(git_url):248 repo_name = git_url.split("/")[-1]249 if repo_name.endswith(".git"):250 repo_name = repo_name[:-4]251 252 repo_dir = os.path.join(temp_dir, repo_name)253 254 if os.path.exists(repo_dir):255 try:256 repo = Repo(repo_dir)257 origin = repo.remote(name="origin")258 origin.pull()259 repo.git.submodule('update', '--init', '--recursive')260 print(f"Pulling {repo_name}...")261 except Exception as e:262 print(f"Pulling {repo_name} failed: {e}")263 else:264 try:265 Repo.clone_from(git_url, repo_dir, recursive=True)266 print(f"Cloning {repo_name}...")267 except Exception as e:268 print(f"Cloning {repo_name} failed: {e}")269 270 271def update_custom_nodes():272 if not os.path.exists(temp_dir):273 os.makedirs(temp_dir)274 275 node_info = {}276 277 git_url_titles_preemptions = get_git_urls_from_json('custom-node-list.json')278 279 def process_git_url_title(url, title, preemptions, node_pattern):280 name = os.path.basename(url)281 if name.endswith(".git"):282 name = name[:-4]283 284 node_info[name] = (url, title, preemptions, node_pattern)285 if not skip_update:286 clone_or_pull_git_repository(url)287 288 def process_git_stats(git_url_titles_preemptions):289 GITHUB_STATS_CACHE_FILENAME = 'github-stats-cache.json'290 GITHUB_STATS_FILENAME = 'github-stats.json'291 292 github_stats = {}293 try:294 with open(GITHUB_STATS_CACHE_FILENAME, 'r', encoding='utf-8') as file:295 github_stats = json.load(file)296 except FileNotFoundError:297 pass298 299 def is_rate_limit_exceeded():300 return g.rate_limiting[0] == 0301 302 if is_rate_limit_exceeded():303 print(f"GitHub API Rate Limit Exceeded: remained - {(g.rate_limiting_resettime - datetime.datetime.now().timestamp())/60:.2f} min")304 else:305 def renew_stat(url):306 if is_rate_limit_exceeded():307 return308 309 if 'github.com' not in url:310 return None311 312 print('.', end="")313 sys.stdout.flush()314 try:315 # Parsing the URL316 parsed_url = urlparse(url)317 domain = parsed_url.netloc318 path = parsed_url.path319 path_parts = path.strip("/").split("/")320 if len(path_parts) >= 2 and domain == "github.com":321 owner_repo = "/".join(path_parts[-2:])322 repo = g.get_repo(owner_repo)323 owner = repo.owner324 now = datetime.datetime.now(datetime.timezone.utc)325 author_time_diff = now - owner.created_at326 327 last_update = repo.pushed_at.strftime("%Y-%m-%d %H:%M:%S") if repo.pushed_at else 'N/A'328 item = {329 "stars": repo.stargazers_count,330 "last_update": last_update,331 "cached_time": now.timestamp(),332 "author_account_age_days": author_time_diff.days,333 }334 return url, item335 else:336 print(f"\nInvalid URL format for GitHub repository: {url}\n")337 except Exception as e:338 print(f"\nERROR on {url}\n{e}")339 340 return None341 342 # resolve unresolved urls343 with concurrent.futures.ThreadPoolExecutor(11) as executor:344 futures = []345 for url, title, preemptions, node_pattern in git_url_titles_preemptions:346 if url not in github_stats:347 futures.append(executor.submit(renew_stat, url))348 349 for future in concurrent.futures.as_completed(futures):350 url_item = future.result()351 if url_item is not None:352 url, item = url_item353 github_stats[url] = item354 355 # renew outdated cache356 outdated_urls = []357 for k, v in github_stats.items():358 elapsed = (datetime.datetime.now().timestamp() - v['cached_time'])359 if elapsed > 60*60*12: # 12 hours360 outdated_urls.append(k)361 362 with concurrent.futures.ThreadPoolExecutor(11) as executor:363 for url in outdated_urls:364 futures.append(executor.submit(renew_stat, url))365 366 for future in concurrent.futures.as_completed(futures):367 url_item = future.result()368 if url_item is not None:369 url, item = url_item370 github_stats[url] = item371 372 with open('github-stats-cache.json', 'w', encoding='utf-8') as file:373 json.dump(github_stats, file, ensure_ascii=False, indent=4)374 375 with open(GITHUB_STATS_FILENAME, 'w', encoding='utf-8') as file:376 for v in github_stats.values():377 if "cached_time" in v:378 del v["cached_time"]379 380 github_stats = dict(sorted(github_stats.items()))381 382 json.dump(github_stats, file, ensure_ascii=False, indent=4)383 384 print(f"Successfully written to {GITHUB_STATS_FILENAME}.")385 386 if not skip_stat_update:387 process_git_stats(git_url_titles_preemptions)388 389 with concurrent.futures.ThreadPoolExecutor(11) as executor:390 for url, title, preemptions, node_pattern in git_url_titles_preemptions:391 executor.submit(process_git_url_title, url, title, preemptions, node_pattern)392 393 py_url_titles_and_pattern = get_py_urls_from_json('custom-node-list.json')394 395 def download_and_store_info(url_title_preemptions_and_pattern):396 url, title, preemptions, node_pattern = url_title_preemptions_and_pattern397 name = os.path.basename(url)398 if name.endswith(".py"):399 node_info[name] = (url, title, preemptions, node_pattern)400 401 try:402 download_url(url, temp_dir)403 except:404 print(f"[ERROR] Cannot download '{url}'")405 406 with concurrent.futures.ThreadPoolExecutor(10) as executor:407 executor.map(download_and_store_info, py_url_titles_and_pattern)408 409 return node_info410 411 412def gen_json(node_info):413 # scan from .py file414 node_files, node_dirs = get_nodes(temp_dir)415 416 comfyui_path = os.path.abspath(os.path.join(temp_dir, "ComfyUI"))417 node_dirs.remove(comfyui_path)418 node_dirs = [comfyui_path] + node_dirs419 420 data = {}421 for dirname in node_dirs:422 py_files = get_py_file_paths(dirname)423 metadata = {}424 425 nodes = set()426 for py in py_files:427 nodes_in_file, metadata_in_file = scan_in_file(py, dirname == "ComfyUI")428 nodes.update(nodes_in_file)429 metadata.update(metadata_in_file)430 431 dirname = os.path.basename(dirname)432 433 if 'Jovimetrix' in dirname:434 pass435 436 if len(nodes) > 0 or (dirname in node_info and node_info[dirname][3] is not None):437 nodes = list(nodes)438 nodes.sort()439 440 if dirname in node_info:441 git_url, title, preemptions, node_pattern = node_info[dirname]442 443 metadata['title_aux'] = title444 445 if preemptions is not None:446 metadata['preemptions'] = preemptions447 448 if node_pattern is not None:449 metadata['nodename_pattern'] = node_pattern450 451 data[git_url] = (nodes, metadata)452 else:453 print(f"WARN: {dirname} is removed from custom-node-list.json")454 455 for file in node_files:456 nodes, metadata = scan_in_file(file)457 458 if len(nodes) > 0 or (dirname in node_info and node_info[dirname][3] is not None):459 nodes = list(nodes)460 nodes.sort()461 462 file = os.path.basename(file)463 464 if file in node_info:465 url, title, preemptions, node_pattern = node_info[file]466 metadata['title_aux'] = title467 468 if preemptions is not None:469 metadata['preemptions'] = preemptions470 471 if node_pattern is not None:472 metadata['nodename_pattern'] = node_pattern473 474 data[url] = (nodes, metadata)475 else:476 print(f"Missing info: {file}")477 478 # scan from node_list.json file479 extensions = [name for name in os.listdir(temp_dir) if os.path.isdir(os.path.join(temp_dir, name))]480 481 for extension in extensions:482 node_list_json_path = os.path.join(temp_dir, extension, 'node_list.json')483 if os.path.exists(node_list_json_path):484 git_url, title, preemptions, node_pattern = node_info[extension]485 486 with open(node_list_json_path, 'r', encoding='utf-8') as f:487 try:488 node_list_json = json.load(f)489 except Exception as e:490 print(f"\nERROR: Invalid json format '{node_list_json_path}'")491 print("------------------------------------------------------")492 print(e)493 print("------------------------------------------------------")494 node_list_json = {}495 496 metadata_in_url = {}497 if git_url not in data:498 nodes = set()499 else:500 nodes_in_url, metadata_in_url = data[git_url]501 nodes = set(nodes_in_url)502 503 for x, desc in node_list_json.items():504 nodes.add(x.strip())505 506 metadata_in_url['title_aux'] = title507 508 if preemptions is not None:509 metadata['preemptions'] = preemptions510 511 if node_pattern is not None:512 metadata_in_url['nodename_pattern'] = node_pattern513 514 nodes = list(nodes)515 nodes.sort()516 data[git_url] = (nodes, metadata_in_url)517 518 json_path = "extension-node-map.json"519 with open(json_path, "w", encoding='utf-8') as file:520 json.dump(data, file, indent=4, sort_keys=True)521 522 523print("### ComfyUI Manager Node Scanner ###")524 525print("\n# Updating extensions\n")526updated_node_info = update_custom_nodes()527 528print("\n# 'extension-node-map.json' file is generated.\n")529gen_json(updated_node_info)530 531print("\nDONE.\n")