mirror of
https://github.com/storytold/storyteller-ml.git
synced 2026-10-09 00:09:55 +00:00
comfy: add support for live previews
This commit is contained in:
Kasisnu
@@ -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
|
||||
+18
-17
@@ -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
|
||||
@@ -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",
|
||||
|
||||
@@ -12,4 +12,5 @@ addict
|
||||
yapf
|
||||
deepdiff
|
||||
pynvml
|
||||
pydantic
|
||||
pydantic
|
||||
insightface==0.7.3
|
||||
Reference in New Issue
Block a user