fred-dev/comfy_ui_ali
0
1import subprocess2import sys3import os4import traceback5 6import git7import json8import yaml9import requests10from tqdm.auto import tqdm11from git.remote import RemoteProgress12 13 14comfy_path = os.environ.get('COMFYUI_PATH')15git_exe_path = os.environ.get('GIT_EXE_PATH')16 17if comfy_path is None:18 print("\nWARN: The `COMFYUI_PATH` environment variable is not set. Assuming `custom_nodes/ComfyUI-Manager/../../` as the ComfyUI path.", file=sys.stderr)19 comfy_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..'))20 21 22def download_url(url, dest_folder, filename=None):23 # Ensure the destination folder exists24 if not os.path.exists(dest_folder):25 os.makedirs(dest_folder)26 27 # Extract filename from URL if not provided28 if filename is None:29 filename = os.path.basename(url)30 31 # Full path to save the file32 dest_path = os.path.join(dest_folder, filename)33 34 # Download the file35 response = requests.get(url, stream=True)36 if response.status_code == 200:37 with open(dest_path, 'wb') as file:38 for chunk in response.iter_content(chunk_size=1024):39 if chunk:40 file.write(chunk)41 else:42 print(f"Failed to download file from {url}")43 44 45nodelist_path = os.path.join(os.path.dirname(__file__), "custom-node-list.json")46working_directory = os.getcwd()47 48if os.path.basename(working_directory) != 'custom_nodes':49 print("WARN: This script should be executed in custom_nodes dir")50 print(f"DBG: INFO {working_directory}")51 print(f"DBG: INFO {sys.argv}")52 # exit(-1)53 54 55class GitProgress(RemoteProgress):56 def __init__(self):57 super().__init__()58 self.pbar = tqdm(ascii=True)59 60 def update(self, op_code, cur_count, max_count=None, message=''):61 self.pbar.total = max_count62 self.pbar.n = cur_count63 self.pbar.pos = 064 self.pbar.refresh()65 66 67def gitclone(custom_nodes_path, url, target_hash=None, repo_path=None):68 repo_name = os.path.splitext(os.path.basename(url))[0]69 70 if repo_path is None:71 repo_path = os.path.join(custom_nodes_path, repo_name)72 73 # Clone the repository from the remote URL74 repo = git.Repo.clone_from(url, repo_path, recursive=True, progress=GitProgress())75 76 if target_hash is not None:77 print(f"CHECKOUT: {repo_name} [{target_hash}]")78 repo.git.checkout(target_hash)79 80 repo.git.clear_cache()81 repo.close()82 83 84def gitcheck(path, do_fetch=False):85 try:86 # Fetch the latest commits from the remote repository87 repo = git.Repo(path)88 89 if repo.head.is_detached:90 print("CUSTOM NODE CHECK: True")91 return92 93 current_branch = repo.active_branch94 branch_name = current_branch.name95 96 remote_name = current_branch.tracking_branch().remote_name97 remote = repo.remote(name=remote_name)98 99 if do_fetch:100 remote.fetch()101 102 # Get the current commit hash and the commit hash of the remote branch103 commit_hash = repo.head.commit.hexsha104 105 if f'{remote_name}/{branch_name}' in repo.refs:106 remote_commit_hash = repo.refs[f'{remote_name}/{branch_name}'].object.hexsha107 else:108 print("CUSTOM NODE CHECK: True") # non default branch is treated as updatable109 return110 111 # Compare the commit hashes to determine if the local repository is behind the remote repository112 if commit_hash != remote_commit_hash:113 # Get the commit dates114 commit_date = repo.head.commit.committed_datetime115 remote_commit_date = repo.refs[f'{remote_name}/{branch_name}'].object.committed_datetime116 117 # Compare the commit dates to determine if the local repository is behind the remote repository118 if commit_date < remote_commit_date:119 print("CUSTOM NODE CHECK: True")120 else:121 print("CUSTOM NODE CHECK: False")122 except Exception as e:123 print(e)124 print("CUSTOM NODE CHECK: Error")125 126 127def get_remote_name(repo):128 available_remotes = [remote.name for remote in repo.remotes]129 if 'origin' in available_remotes:130 return 'origin'131 elif 'upstream' in available_remotes:132 return 'upstream'133 elif len(available_remotes) > 0:134 return available_remotes[0]135 136 if not available_remotes:137 print(f"[ComfyUI-Manager] No remotes are configured for this repository: {repo.working_dir}")138 else:139 print(f"[ComfyUI-Manager] Available remotes in '{repo.working_dir}': ")140 for remote in available_remotes:141 print(f"- {remote}")142 143 return None144 145 146def switch_to_default_branch(repo):147 remote_name = get_remote_name(repo)148 149 try:150 if remote_name is None:151 return False152 153 default_branch = repo.git.symbolic_ref(f'refs/remotes/{remote_name}/HEAD').replace(f'refs/remotes/{remote_name}/', '')154 repo.git.checkout(default_branch)155 return True156 except:157 # try checkout master158 # try checkout main if failed159 try:160 repo.git.checkout(repo.heads.master)161 return True162 except:163 try:164 if remote_name is not None:165 repo.git.checkout('-b', 'master', f'{remote_name}/master')166 return True167 except:168 try:169 repo.git.checkout(repo.heads.main)170 return True171 except:172 try:173 if remote_name is not None:174 repo.git.checkout('-b', 'main', f'{remote_name}/main')175 return True176 except:177 pass178 179 print("[ComfyUI Manager] Failed to switch to the default branch")180 return False181 182 183def gitpull(path):184 # Check if the path is a git repository185 if not os.path.exists(os.path.join(path, '.git')):186 raise ValueError('Not a git repository')187 188 # Pull the latest changes from the remote repository189 repo = git.Repo(path)190 if repo.is_dirty():191 print(f"STASH: '{path}' is dirty.")192 repo.git.stash()193 194 commit_hash = repo.head.commit.hexsha195 try:196 if repo.head.is_detached:197 switch_to_default_branch(repo)198 199 current_branch = repo.active_branch200 branch_name = current_branch.name201 202 remote_name = current_branch.tracking_branch().remote_name203 remote = repo.remote(name=remote_name)204 205 if f'{remote_name}/{branch_name}' not in repo.refs:206 switch_to_default_branch(repo)207 current_branch = repo.active_branch208 branch_name = current_branch.name209 210 remote.fetch()211 if f'{remote_name}/{branch_name}' in repo.refs:212 remote_commit_hash = repo.refs[f'{remote_name}/{branch_name}'].object.hexsha213 else:214 print("CUSTOM NODE PULL: Fail") # update fail215 return216 217 if commit_hash == remote_commit_hash:218 print("CUSTOM NODE PULL: None") # there is no update219 repo.close()220 return221 222 remote.pull()223 224 repo.git.submodule('update', '--init', '--recursive')225 new_commit_hash = repo.head.commit.hexsha226 227 if commit_hash != new_commit_hash:228 print("CUSTOM NODE PULL: Success") # update success229 else:230 print("CUSTOM NODE PULL: Fail") # update fail231 except Exception as e:232 print(e)233 print("CUSTOM NODE PULL: Fail") # unknown git error234 235 repo.close()236 237 238def checkout_comfyui_hash(target_hash):239 repo = git.Repo(comfy_path)240 commit_hash = repo.head.commit.hexsha241 242 if commit_hash != target_hash:243 try:244 print(f"CHECKOUT: ComfyUI [{target_hash}]")245 repo.git.checkout(target_hash)246 except git.GitCommandError as e:247 print(f"Error checking out the ComfyUI: {str(e)}")248 249 250def checkout_custom_node_hash(git_custom_node_infos):251 repo_name_to_url = {}252 253 for url in git_custom_node_infos.keys():254 repo_name = url.split('/')[-1]255 256 if repo_name.endswith('.git'):257 repo_name = repo_name[:-4]258 259 repo_name_to_url[repo_name] = url260 261 for path in os.listdir(working_directory):262 if path.endswith("ComfyUI-Manager"):263 continue264 265 fullpath = os.path.join(working_directory, path)266 267 if os.path.isdir(fullpath):268 is_disabled = path.endswith(".disabled")269 270 try:271 git_dir = os.path.join(fullpath, '.git')272 if not os.path.exists(git_dir):273 continue274 275 need_checkout = False276 repo_name = os.path.basename(fullpath)277 278 if repo_name.endswith('.disabled'):279 repo_name = repo_name[:-9]280 281 if repo_name not in repo_name_to_url:282 if not is_disabled:283 # should be disabled284 print(f"DISABLE: {repo_name}")285 new_path = fullpath + ".disabled"286 os.rename(fullpath, new_path)287 need_checkout = False288 else:289 item = git_custom_node_infos[repo_name_to_url[repo_name]]290 if item['disabled'] and is_disabled:291 pass292 elif item['disabled'] and not is_disabled:293 # disable294 print(f"DISABLE: {repo_name}")295 new_path = fullpath + ".disabled"296 os.rename(fullpath, new_path)297 298 elif not item['disabled'] and is_disabled:299 # enable300 print(f"ENABLE: {repo_name}")301 new_path = fullpath[:-9]302 os.rename(fullpath, new_path)303 fullpath = new_path304 need_checkout = True305 else:306 need_checkout = True307 308 if need_checkout:309 repo = git.Repo(fullpath)310 commit_hash = repo.head.commit.hexsha311 312 if commit_hash != item['hash']:313 print(f"CHECKOUT: {repo_name} [{item['hash']}]")314 repo.git.checkout(item['hash'])315 316 except Exception:317 print(f"Failed to restore snapshots for the custom node '{path}'")318 319 # clone missing320 for k, v in git_custom_node_infos.items():321 if 'ComfyUI-Manager' in k:322 continue323 324 if not v['disabled']:325 repo_name = k.split('/')[-1]326 if repo_name.endswith('.git'):327 repo_name = repo_name[:-4]328 329 path = os.path.join(working_directory, repo_name)330 if not os.path.exists(path):331 print(f"CLONE: {path}")332 gitclone(working_directory, k, target_hash=v['hash'])333 334 335def invalidate_custom_node_file(file_custom_node_infos):336 global nodelist_path337 338 enabled_set = set()339 for item in file_custom_node_infos:340 if not item['disabled']:341 enabled_set.add(item['filename'])342 343 for path in os.listdir(working_directory):344 fullpath = os.path.join(working_directory, path)345 346 if not os.path.isdir(fullpath) and fullpath.endswith('.py'):347 if path not in enabled_set:348 print(f"DISABLE: {path}")349 new_path = fullpath+'.disabled'350 os.rename(fullpath, new_path)351 352 elif not os.path.isdir(fullpath) and fullpath.endswith('.py.disabled'):353 path = path[:-9]354 if path in enabled_set:355 print(f"ENABLE: {path}")356 new_path = fullpath[:-9]357 os.rename(fullpath, new_path)358 359 # download missing: just support for 'copy' style360 py_to_url = {}361 362 with open(nodelist_path, 'r', encoding="UTF-8") as json_file:363 info = json.load(json_file)364 for item in info['custom_nodes']:365 if item['install_type'] == 'copy':366 for url in item['files']:367 if url.endswith('.py'):368 py = url.split('/')[-1]369 py_to_url[py] = url370 371 for item in file_custom_node_infos:372 filename = item['filename']373 if not item['disabled']:374 target_path = os.path.join(working_directory, filename)375 376 if not os.path.exists(target_path) and filename in py_to_url:377 url = py_to_url[filename]378 print(f"DOWNLOAD: {filename}")379 download_url(url, working_directory)380 381 382def apply_snapshot(path):383 try:384 if os.path.exists(path):385 if not path.endswith('.json') and not path.endswith('.yaml'):386 print(f"Snapshot file not found: `{path}`")387 print("APPLY SNAPSHOT: False")388 return None389 390 with open(path, 'r', encoding="UTF-8") as snapshot_file:391 if path.endswith('.json'):392 info = json.load(snapshot_file)393 elif path.endswith('.yaml'):394 info = yaml.load(snapshot_file, Loader=yaml.SafeLoader)395 info = info['custom_nodes']396 else:397 # impossible case398 print("APPLY SNAPSHOT: False")399 return None400 401 comfyui_hash = info['comfyui']402 git_custom_node_infos = info['git_custom_nodes']403 file_custom_node_infos = info['file_custom_nodes']404 405 if comfyui_hash:406 checkout_comfyui_hash(comfyui_hash)407 checkout_custom_node_hash(git_custom_node_infos)408 invalidate_custom_node_file(file_custom_node_infos)409 410 print("APPLY SNAPSHOT: True")411 if 'pips' in info and info['pips']:412 return info['pips']413 else:414 return None415 416 print(f"Snapshot file not found: `{path}`")417 print("APPLY SNAPSHOT: False")418 419 return None420 except Exception as e:421 print(e)422 traceback.print_exc()423 print("APPLY SNAPSHOT: False")424 425 return None426 427 428def restore_pip_snapshot(pips, options):429 non_url = []430 local_url = []431 non_local_url = []432 for k, v in pips.items():433 if v == "":434 non_url.append(k)435 else:436 if v.startswith('file:'):437 local_url.append(v)438 else:439 non_local_url.append(v)440 441 failed = []442 if '--pip-non-url' in options:443 # try all at once444 res = 1445 try:446 res = subprocess.check_call([sys.executable, '-m', 'pip', 'install'] + non_url)447 except:448 pass449 450 # fallback451 if res != 0:452 for x in non_url:453 res = 1454 try:455 res = subprocess.check_call([sys.executable, '-m', 'pip', 'install', x])456 except:457 pass458 459 if res != 0:460 failed.append(x)461 462 if '--pip-non-local-url' in options:463 for x in non_local_url:464 res = 1465 try:466 res = subprocess.check_call([sys.executable, '-m', 'pip', 'install', x])467 except:468 pass469 470 if res != 0:471 failed.append(x)472 473 if '--pip-local-url' in options:474 for x in local_url:475 res = 1476 try:477 res = subprocess.check_call([sys.executable, '-m', 'pip', 'install', x])478 except:479 pass480 481 if res != 0:482 failed.append(x)483 484 print(f"Installation failed for pip packages: {failed}")485 486 487def setup_environment():488 if git_exe_path is not None:489 git.Git().update_environment(GIT_PYTHON_GIT_EXECUTABLE=git_exe_path)490 491 492setup_environment()493 494 495try:496 if sys.argv[1] == "--clone":497 repo_path = None498 if len(sys.argv) > 4:499 repo_path = sys.argv[4]500 501 gitclone(sys.argv[2], sys.argv[3], repo_path=repo_path)502 elif sys.argv[1] == "--check":503 gitcheck(sys.argv[2], False)504 elif sys.argv[1] == "--fetch":505 gitcheck(sys.argv[2], True)506 elif sys.argv[1] == "--pull":507 gitpull(sys.argv[2])508 elif sys.argv[1] == "--apply-snapshot":509 options = set()510 for x in sys.argv:511 if x in ['--pip-non-url', '--pip-local-url', '--pip-non-local-url']:512 options.add(x)513 514 pips = apply_snapshot(sys.argv[2])515 516 if pips and len(options) > 0:517 restore_pip_snapshot(pips, options)518 sys.exit(0)519except Exception as e:520 print(e)521 sys.exit(-1)522 523 524