build preview images better

workflows/preview_server/install.py

@@ -3,16 +3,27 @@ import subprocess
 import json
 import os
 from pathlib import Path
-import requests
+from urllib.parse import urlparse
 from multiprocessing import Pool
+import argparse
+
+parser = argparse.ArgumentParser(description='Install ComfyUI and its dependencies')
+parser.add_argument('--source', type=str, help='The source directory of the dependencies', required=False)
+parser.add_argument('--dependencies-json', type=str, help='The path to the json file containing the dependencies', required=False)
+parser.add_argument('--repos-json', type=str, help='The path to the json file containing the repositories', required=False)
+
+parser.add_argument('--concurrency', type=int, help='The number of concurrent processes to use', default=8)
+
+parser.add_argument('--comfy-ui-url', type=str, help='The url of the ComfyUI repository', required=False)
+parser.add_argument('--comfy-ui-commit', type=str, help='The commit of the ComfyUI repository', required=False)
 
 
+args = parser.parse_args()
 
 
 # set cwd to script location
 os.chdir(os.path.dirname(os.path.realpath(__file__)))
 
-json_file_path = 'repos.json'
 venv_python = os.path.join('ComfyUI', 'venv', 'Scripts' if os.name == 'nt' else 'bin', 'python')
 
 
@@ -33,7 +44,6 @@ def run_command(command):
             print(line, end='')
         print(f"Error executing command: {command}")
 
-
 def setup_repo(repo_info):
     url, commit, folder = repo_info['url'], repo_info['commit'], repo_info['folder']
     # save the current working directory
@@ -46,6 +56,15 @@ def setup_repo(repo_info):
         os.chdir(original_cwd)
 
 
+def setup_repos(json_file_path, concurrency=4):
+    with open(json_file_path, 'r') as file:
+        repos = json.load(file)
+    with Pool(concurrency) as pool:
+        pool.map(setup_repo, repos)
+
+def setup_comfy_ui(comfyui_info):
+    setup_repo(comfyui_info)
+
 def setup_venv_and_install_requirements():
     run_command("python3 -m venv ComfyUI/venv")
     print('=============# Installing ComfyUI Requirements #=============')
@@ -54,58 +73,55 @@ def setup_venv_and_install_requirements():
     run_command(f"{venv_python} -m pip install -r required_requirements.txt")
 
 
-def run_comfyui_setup():
+def initialize_comfyui():
     run_command(f"{venv_python} initialize_comfy_extensions.py")
 
-
-
-def install_dependency(dep):
+def install_dependency(dep, source_dir):
     location = dep['location']
     url = dep['url']
 
+    # separate the pathname from the url host
+    source_pathname = urlparse(url).path.lstrip('/')
+
     # Convert the location to a Path object and create parent directories if they don't exist
     path = Path(location)
     path.parent.mkdir(parents=True, exist_ok=True)
 
-    if path.is_file():
-        print(f"{path} is already installed")
-        return
+    source_path = Path() / source_dir /  source_pathname
+    # Copy the dependency from source dir to the location
+    print(f"Copying checkpoint from {source_path} to {location}")
+    run_command(f"cp {source_path} {location}")
 
-    # Download the file
-    response = requests.get(url)
-    if response.status_code == 200:
-        # Write the content to the file in the specified location
-        with open(path, 'wb') as file:
-            file.write(response.content)
-        print(f"Installed {path}")
-    else:
-        print(f"Failed to download {url}")
-
-
-
-def install_dependencies(json_file_path='dependencies.json'):
+def install_dependencies(json_file_path, source_dir, concurrency=4):
     # Load the dependency JSON
     with open(json_file_path, 'r') as file:
         dependencies = json.load(file)
 
-
-    with Pool(5) as pool:
-        pool.starmap(install_dependency, [(dep,) for dep in dependencies])
-
+    with Pool(concurrency) as pool:
+        pool.starmap(install_dependency, [(dep, source_dir) for dep in dependencies])
 
 def main():
-    with open(json_file_path, 'r') as file:
-        repos = json.load(file)
+    if args.comfy_ui_url and args.comfy_ui_commit:
+        comfyui_info = {
+            "url": args.comfy_ui_url,
+            "commit": args.comfy_ui_commit,
+            "folder": "ComfyUI",
+        }
+        setup_comfy_ui(comfyui_info)
 
-    for repo_info in repos:
-        print(f"installing {repo_info['url']} at commit {repo_info['commit']} to {repo_info['folder']}")
-        setup_repo(repo_info)
+    if args.repos_json:
+        setup_repos(args.repos_json, args.concurrency)
+    else:
+        print("Skipping repository setup as repos json is not provided")
 
-    setup_venv_and_install_requirements()
-    install_dependencies()
-    run_comfyui_setup()
+
+    #initialize_comfyui()
+    if args.source and args.dependencies_json:
+        install_dependencies(args.dependencies_json, args.source, args.concurrency)
+    else:
+        print("Skipping dependency installation as source and dependencies json are not provided")
 
 
 if __name__ == "__main__":
     main()
-    open('ALREADY_SETUP.flag', 'w').close()
+    open('ALREADY_SETUP.flag', 'w').close()