mirror of
https://github.com/storytold/storyteller-ml.git
synced 2026-10-28 09:08:54 +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
|
||||
Reference in New Issue
Block a user