build preview images better

workflows/comfy/Dockerfile

@@ -52,7 +52,6 @@ RUN cd /tmp && \
         rm -rf /tmp/rclone*
 
 RUN mkdir -p /root/.config/pip
-
 COPY ./pip.conf /root/.config/pip/pip.conf
 
 WORKDIR /app

workflows/preview_server/Dockerfile

@@ -31,23 +31,63 @@ RUN apt-get update \
         python3.10 \
         python3.10-venv \
         unzip \
+        htop \
+        vim \
+        tmux \
+        rsync \
         wget \
         git \
         --no-install-recommends
 
 RUN addgroup --system appgroup && adduser --system --ingroup appgroup appuser
-# Copy files to container
-COPY --chown=appuser:appgroup initialize_comfy_extensions.py /app/initialize_comfy_extensions.py
-COPY --chown=appuser:appgroup install.py /app/install.py
-COPY --chown=appuser:appgroup repos.json /app/repos.json
-COPY --chown=appuser:appgroup dependencies.json /app/dependencies.json
-COPY --chown=appuser:appgroup server_requirements.txt /app/server_requirements.txt
-COPY --chown=appuser:appgroup required_requirements.txt /app/required_requirements.txt
-COPY --chown=appuser:appgroup preview_server.py /app/preview_server.py
 
 RUN pip3 install --no-cache-dir --upgrade pip setuptools wheel requests argparse
 
+RUN cd /tmp && \
+    curl -O https://downloads.rclone.org/rclone-current-linux-amd64.zip && \
+        unzip rclone-current-linux-amd64.zip && \
+        cd rclone-*-linux-amd64 && \
+        cp rclone /usr/bin/ && \
+        chmod 755 /usr/bin/rclone && \
+        rm -rf /tmp/rclone*
+
+RUN mkdir -p /root/.config/pip
+COPY ./pip.conf /root/.config/pip/pip.conf
+
+WORKDIR /app
+
+# Copy files to container
+COPY --chown=appuser:appgroup initialize_comfy_extensions.py /app/initialize_comfy_extensions.py
+COPY --chown=appuser:appgroup install.py /app/install.py
 RUN chmod -R 777 /app
 
+
+
+
+RUN python3 install.py --comfy-ui-url https://github.com/comfyanonymous/ComfyUI.git --comfy-ui-commit eda87043862f743b0a0467735f8531f7c4709b3a
+RUN python3 -m venv ComfyUI/venv
+RUN ComfyUI/venv/bin/python -m pip install -r ComfyUI/requirements.txt --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/ && \
+    ComfyUI/venv/bin/python -m pip cache purge
+
+COPY --chown=appuser:appgroup required_requirements.txt /app/required_requirements.txt
+RUN ComfyUI/venv/bin/python -m pip install -r required_requirements.txt --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/ && \
+    ComfyUI/venv/bin/python -m pip uninstall onnxruntime onnxruntime-gpu -y && \
+    ComfyUI/venv/bin/python -m pip install onnxruntime-gpu ort-nightly-gpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/ && \
+    ComfyUI/venv/bin/python -m pip cache purge
+
+COPY --chown=appuser:appgroup repos.json /app/repos.json
+RUN python3 install.py --repos-json repos.json
+
+COPY --chown=appuser:appgroup dependencies.json /app/dependencies.json
+COPY --chown=appuser:appgroup preview_server.py /app/preview_server.py
+COPY --chown=appuser:appgroup server_requirements.txt /app/server_requirements.txt
+ENV COMFY_ROOT_DIRECTORY=/app/ComfyUI
+ENV V2V_WORKFLOWS_DIRECTORY=/workflow_configs
+ENV PREVIEW_WORKFLOW_NAME=workflow_api_instantpreview.json
+
 # protobuf use python implementation (weird error without)
 ENV PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python
+
+# rclone copy --progress --transfers=40 vocodes-s3:kasisnu-test-area-private/comfyui-deps/ vocodes-s3:kasisnu-test-area-private/comfyui-preview-deps/
+
+# rclone copy --progress --transfers=40 . --include 'ComfyUI/models' vocodes-s3:kasisnu-test-area-private/comfyui-preview-deps/

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, source_dir):
-
-def install_dependency(dep):
     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():
