storytold/storyteller-ml
2
0
Fork 0
mirror of https://github.com/storytold/storyteller-ml.git synced 2026-10-09 00:09:55 +00:00
Code Issues Packages Projects Releases Wiki Activity

add support for cogvide

Browse Source
This commit is contained in:
Kasisnu
2024-11-14 17:19:55 +05:30
parent 7e12c7c3a5
commit e95ea00436
1 changed files with 63 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
workflows/comfy/ComfyLauncher/ComfyRunnerForSidecar.py
+63 -15
View File
@@ -4,13 +4,12 @@ import os
import json
import time
from typing import Dict, Any, Optional
from jsonpath_ng import jsonpath, parse
from jsonpath_ng import parse
import shutil
from pathlib import Path
from enum import IntEnum, auto
import os
import subprocess
import concurrent.futures
import signal
import websockets
import httpx
@@ -18,7 +17,6 @@ import asyncio
import urllib
import uuid
import re
import imageio.v3 as iio
client_id = str(uuid.uuid4())
@@ -37,6 +35,7 @@ MAIN_IPA_WORKFLOW = os.environ.get("MAIN_IPA_WORKFLOW", "06-08-2024/yae_vid2vid_
FACE_DETAILER_WORKFLOW = os.environ.get("FACE_DETAILER_WORKFLOW", "3-06-2024/yae_vid2vid_FaceDetailer_1-07_API.json")
UPSCALER_WORKFLOW = os.environ.get("UPSCALER_WORKFLOW", "06-08-2024/yae_vid2vid_Upscale_25-07_API.json")
CINEMATIC_WORKFLOW = os.environ.get("CINEMATIC_WORKFLOW", "3-06-2024/yae_vid2vid_better_main+upscaler_04-07_API.json")
COG_V1_WORKFLOW = os.environ.get("COG_V1_WORKFLOW", "14-11-2024/yae_15-10_COG_API.json")
class PipelineType(IntEnum):
BASE = auto()
@@ -45,6 +44,7 @@ class PipelineType(IntEnum):
UPSCALER = auto()
CINEMATIC = auto()
PREVIEW = auto()
COG_V1 = auto()
style_to_filename = {
@@ -114,9 +114,6 @@ def generated_frames_count(frames_dir):
async def extract_frames(input_video, output_directory):
if not output_directory.exists():
output_directory.mkdir(parents=True)
# for idx, frame in enumerate(iio.imiter(input_video)):
# frame_idx = idx + 1
# iio.imwrite(f"{output_directory}/Frame_{frame_idx:05d}.jpg", frame)
command = [
'ffmpeg',
'-i', str(input_video),
@@ -168,7 +165,6 @@ def ensure_output_directory_empty():
(Path(COMFY_ROOT_DIRECTORY) / "output").mkdir(parents=True)
def parse_args():
parser = argparse.ArgumentParser(description='Run Comfy')
parser.add_argument('--prompt', type=str, help='location of the prompt json', required=False)
parser.add_argument('--style', type=str, help='style name', required=False)
@@ -191,6 +187,8 @@ def parse_args():
parser.add_argument('--generate-previews', help='generate previews', required=False, default=False, action=argparse.BooleanOptionalAction)
parser.add_argument('--preview-frames-directory', type=str, help='previews directory', required=False)
parser.add_argument('--enable-cinematic', help='enable cinematic', required=False, default=False, action=argparse.BooleanOptionalAction)
parser.add_argument('--enable-cogvideo', help='enable cog video', required=False, default=False, action=argparse.BooleanOptionalAction)
parser.add_argument('--skip-comfy-startup', action='store_true', help='controls if this script tries to start comfy', default=False)
parser.add_argument('--server-restart-trigger-file', type=str, help='server restart trigger file', required=False)
return parser.parse_args()
@@ -258,7 +256,6 @@ def invoke_facefusion(input_audio, input_video, output_video):
subprocess.check_output(command)
print("Facefusion completed successfully.")
async def main():
if not args.skip_comfy_startup:
print("Starting comfy server")
@@ -280,11 +277,49 @@ async def main():
# interrupt_comfy_queue(SERVER_IP)
wait_for_comfy_startup(SERVER_IP)
# print our process id, parent group id, comfyui process id, comfyui process group id, comfyui process group leader id, comfyui process group leader parent id
# print(f"Process ID: {os.getpid()}")
# print(f"Parent Group ID: {os.getpgrp()}")
# print(f"ComfyUI Process ID: {comfyui_process.pid}")
# print(f"ComfyUI Process Group ID: {os.getpgid(comfyui_process.pid)}")
# print(f"ComfyUI Process Group Leader ID: {os.getpgid(comfyui_process.pid)}")
# print(f"ComfyUI Process Group Leader Parent ID: {os.getppid()}")
pipeline_type = PipelineType.IPA
output_index = 0
file_path = args.prompt
prompt = None
if args.enable_cogvideo:
print("Generating prompt json for COG V1")
workflow_filename = COG_V1_WORKFLOW
positive_prompt_filename = args.positive_prompt_filename
step_count = 25
strength = 1.4
if args.strength != 1.0:
# Could we even make better spaghetti?
strength = args.strength
positive_prompt = ''
if positive_prompt_filename:
positive_prompt = open(positive_prompt_filename).read()
ensure_output_directory_empty()
prompt = generate_prompt_for_cog_v1(
workflow_filename,
positive_prompt,
step_count,
strength
)
print("running pipeline" + str(pipeline_type))
await enqueue_prompt_and_wait(prompt)
left = Path(COMFY_ROOT_DIRECTORY) / "output" / f"vid2vid/CogVideoX-I2V_00001.mp4"
if left.exists():
print("Pipeline completed successfully to output file: " + str(left))
else:
print("Pipeline failed to produce output")
exit(1)
restore_last_pipeline_output(left)
return
if args.style is not None:
print("Generating prompt json")
style = args.style
@@ -351,7 +386,7 @@ async def main():
time_before = time.perf_counter()
validate_style_name(style)
prompt = generate_prompt_for_style(
style,
positive_prompt,
@@ -369,7 +404,7 @@ async def main():
outline_video_filename=outline_video_filename,
normals_video_filename=normals_video_filename
)
print("running pipeline" + str(pipeline_type))
await enqueue_prompt_and_wait(prompt)
output_index = output_index + 1
@@ -547,8 +582,6 @@ async def enqueue_prompt_and_wait(prompt):
# Wait for the job to finish
print(f"Waiting for job to finish (timeout: {TIMEOUT_SECONDS} seconds)")
execution_start_time = time.perf_counter()
# import ipdb; ipdb.set_trace()
# max_size = 1048576 * 10 allows us to receive a really large payload from comfy
async with websockets.connect("ws://{}/ws?clientId={}".format(SERVER_IP, client_id), max_size=1048576 * 10) as ws:
await smart_wait(ws, prompt_id)
current_time = time.perf_counter()
@@ -557,6 +590,21 @@ async def enqueue_prompt_and_wait(prompt):
def generate_prompt_for_cog_v1(workflow_filename, positive_prompt, step_count, strength=1.4):
workflow_directory = Path(V2V_WORKFLOWS_DIRECTORY) / "workflows"
workflow_json = json.load(open(workflow_directory / workflow_filename))
modifications = {
"$.320.inputs.text": positive_prompt,
"$.464.inputs.value": step_count,
"$.465.inputs.value": strength,
}
json_prompt = apply_jsonpath_mods(modifications, workflow_json)
p = {"prompt": json_prompt, "client_id": client_id}
return p
def generate_prompt_for_style(style_name,
positive_prompt,
negative_prompt,
@@ -653,7 +701,7 @@ def get_jsonpath_mods(style_mods: Dict[str, Any],
outline_video_filename: str = None,
normals_video_filename: str = None
) -> Dict[str, Any]:
modifications = {}
new_mod_json = style_mods.copy()
# Process "loras" differently
@@ -686,8 +734,8 @@ def get_jsonpath_mods(style_mods: Dict[str, Any],
if frame_skip:
new_mod_json["every_nth_frame"] = frame_skip
for key, value in new_mod_json.items():
mapping_key = f"$.{key}"
jsonpath_expr = parse(mapping_key)