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", 180)) 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", "17-07-2024/yae_LivingPortrait_17-07_API.json") # V2V_WORKFLOWS_DIRECTORY = "/app/ComfyUI/exp" # BASE_LIVE_PORTRAIT_WORKFLOW = "yae_LivingPortrait_17-07_API.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('--skip-comfy-startup', action='store_true', help='controls if this script tries to start comfy', default=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 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() print("Generating prompt json using workflow: " + workflow_filename) 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 jsonpath_mods["$.46.inputs.boolean"] = 0 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: if not args.skip_comfy_startup: 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