+    source_path = Path() / source_dir /  source_pathname
-        print(f"{path} is already installed")
+    # Copy the dependency from source dir to the location
-        return
+    print(f"Copying checkpoint from {source_path} to {location}")
+    run_command(f"cp {source_path} {location}")
 
-    # Download the file
+def install_dependencies(json_file_path, source_dir, concurrency=4):
-    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'):
     # Load the dependency JSON
     with open(json_file_path, 'r') as file:
         dependencies = json.load(file)
 
-
+    with Pool(concurrency) as pool:
-    with Pool(5) as pool:
+        pool.starmap(install_dependency, [(dep, source_dir) for dep in dependencies])
-        pool.starmap(install_dependency, [(dep,) for dep in dependencies])
-
 
 def main():
-    with open(json_file_path, 'r') as file:
+    if args.comfy_ui_url and args.comfy_ui_commit:
-        repos = json.load(file)
+        comfyui_info = {
+            "url": args.comfy_ui_url,
+            "commit": args.comfy_ui_commit,
+            "folder": "ComfyUI",
+        }
+        setup_comfy_ui(comfyui_info)
 
-    for repo_info in repos:
+    if args.repos_json:
-        print(f"installing {repo_info['url']} at commit {repo_info['commit']} to {repo_info['folder']}")
+        setup_repos(args.repos_json, args.concurrency)
-        setup_repo(repo_info)
+    else:
+        print("Skipping repository setup as repos json is not provided")
 
-    setup_venv_and_install_requirements()
+
-    install_dependencies()
+    #initialize_comfyui()
-    run_comfyui_setup()
+    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()

workflows/preview_server/repos.json

@@ -1,9 +1,4 @@
 [
-    {
-        "url": "https://github.com/sALTaccount/ComfyUI.git",
-        "commit": "3f56e649595967624682dcbe3bded975f3a0dae8",
-        "folder": "ComfyUI"
-    },
     {
         "url": "https://github.com/ltdrdata/ComfyUI-Manager.git",
         "commit": "6c553d1228978352767f4914c4dd7787b00f6769",

workflows/preview_server/test_request.py

@@ -21,7 +21,7 @@ style_to_filename = {
     "ink_bw_style": "8_ink_w_and_b_style.json",
     "jojo_style": "9_jojo_style.json",
     "paper_origami": "10_paper_origami.json",
-    # "pixel_art": "11_pixel_art.json",
+    "pixel_art": "11_pixel_art.json",
     "pop_art": "12_pop_art.json",
     "realistic_1": "13_realistic_1.json",
     "realistic_2": "14_realistic_2.json",
@@ -46,12 +46,13 @@ for style in all_styles:
         with open(TEST_INPUT_FILE , 'rb') as input_file:
             # files = {'input_file': ('input_file', input_file, 'image/jpeg')}
             files = {'input_file':  input_file}
+            # write output file suffixed with timestamp
+            timestamp = time.strftime('%Y%m%d%H%M%S')
+            output_file_name = f'outputs/output_image_{style}_{timestamp}.jpeg'
             execution_start_time = time.perf_counter()
             with requests.post(url, data={'request': json.dumps(payload)}, files=files) as response:
                 if response.status_code == 200:
-                    # write output file suffixed with timestamp
+                    with open(output_file_name, 'wb') as f:
-                    timestamp = time.strftime('%Y%m%d%H%M%S')
-                    with open(f'outputs/preview_image_{style}_{timestamp}.jpeg', 'wb') as f:
                         f.write(response.content)
                 else:
                     # Handle request error
@@ -59,4 +60,4 @@ for style in all_styles:
                     print(response.text)
             current_time = time.perf_counter()
             execution_time = current_time - execution_start_time
-            print('{} executed in {:.2f} seconds'.format(style, execution_time))
+            print('{} generated {} in {:.2f} seconds'.format(style, output_file_name, execution_time))