From 5fc27f150d859fee48e8ccdc441a93ac90675eb1 Mon Sep 17 00:00:00 2001 From: Kasisnu <5044120+kasisnu@users.noreply.github.com> Date: Tue, 16 Jul 2024 17:31:44 +0530 Subject: [PATCH] comfy: add support for live previews --- .../ComfyLauncher/ComfyLivePortraitRunner.py | 254 ++++++++++++++++++ workflows/comfy/Dockerfile | 35 +-- workflows/comfy/repos.json | 5 + workflows/comfy/required_requirements.txt | 3 +- 4 files changed, 279 insertions(+), 18 deletions(-) create mode 100644 workflows/comfy/ComfyLauncher/ComfyLivePortraitRunner.py diff --git a/workflows/comfy/ComfyLauncher/ComfyLivePortraitRunner.py b/workflows/comfy/ComfyLauncher/ComfyLivePortraitRunner.py new file mode 100644 index 0000000..abc2f57 --- /dev/null +++ b/workflows/comfy/ComfyLauncher/ComfyLivePortraitRunner.py @@ -0,0 +1,254 @@ +import requests +import argparse +import os +import json +import time +from typing import Dict, Any, Optional +from jsonpath_ng import jsonpath, parse +import shutil +from pathlib import Path +from enum import IntEnum, auto +import os +import subprocess +import signal + +SERVER_IP = os.environ.get("SERVER_IP", "127.0.0.1:8188") +SERVER_START_TIMEOUT = int(os.environ.get("SERVER_START_TIMEOUT", 60)) + +TIMEOUT_SECONDS = int(os.environ.get("TIMEOUT_SECONDS", 1000)) +PROMPT_ENDPOINT = f"http://{SERVER_IP}/prompt" +V2V_WORKFLOWS_DIRECTORY = os.environ.get("V2V_WORKFLOWS_DIRECTORY", + "/workflow_configs") + + +COMFY_ROOT_DIRECTORY = os.environ.get("COMFY_ROOT_DIRECTORY", + "/app/ComfyUI") + + +BASE_LIVE_PORTRAIT_WORKFLOW = os.environ.get("BASE_LIVE_PORTRAIT_WORKFLOW", "16-07-2024/workflow_api_live_portrait-16-07-24.json") + +# V2V_WORKFLOWS_DIRECTORY = "/app/ComfyUI/exp" +# BASE_LIVE_PORTRAIT_WORKFLOW = "workflow_api_live_portrait-16-07-24.json" + +class PipelineType(IntEnum): + BASE = auto() + + +def ensure_input_directory_empty(): + input_directory = Path(COMFY_ROOT_DIRECTORY) / "input" + for file in input_directory.iterdir(): + if file.is_file(): + file.unlink() + +def ensure_output_directory_empty(): + output_directory = Path(COMFY_ROOT_DIRECTORY) / "output" / "vid2vid" + for file in output_directory.iterdir(): + if file.is_file(): + print('deleteing file: ' + str(file) ) + file.unlink() + +def ensure_pipeline_input_present(source_path, input_name=None): + left = Path(source_path) + left_filename = left.name + right_filename = left_filename + if input_name is not None: + right_filename = input_name + right = Path(COMFY_ROOT_DIRECTORY) / "input" / right_filename + shutil.copy(left, right) + + +def ensure_output_present_at(right): + left = Path(COMFY_ROOT_DIRECTORY) / "output" / "vid2vid" / "LivePortrait_00001.mp4" + shutil.copy(left, right) + +def parse_args(): + parser = argparse.ArgumentParser(description='Run Comfy') + parser.add_argument('--portrait-media-filename', type=str, help='path of mp4 depth from engine for preprocessing', required=True) + parser.add_argument('--driver-media-filename', type=str, help='path of mp4 outline from engine for preprocessing', required=True) + parser.add_argument('--input-is-image', action='store_true', help='input is image', required=True) + parser.add_argument('--tmpdir', type=str, help='path of tmpdir for preprocessing', required=False) + parser.add_argument('--output-filename', type=str, help='path of output filename', required=True) + + + return parser.parse_args() + +args = parse_args() + +def ensure_server_is_healthy(server_address): + # check the health of the server (5 second timeout) + response = requests.get(f"http://{server_address}/queue", timeout=5) + if response.status_code == 200: + print("Server is healthy") + else: + print("Server is unhealthy") + print(f"Status code: {response.status_code}") + print(f"Response: {response.text}") + + # make sure there are no jobs already + queue_json = response.json() + if len(queue_json.get("queue_pending", [])) > 0: + print("There are jobs in the queue already. Exiting...") + exit(1) + else: + print("Queue is empty") + + + +def main(): + ensure_server_is_healthy(SERVER_IP) + pipeline_type = PipelineType.BASE + print("Generating prompt json") + workflow_filename = BASE_LIVE_PORTRAIT_WORKFLOW + input_media_filename = args.portrait_media_filename + input_driver_filename = args.driver_media_filename + input_is_image = args.input_is_image + ensure_input_directory_empty() + ensure_output_directory_empty() + input_filename = "input.mp4" + if input_media_filename: + if input_is_image: + ensure_pipeline_input_present(input_media_filename) + input_filename = Path(input_media_filename).name + else: + ensure_pipeline_input_present(input_media_filename, "input.mp4") + + if input_driver_filename: + ensure_pipeline_input_present(input_driver_filename, "driver.mp4") + + + time_before = time.perf_counter() + + prompt = generate_prompt_for_style( + workflow_filename=workflow_filename, + input_media_filename="input/" + input_filename, + input_driver_filename="input/driver.mp4", + input_is_image=input_is_image, + ) + print("running pipeline" + str(pipeline_type)) + enqueue_prompt_and_wait(prompt) + + ensure_output_present_at(args.output_filename) + time_after = time.perf_counter() + print(f"Time taken: {time_after - time_before:.2f} seconds") + + +def enqueue_prompt_and_wait(prompt): + data = json.dumps(prompt).encode('utf-8') + # Send HTTP post request + response = requests.post(PROMPT_ENDPOINT, data=data) + if response.status_code == 200: + print("Job enqueued successfully") + else: + print("Job failed to queue") + print(f"Status code: {response.status_code}") + print(f"Response: {response.text}") + exit(1) + + # Wait for the job to finish + print(f"Waiting for job to finish (timeout: {TIMEOUT_SECONDS} seconds)") + start_time = time.time() + # check queue endpoint in loop + while True: + try: + response = requests.get(PROMPT_ENDPOINT, timeout=5) + queue_json = response.json() + if queue_json["exec_info"]["queue_remaining"]: + pass + else: + print("Job is done") + break + except Exception: + print("Error while checking queue") + time.sleep(1) + continue + if time.time() - start_time > TIMEOUT_SECONDS: + print("Timeout reached and job didn't finish!") + exit(1) + time.sleep(1) + + + +def generate_prompt_for_style( + workflow_filename: str, + input_media_filename: Optional[str] = None, + input_driver_filename: Optional[str] = None, + input_is_image: Optional[bool] = False, + ): + workflow_directory = Path(V2V_WORKFLOWS_DIRECTORY) / "workflows" + + + if workflow_filename: + workflow_json_filename = workflow_filename + + workflow_json = json.load(open(workflow_directory / workflow_json_filename)) + + + jsonpath_mods = { + "$.25.inputs.video": input_driver_filename + } + + if input_is_image: + jsonpath_mods["$.43.inputs.image_path"] = input_media_filename + jsonpath_mods["$.26.inputs.video"] = input_driver_filename + else: + jsonpath_mods["$.26.inputs.video"] = input_media_filename + jsonpath_mods["$.46.inputs.boolean"] = 1 + + print(*jsonpath_mods.items(), sep="\n") + + json_prompt = apply_jsonpath_mods(jsonpath_mods, workflow_json) + + p = {"prompt": json_prompt} + return p + +def apply_jsonpath_mods(jsonpath_mods: Dict[str, Any], workflow_json: Dict[str, Any]) -> Dict[str, Any]: + for key, value in jsonpath_mods.items(): + # jsonpath-ng expects numbers to be wrapped in quotes + parts = key.split(".") + new_parts = [f"'{part}'" if part.isdigit() else part for part in parts] + key = ".".join(new_parts) + + jsonpath_expr = parse(key) + matches = jsonpath_expr.find(workflow_json) + + if matches: + parent = matches[0].context.value + + if isinstance(parent, dict): + last_part = new_parts[-1].strip("'") + parent[last_part] = value + + elif isinstance(parent, list): + index = int(new_parts[-1].strip("'")) + if 0 <= index < len(parent): + parent[index] = value + else: + print(f"Index {index} is out of bounds for key '{key}'") + else: + print(f"No mapping found for key '{key}'") + + return workflow_json + + +if __name__ == "__main__": + os.setpgrp() + try: + comfyui_process = subprocess.Popen(["../venv/bin/python", "main.py", "--highvram", "--disable-metadata"], cwd=COMFY_ROOT_DIRECTORY) + start_time = time.time() + started = False + while time.time() - start_time < SERVER_START_TIMEOUT: + try: + response = requests.get(PROMPT_ENDPOINT, timeout=5) + if response.status_code == 200: + print("Server is healthy") + started = True + break + except requests.exceptions.ConnectionError: + print(f"Server is not up yet. Time elapsed: {time.time() - start_time:.2f}") + time.sleep(1) + if not started: + print(f"Server failed to start in {SERVER_START_TIMEOUT} seconds") + exit(1) + main() + finally: + os.killpg(0, signal.SIGKILL) # kill all processes in my group diff --git a/workflows/comfy/Dockerfile b/workflows/comfy/Dockerfile index e8c6ede..da5e5f5 100644 --- a/workflows/comfy/Dockerfile +++ b/workflows/comfy/Dockerfile @@ -25,6 +25,7 @@ RUN apt-get update \ python3-pip \ python3.10 \ python3.10-venv \ + python3.10-dev \ unzip \ htop \ vim \ @@ -49,37 +50,37 @@ RUN cd /tmp && \ RUN mkdir -p /root/.config/pip COPY ./pip.conf /root/.config/pip/pip.conf -WORKDIR /app # Copy files to container - -COPY --chown=appuser:appgroup install.py /app/install.py -RUN chmod -R 777 /app +USER appuser +WORKDIR /app +COPY install.py /app/install.py # download & install comfy ui and its dependencies -RUN python3 install.py --comfy-ui-url https://github.com/comfyanonymous/ComfyUI.git --comfy-ui-commit ece5acb8e8025d8ca26aa880f604d971d245475d +RUN python3 install.py --comfy-ui-url ${COMFY_GIT_URL} --comfy-ui-commit ${COMFY_GIT_SHA} RUN python3 -m venv venv RUN 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/ && \ venv/bin/python -m pip cache purge -COPY --chown=appuser:appgroup required_requirements.txt /app/required_requirements.txt +COPY required_requirements.txt /app/required_requirements.txt RUN 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/ && \ venv/bin/python -m pip uninstall onnxruntime onnxruntime-gpu -y && \ venv/bin/python -m pip install onnxruntime-gpu==1.17.1 ort-nightly-gpu --index-url=https://pkgs.dev.azure.com/onnxruntime/onnxruntime/_packaging/onnxruntime-cuda-12/pypi/simple/ --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/ && \ venv/bin/python -m pip cache purge -COPY --chown=appuser:appgroup repos.json /app/repos.json +COPY repos.json /app/repos.json RUN python3 install.py --repos-json repos.json -COPY --chown=appuser:appgroup custom_nodes/Yae_CustomNodes.py /app/ComfyUI/custom_nodes/Yae_CustomNodes.py +COPY custom_nodes/Yae_CustomNodes.py /app/ComfyUI/custom_nodes/Yae_CustomNodes.py -COPY --chown=appuser:appgroup dependencies.json /app/dependencies.json -COPY --chown=appuser:appgroup server_requirements.txt /app/server_requirements.txt - -COPY --chown=appuser:appgroup initialize_comfy_extensions.py /app/initialize_comfy_extensions.py -COPY --chown=appuser:appgroup list-repos.sh /app/list-repos.sh -COPY --chown=appuser:appgroup ComfyLauncher /app/ComfyLauncher -COPY --chown=appuser:appgroup format_video.py /app/format_video.py -COPY --chown=appuser:appgroup facefusion-runner.sh /app/facefusion-runner.sh +COPY dependencies.json /app/dependencies.json +COPY server_requirements.txt /app/server_requirements.txt +COPY initialize_comfy_extensions.py /app/initialize_comfy_extensions.py +COPY list-repos.sh /app/list-repos.sh +COPY ComfyLauncher /app/ComfyLauncher +COPY format_video.py /app/format_video.py +COPY facefusion-runner.sh /app/facefusion-runner.sh +#TODO(kasisnu, 15/07/2024): remove root user after confirming that it is not needed anymore +USER root # protobuf use python implementation (weird error without) -ENV PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python +ENV PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python \ No newline at end of file diff --git a/workflows/comfy/repos.json b/workflows/comfy/repos.json index f9e367b..abaf9ee 100644 --- a/workflows/comfy/repos.json +++ b/workflows/comfy/repos.json @@ -94,6 +94,11 @@ "commit": "2ead4fae1dddfe65d951248d98a1ff11dd50ab7a", "url": "https://github.com/kijai/ComfyUI-KJNodes.git" }, + { + "folder": "ComfyUI/custom_nodes/ComfyUI-LivePortraitKJ", + "commit": "2a7bd6116f32e99d3a40057eac222c3666c1dfb1", + "url": "https://github.com/kijai/ComfyUI-LivePortraitKJ" + }, { "folder": "ComfyUI/custom_nodes/ComfyUI-Impact-Pack", "commit": "a26fe487ef5c46506bd75872c82a3cfb002f7ea1", diff --git a/workflows/comfy/required_requirements.txt b/workflows/comfy/required_requirements.txt index eb901c9..302f99e 100644 --- a/workflows/comfy/required_requirements.txt +++ b/workflows/comfy/required_requirements.txt @@ -12,4 +12,5 @@ addict yapf deepdiff pynvml -pydantic \ No newline at end of file +pydantic +insightface==0.7.3 \ No newline at end of file