245 lines
8.1 KiB
Python
245 lines
8.1 KiB
Python
import os
|
|
import shutil
|
|
import subprocess
|
|
import threading
|
|
from pathlib import Path
|
|
|
|
from .ffmpeg_utils import check_ffmpeg, convert_audio_to_16k_wav, convert_video_to_25fps
|
|
from .file_utils import create_task_dir, ensure_exists, find_files_by_ext, safe_copy
|
|
from .media_utils import (
|
|
load_video_to_images,
|
|
save_audio_to_wav,
|
|
save_images_to_video,
|
|
trim_or_pad_audio,
|
|
)
|
|
|
|
|
|
def _bool_for_script(value):
|
|
return "True" if bool(value) else "False"
|
|
|
|
|
|
def _clean_position(position):
|
|
if position is None or str(position).strip() == "":
|
|
return "None"
|
|
return str(position).strip()
|
|
|
|
|
|
def find_latest_video(output_dir):
|
|
output_path = Path(output_dir)
|
|
if not output_path.is_dir():
|
|
raise RuntimeError(f"KeySync finished but output folder was not found: {output_path}")
|
|
videos = find_files_by_ext(output_path, ["mp4"])
|
|
if not videos:
|
|
raise RuntimeError(f"KeySync finished but no output mp4 was found in {output_path}")
|
|
return max(videos, key=lambda path: path.stat().st_mtime)
|
|
|
|
|
|
def build_command(
|
|
task_data_videos_dir,
|
|
task_data_audios_dir,
|
|
output_name,
|
|
keyframes_ckpt,
|
|
interpolation_ckpt,
|
|
compute_until,
|
|
fix_occlusion,
|
|
position,
|
|
start_frame,
|
|
):
|
|
return [
|
|
"bash",
|
|
"scripts/infer_raw.sh",
|
|
"--file_list",
|
|
str(task_data_videos_dir),
|
|
"--file_list_audio",
|
|
str(task_data_audios_dir),
|
|
"--output_folder",
|
|
str(output_name),
|
|
"--keyframes_ckpt",
|
|
str(keyframes_ckpt),
|
|
"--interpolation_ckpt",
|
|
str(interpolation_ckpt),
|
|
"--compute_until",
|
|
str(int(compute_until)),
|
|
"--fix_occlusion",
|
|
_bool_for_script(fix_occlusion),
|
|
"--position",
|
|
_clean_position(position),
|
|
"--start_frame",
|
|
str(int(start_frame)),
|
|
]
|
|
|
|
|
|
def _stream_pipe(pipe, log_file, console_prefix):
|
|
with pipe, open(log_file, "w", encoding="utf-8", errors="replace") as handle:
|
|
for line in iter(pipe.readline, ""):
|
|
handle.write(line)
|
|
handle.flush()
|
|
print(f"{console_prefix}{line}", end="")
|
|
|
|
|
|
def _subprocess_env(repo_path):
|
|
env = os.environ.copy()
|
|
venv_bin = Path(repo_path) / ".venv" / ("Scripts" if os.name == "nt" else "bin")
|
|
if venv_bin.exists():
|
|
env["PATH"] = str(venv_bin) + os.pathsep + env.get("PATH", "")
|
|
return env
|
|
|
|
|
|
def _ensure_script_python_available(env):
|
|
if shutil.which("python", path=env.get("PATH")) is None:
|
|
raise RuntimeError(
|
|
"KeySync official infer_raw.sh calls 'python', but python was not found on PATH. "
|
|
"Activate the KeySync Python environment before starting ComfyUI, create "
|
|
"keysync_repo_path/.venv, or install a python command such as python-is-python3."
|
|
)
|
|
|
|
|
|
def run_subprocess_with_logs(command, cwd, stdout_log, stderr_log, timeout_sec=0):
|
|
if shutil.which("bash") is None and shutil.which("bash.exe") is None:
|
|
raise RuntimeError(
|
|
"KeySync official inference script requires bash. Please run in WSL2 / Linux / Git Bash, "
|
|
"or implement direct python pipeline mode later."
|
|
)
|
|
|
|
stdout_log = Path(stdout_log)
|
|
stderr_log = Path(stderr_log)
|
|
stdout_log.parent.mkdir(parents=True, exist_ok=True)
|
|
stderr_log.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
print(f"[KeySync] Running command in {cwd}: {' '.join(command)}")
|
|
print(f"[KeySync] stdout log: {stdout_log}")
|
|
print(f"[KeySync] stderr log: {stderr_log}")
|
|
|
|
env = _subprocess_env(cwd)
|
|
_ensure_script_python_available(env)
|
|
|
|
process = subprocess.Popen(
|
|
command,
|
|
cwd=str(cwd),
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
text=True,
|
|
bufsize=1,
|
|
env=env,
|
|
)
|
|
|
|
stdout_thread = threading.Thread(target=_stream_pipe, args=(process.stdout, stdout_log, "[KeySync] "), daemon=True)
|
|
stderr_thread = threading.Thread(target=_stream_pipe, args=(process.stderr, stderr_log, "[KeySync:stderr] "), daemon=True)
|
|
stdout_thread.start()
|
|
stderr_thread.start()
|
|
|
|
try:
|
|
return_code = process.wait(timeout=None if int(timeout_sec) <= 0 else int(timeout_sec))
|
|
except subprocess.TimeoutExpired as exc:
|
|
process.kill()
|
|
stdout_thread.join(timeout=5)
|
|
stderr_thread.join(timeout=5)
|
|
raise RuntimeError(
|
|
f"KeySync subprocess timed out after {timeout_sec} seconds. "
|
|
f"stdout log: {stdout_log}; stderr log: {stderr_log}"
|
|
) from exc
|
|
|
|
stdout_thread.join(timeout=5)
|
|
stderr_thread.join(timeout=5)
|
|
|
|
if return_code != 0:
|
|
raise RuntimeError(
|
|
f"KeySync subprocess failed with exit code {return_code}. "
|
|
f"stdout log: {stdout_log}; stderr log: {stderr_log}"
|
|
)
|
|
|
|
return return_code
|
|
|
|
|
|
def run_keysync(
|
|
keysync_model,
|
|
images,
|
|
audio,
|
|
input_fps=25,
|
|
compute_until=30,
|
|
fix_occlusion=False,
|
|
position="None",
|
|
start_frame=0,
|
|
force_25fps_16k_audio=True,
|
|
trim_audio_to_frames=True,
|
|
keep_intermediate_outputs=True,
|
|
subprocess_timeout_sec=0,
|
|
use_subprocess=True,
|
|
):
|
|
if not use_subprocess:
|
|
raise NotImplementedError("KeySync v1 only supports subprocess mode with the official scripts/infer_raw.sh pipeline.")
|
|
|
|
check_ffmpeg()
|
|
|
|
repo_path = ensure_exists(keysync_model.get("keysync_repo_path"), "KeySync repo path", is_dir=True)
|
|
infer_script = repo_path / "scripts" / "infer_raw.sh"
|
|
if not infer_script.is_file():
|
|
raise FileNotFoundError("scripts/infer_raw.sh not found. Please clone antonibigata/keysync into keysync_repo_path.")
|
|
|
|
keyframes_ckpt = ensure_exists(keysync_model.get("keyframes_ckpt"), "keyframes_ckpt", is_dir=False)
|
|
interpolation_ckpt = ensure_exists(keysync_model.get("interpolation_ckpt"), "interpolation_ckpt", is_dir=False)
|
|
|
|
task = create_task_dir()
|
|
task_id = task["task_id"]
|
|
raw_video_path = task["task_dir"] / "raw_input.mp4"
|
|
raw_audio_path = task["task_dir"] / "raw_input.wav"
|
|
keysync_video_path = task["videos_dir"] / "input.mp4"
|
|
keysync_audio_path = task["audios_dir"] / "input.wav"
|
|
final_video_path = task["task_dir"] / "final.mp4"
|
|
|
|
input_fps = int(input_fps)
|
|
output_fps = 25 if force_25fps_16k_audio else input_fps
|
|
|
|
save_images_to_video(images, raw_video_path, fps=input_fps)
|
|
save_audio_to_wav(
|
|
audio,
|
|
raw_audio_path,
|
|
duration_sec=float(compute_until) if trim_audio_to_frames else None,
|
|
pad_to_duration=False,
|
|
)
|
|
|
|
if force_25fps_16k_audio:
|
|
convert_video_to_25fps(raw_video_path, keysync_video_path)
|
|
convert_audio_to_16k_wav(raw_audio_path, keysync_audio_path)
|
|
else:
|
|
safe_copy(raw_video_path, keysync_video_path)
|
|
safe_copy(raw_audio_path, keysync_audio_path)
|
|
|
|
command = build_command(
|
|
task_data_videos_dir=task["videos_dir"],
|
|
task_data_audios_dir=task["audios_dir"],
|
|
output_name=task_id,
|
|
keyframes_ckpt=keyframes_ckpt,
|
|
interpolation_ckpt=interpolation_ckpt,
|
|
compute_until=compute_until,
|
|
fix_occlusion=fix_occlusion,
|
|
position=position,
|
|
start_frame=start_frame,
|
|
)
|
|
|
|
run_subprocess_with_logs(
|
|
command=command,
|
|
cwd=repo_path,
|
|
stdout_log=task["logs_dir"] / "keysync_stdout.log",
|
|
stderr_log=task["logs_dir"] / "keysync_stderr.log",
|
|
timeout_sec=subprocess_timeout_sec,
|
|
)
|
|
|
|
repo_output_dir = repo_path / "outputs" / task_id
|
|
latest_video = find_latest_video(repo_output_dir)
|
|
safe_copy(latest_video, final_video_path)
|
|
|
|
output_images, decoded_fps = load_video_to_images(final_video_path)
|
|
output_fps = decoded_fps or output_fps
|
|
output_duration = output_images.shape[0] / float(output_fps)
|
|
output_audio = trim_or_pad_audio(audio, duration_sec=output_duration, pad_to_duration=bool(trim_audio_to_frames))
|
|
|
|
if not keep_intermediate_outputs:
|
|
shutil.rmtree(task["data_dir"], ignore_errors=True)
|
|
shutil.rmtree(task["outputs_dir"], ignore_errors=True)
|
|
shutil.rmtree(repo_output_dir, ignore_errors=True)
|
|
raw_video_path.unlink(missing_ok=True)
|
|
raw_audio_path.unlink(missing_ok=True)
|
|
|
|
return output_images, output_audio, int(output_fps), str(final_video_path.resolve())
|