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())