From e08f887b247a1f8eabef5a046b840f6bbb9fdda5 Mon Sep 17 00:00:00 2001 From: magicwang1111 Date: Sat, 23 May 2026 14:05:13 +0800 Subject: [PATCH] Implement KeySync ComfyUI wrapper --- .gitignore | 2 + README.md | 173 +++++++++++++++++++++ __init__.py | 3 + keysync_wrapper/__init__.py | 1 + keysync_wrapper/download_utils.py | 50 ++++++ keysync_wrapper/ffmpeg_utils.py | 69 +++++++++ keysync_wrapper/file_utils.py | 105 +++++++++++++ keysync_wrapper/media_utils.py | 156 +++++++++++++++++++ keysync_wrapper/model_utils.py | 37 +++++ keysync_wrapper/runner.py | 244 ++++++++++++++++++++++++++++++ nodes.py | 155 +++++++++++++++++++ requirements.txt | 4 + 12 files changed, 999 insertions(+) create mode 100644 .gitignore create mode 100644 README.md create mode 100644 __init__.py create mode 100644 keysync_wrapper/__init__.py create mode 100644 keysync_wrapper/download_utils.py create mode 100644 keysync_wrapper/ffmpeg_utils.py create mode 100644 keysync_wrapper/file_utils.py create mode 100644 keysync_wrapper/media_utils.py create mode 100644 keysync_wrapper/model_utils.py create mode 100644 keysync_wrapper/runner.py create mode 100644 nodes.py create mode 100644 requirements.txt diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..43ae0e2 --- /dev/null +++ b/.gitignore @@ -0,0 +1,2 @@ +__pycache__/ +*.py[cod] diff --git a/README.md b/README.md new file mode 100644 index 0000000..c46cc4d --- /dev/null +++ b/README.md @@ -0,0 +1,173 @@ +# ComfyUI KeySync Wrapper + +This custom node package wraps the official KeySync inference pipeline for ComfyUI. + +It does not include KeySync weights and it does not rewrite KeySync model code. The first version calls the official `scripts/infer_raw.sh` script through `subprocess`. + +## Nodes + +- `KeySync Download Models` +- `KeySync Model Loader` +- `KeySync LipSync` +- `KeySync LipSync Advanced` + +Typical workflow: + +```text +Load Video (Path) -> KeySync LipSync.images +Load Audio (Path) -> KeySync LipSync.audio +KeySync LipSync.images/audio -> VideoHelperSuite Video Combine +``` + +`KeySync LipSync` returns: + +- `images`: ComfyUI `IMAGE` frames +- `audio`: ComfyUI `AUDIO` +- `frame_rate`: usually `25` +- `output_video_path`: path to the intermediate `final.mp4` + +## Install The Official KeySync Repo + +Clone the original KeySync repository into this wrapper folder: + +```bash +cd D:/ComfyUI/custom_nodes/ComfyUI-Keysync-wrapper +git clone https://github.com/antonibigata/keysync +``` + +Default repo path used by the loader: + +```text +D:/ComfyUI/custom_nodes/ComfyUI-Keysync-wrapper/keysync +``` + +## Download Models + +Recommended model directory: + +```text +D:/ComfyUI/models/keysync +``` + +You can use the `KeySync Download Models` node with `download=True`. + +Manual download is also supported: + +```bash +git lfs install +git clone https://huggingface.co/toninio19/keysync D:/ComfyUI/models/keysync +``` + +The loader expects: + +```text +D:/ComfyUI/models/keysync/checkpoints/keyframe_dub.pt +D:/ComfyUI/models/keysync/checkpoints/interpolation_dub.pt +``` + +The default UI intentionally hides repo and model directory fields. The wrapper checks these fixed paths: + +```text +KeySync repo: D:/ComfyUI/custom_nodes/ComfyUI-Keysync-wrapper/keysync +Model dir: D:/ComfyUI/models/keysync +``` + +## Runtime Notes + +KeySync has heavy dependencies. This wrapper intentionally does not install the full official KeySync dependency set into ComfyUI. + +Recommended runtime: + +- WSL2 or Linux +- `ffmpeg` available on `PATH` +- `bash` available on `PATH` +- official KeySync environment/dependencies installed according to the upstream README + +If you run native Windows ComfyUI, install Git Bash or use WSL2. The official inference script depends on bash. + +## Output Layout + +Each run creates an isolated task folder: + +```text +D:/ComfyUI/output/keysync/task_YYYYMMDD_HHMMSS_xxxxxxxx/ + data/ + videos/input.mp4 + audios/input.wav + logs/ + keysync_stdout.log + keysync_stderr.log + final.mp4 +``` + +The official KeySync repo also receives its normal output under: + +```text +/outputs// +``` + +Set `keep_intermediate_outputs=False` to remove the prepared input data and upstream output folder after `final.mp4` is copied. + +## Parameters + +The simple `KeySync LipSync` node exposes only the everyday parameters: + +- `input_fps` +- `compute_until` + +Use `KeySync LipSync Advanced` only when you need the official occlusion path: + +- `fix_occlusion`: enables KeySync's official SAM2-based occlusion mask path. +- `position`: point prompt for the occluding object, for example `450,450`. +- `start_frame`: frame index where the point prompt applies. + +Internal wrapper behavior is fixed for the simple node: + +- video is normalized to 25 fps +- audio is normalized to 16 kHz mono WAV +- subprocess mode is always used +- intermediate task files are kept for debugging +- no subprocess timeout is enforced by default + +Advanced KeySync parameters such as `resize_size`, `chunk_size`, `cond_aug`, and `decoding_t` are not exposed in v1 because the official `infer_raw.sh` script currently hardcodes them. + +## Occlusion And SAM + +The official KeySync code uses SAM2 for occlusion handling. The `position` field is a SAM2 point prompt used to track the object that blocks the face, such as a hand or microphone. This is powerful but awkward as a text field, so it is hidden from the simple node. + +SAM3-based tracking may be added later as a separate preprocessing node that outputs an occlusion mask or point prompt. It is not wired into v1 because the official KeySync `infer_raw.sh` path expects its own SAM2 implementation and checkpoint layout. + +## Common Errors + +Missing KeySync repo: + +```text +KeySync repo path does not exist +``` + +Missing official script: + +```text +scripts/infer_raw.sh not found +``` + +Missing checkpoints: + +```text +keyframes_ckpt does not exist +interpolation_ckpt does not exist +``` + +Missing ffmpeg: + +```text +ffmpeg not found. Please install ffmpeg and add it to PATH. +``` + +Missing bash: + +```text +KeySync official inference script requires bash. +``` + +When KeySync fails, the error message includes the stdout and stderr log paths. diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..39a8c6b --- /dev/null +++ b/__init__.py @@ -0,0 +1,3 @@ +from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/keysync_wrapper/__init__.py b/keysync_wrapper/__init__.py new file mode 100644 index 0000000..6a10748 --- /dev/null +++ b/keysync_wrapper/__init__.py @@ -0,0 +1 @@ +"""Utilities for the ComfyUI KeySync wrapper.""" diff --git a/keysync_wrapper/download_utils.py b/keysync_wrapper/download_utils.py new file mode 100644 index 0000000..b0dd8aa --- /dev/null +++ b/keysync_wrapper/download_utils.py @@ -0,0 +1,50 @@ +from pathlib import Path + +from .file_utils import normalize_path, resolve_path + + +REQUIRED_CHECKPOINTS = ( + "checkpoints/keyframe_dub.pt", + "checkpoints/interpolation_dub.pt", +) + + +def validate_required_models(models_dir): + root = resolve_path(models_dir) + missing = [rel for rel in REQUIRED_CHECKPOINTS if not (root / rel).is_file()] + if missing: + missing_text = ", ".join(missing) + raise FileNotFoundError( + f"KeySync model files are missing in {root}: {missing_text}. " + "Run the KeySync Download Models node with download=True, or clone " + "https://huggingface.co/toninio19/keysync into this folder." + ) + return root + + +def download_or_check_models(models_dir, repo_id="toninio19/keysync", download=False): + root = Path(normalize_path(models_dir)) + + if download: + try: + from huggingface_hub import snapshot_download + except ImportError as exc: + raise RuntimeError( + "huggingface_hub is required for KeySync Download Models. " + "Install this wrapper's requirements.txt first." + ) from exc + + root.mkdir(parents=True, exist_ok=True) + snapshot_download( + repo_id=repo_id, + local_dir=str(root), + allow_patterns=["checkpoints/*"], + repo_type="model", + ) + elif not root.exists(): + raise FileNotFoundError( + f"models_dir does not exist: {root}. " + "Run KeySync Download Models with download=True or download the models manually." + ) + + return str(validate_required_models(root)) diff --git a/keysync_wrapper/ffmpeg_utils.py b/keysync_wrapper/ffmpeg_utils.py new file mode 100644 index 0000000..ef2a411 --- /dev/null +++ b/keysync_wrapper/ffmpeg_utils.py @@ -0,0 +1,69 @@ +import shutil +import subprocess +from pathlib import Path + +from .file_utils import resolve_path + + +def check_ffmpeg(): + ffmpeg = shutil.which("ffmpeg") or shutil.which("ffmpeg.exe") + if ffmpeg is None: + raise RuntimeError("ffmpeg not found. Please install ffmpeg and add it to PATH.") + return ffmpeg + + +def run_ffmpeg(args, error_label): + ffmpeg = check_ffmpeg() + command = [ffmpeg] + [str(arg) for arg in args] + result = subprocess.run(command, capture_output=True, text=True) + if result.returncode != 0: + stderr = result.stderr.strip() + raise RuntimeError(f"{error_label} failed with exit code {result.returncode}: {stderr}") + + +def convert_video_to_25fps(input_path, output_path): + input_file = resolve_path(input_path) + output_file = resolve_path(output_path) + output_file.parent.mkdir(parents=True, exist_ok=True) + run_ffmpeg( + [ + "-y", + "-i", + input_file, + "-r", + "25", + "-an", + "-c:v", + "libx264", + "-pix_fmt", + "yuv420p", + "-crf", + "18", + output_file, + ], + "25fps video conversion", + ) + return output_file + + +def convert_audio_to_16k_wav(input_path, output_path): + input_file = resolve_path(input_path) + output_file = resolve_path(output_path) + output_file.parent.mkdir(parents=True, exist_ok=True) + run_ffmpeg( + [ + "-y", + "-i", + input_file, + "-vn", + "-ac", + "1", + "-ar", + "16000", + "-acodec", + "pcm_s16le", + output_file, + ], + "16kHz audio conversion", + ) + return output_file diff --git a/keysync_wrapper/file_utils.py b/keysync_wrapper/file_utils.py new file mode 100644 index 0000000..183e11e --- /dev/null +++ b/keysync_wrapper/file_utils.py @@ -0,0 +1,105 @@ +import os +import re +import shutil +import uuid +from datetime import datetime +from pathlib import Path + + +WINDOWS_DRIVE_RE = re.compile(r"^([A-Za-z]):[\\/](.*)$") + + +def is_windows_drive_path(path): + return isinstance(path, str) and WINDOWS_DRIVE_RE.match(path.strip()) is not None + + +def normalize_path(path): + if path is None: + return path + + raw_path = str(path).strip().strip('"') + if not raw_path: + return raw_path + + match = WINDOWS_DRIVE_RE.match(raw_path) + if match and os.name != "nt": + drive = match.group(1).lower() + rest = match.group(2).replace("\\", "/") + return f"/mnt/{drive}/{rest}" + + return os.path.abspath(os.path.expanduser(raw_path)) + + +def resolve_path(path): + normalized = normalize_path(path) + return Path(normalized).expanduser().resolve() + + +def get_comfy_output_dir(): + try: + import folder_paths + + return Path(folder_paths.get_output_directory()).resolve() + except Exception: + current = Path(__file__).resolve() + for parent in current.parents: + if (parent / "custom_nodes").exists() and (parent / "output").exists(): + return (parent / "output").resolve() + return Path("/mnt/d/ComfyUI/output").resolve() + + +def ensure_exists(path, label="path", is_dir=None): + if path is None or str(path).strip() == "": + raise FileNotFoundError(f"{label} is empty.") + + resolved = resolve_path(path) + if is_dir is True and not resolved.is_dir(): + raise FileNotFoundError(f"{label} does not exist: {resolved}") + if is_dir is False and not resolved.is_file(): + raise FileNotFoundError(f"{label} does not exist: {resolved}") + if is_dir is None and not resolved.exists(): + raise FileNotFoundError(f"{label} does not exist: {resolved}") + return resolved + + +def safe_copy(src, dst): + src_path = ensure_exists(src, "source file", is_dir=False) + dst_path = resolve_path(dst) + dst_path.parent.mkdir(parents=True, exist_ok=True) + if src_path.resolve() != dst_path.resolve(): + shutil.copy2(src_path, dst_path) + return dst_path + + +def find_files_by_ext(root, extensions): + root_path = ensure_exists(root, "search root", is_dir=True) + wanted = {ext.lower().lstrip(".") for ext in extensions} + matches = [] + for path in root_path.rglob("*"): + if path.is_file() and path.suffix.lower().lstrip(".") in wanted: + matches.append(path) + return matches + + +def create_task_dir(prefix="task"): + output_root = get_comfy_output_dir() / "keysync" + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + suffix = uuid.uuid4().hex[:8] + task_id = f"{prefix}_{timestamp}_{suffix}" + task_dir = output_root / task_id + + paths = { + "task_id": task_id, + "task_dir": task_dir, + "data_dir": task_dir / "data", + "videos_dir": task_dir / "data" / "videos", + "audios_dir": task_dir / "data" / "audios", + "outputs_dir": task_dir / "outputs", + "logs_dir": task_dir / "logs", + } + + for key, value in paths.items(): + if key.endswith("_dir"): + value.mkdir(parents=True, exist_ok=True) + + return paths diff --git a/keysync_wrapper/media_utils.py b/keysync_wrapper/media_utils.py new file mode 100644 index 0000000..8045a3b --- /dev/null +++ b/keysync_wrapper/media_utils.py @@ -0,0 +1,156 @@ +import subprocess +from pathlib import Path + +import cv2 +import numpy as np +import soundfile as sf +import torch + +from .ffmpeg_utils import check_ffmpeg +from .file_utils import resolve_path + + +def _as_image_tensor(images): + if isinstance(images, list): + images = torch.stack(images) + if not isinstance(images, torch.Tensor): + images = torch.as_tensor(images) + if images.dim() != 4: + raise ValueError(f"Expected IMAGE tensor with shape [frames, height, width, channels], got {tuple(images.shape)}") + if images.shape[-1] < 3: + raise ValueError(f"Expected IMAGE tensor with at least 3 channels, got {tuple(images.shape)}") + return images.detach().cpu() + + +def _frame_to_uint8(frame): + frame = frame[..., :3] + if frame.dtype != torch.uint8: + frame = frame.float() + if float(frame.max()) <= 1.0: + frame = frame * 255.0 + frame = frame.clamp(0, 255).byte() + return np.ascontiguousarray(frame.numpy()) + + +def save_images_to_video(images, output_path, fps): + image_tensor = _as_image_tensor(images) + output_file = resolve_path(output_path) + output_file.parent.mkdir(parents=True, exist_ok=True) + + frame_count, height, width, _ = image_tensor.shape + if frame_count <= 0: + raise ValueError("No frames were provided to save_images_to_video.") + + ffmpeg = check_ffmpeg() + command = [ + ffmpeg, + "-y", + "-f", + "rawvideo", + "-vcodec", + "rawvideo", + "-s", + f"{width}x{height}", + "-pix_fmt", + "rgb24", + "-r", + str(fps), + "-i", + "-", + "-an", + "-c:v", + "libx264", + "-pix_fmt", + "yuv420p", + "-crf", + "18", + str(output_file), + ] + + process = subprocess.Popen(command, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE) + try: + for frame in image_tensor: + process.stdin.write(_frame_to_uint8(frame).tobytes()) + stdout, stderr = process.communicate() + except BrokenPipeError: + stdout, stderr = process.communicate() + finally: + if process.stdin and not process.stdin.closed: + process.stdin.close() + + if process.returncode != 0: + message = stderr.decode("utf-8", errors="replace").strip() + raise RuntimeError(f"Failed to encode IMAGE frames to mp4: {message}") + + return output_file + + +def _audio_waveform_to_tensor(audio): + if not isinstance(audio, dict) or "waveform" not in audio or "sample_rate" not in audio: + raise ValueError("AUDIO input must be a dict with 'waveform' and 'sample_rate'.") + + waveform = audio["waveform"] + if not isinstance(waveform, torch.Tensor): + waveform = torch.as_tensor(waveform) + waveform = waveform.detach().cpu().float() + + if waveform.dim() == 3: + waveform = waveform[0] + elif waveform.dim() == 1: + waveform = waveform.unsqueeze(0) + elif waveform.dim() != 2: + raise ValueError(f"Expected AUDIO waveform with shape [batch, channels, samples] or [channels, samples], got {tuple(waveform.shape)}") + + sample_rate = int(audio["sample_rate"]) + return waveform, sample_rate + + +def trim_or_pad_audio(audio, duration_sec=None, pad_to_duration=False): + waveform, sample_rate = _audio_waveform_to_tensor(audio) + + if duration_sec is not None and duration_sec > 0: + target_samples = max(1, int(round(float(duration_sec) * sample_rate))) + if waveform.shape[-1] > target_samples: + waveform = waveform[:, :target_samples] + elif pad_to_duration and waveform.shape[-1] < target_samples: + pad = torch.zeros((waveform.shape[0], target_samples - waveform.shape[-1]), dtype=waveform.dtype) + waveform = torch.cat([waveform, pad], dim=-1) + + return {"waveform": waveform.unsqueeze(0).contiguous(), "sample_rate": sample_rate} + + +def save_audio_to_wav(audio, output_path, duration_sec=None, pad_to_duration=False): + adjusted = trim_or_pad_audio(audio, duration_sec=duration_sec, pad_to_duration=pad_to_duration) + waveform = adjusted["waveform"].squeeze(0).clamp(-1.0, 1.0) + sample_rate = int(adjusted["sample_rate"]) + + output_file = resolve_path(output_path) + output_file.parent.mkdir(parents=True, exist_ok=True) + audio_np = waveform.transpose(0, 1).numpy() + sf.write(str(output_file), audio_np, sample_rate, subtype="PCM_16") + return output_file + + +def load_video_to_images(video_path): + video_file = resolve_path(video_path) + capture = cv2.VideoCapture(str(video_file)) + if not capture.isOpened(): + raise RuntimeError(f"Failed to open output video: {video_file}") + + fps = capture.get(cv2.CAP_PROP_FPS) + frames = [] + try: + while True: + ok, frame_bgr = capture.read() + if not ok: + break + frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB) + frames.append(frame_rgb) + finally: + capture.release() + + if not frames: + raise RuntimeError(f"No frames were decoded from output video: {video_file}") + + array = np.stack(frames).astype(np.float32) / 255.0 + return torch.from_numpy(array), int(round(fps or 25)) diff --git a/keysync_wrapper/model_utils.py b/keysync_wrapper/model_utils.py new file mode 100644 index 0000000..e605133 --- /dev/null +++ b/keysync_wrapper/model_utils.py @@ -0,0 +1,37 @@ +import os +from pathlib import Path + +from .file_utils import ensure_exists, is_windows_drive_path, normalize_path, resolve_path + + +def _resolve_checkpoint(models_dir, checkpoint): + checkpoint_text = str(checkpoint).strip().strip('"') + if os.path.isabs(checkpoint_text) or is_windows_drive_path(checkpoint_text): + return resolve_path(checkpoint_text) + return resolve_path(Path(normalize_path(models_dir)) / checkpoint_text) + + +def load_keysync_model_config(keysync_repo_path, models_dir, keyframes_ckpt, interpolation_ckpt): + repo_path = ensure_exists(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( + f"scripts/infer_raw.sh not found. Please clone antonibigata/keysync into keysync_repo_path: {repo_path}" + ) + + models_root = ensure_exists(models_dir, "models_dir", is_dir=True) + keyframes_path = _resolve_checkpoint(models_root, keyframes_ckpt) + interpolation_path = _resolve_checkpoint(models_root, interpolation_ckpt) + + if not keyframes_path.is_file(): + raise FileNotFoundError(f"keyframes_ckpt does not exist: {keyframes_path}") + if not interpolation_path.is_file(): + raise FileNotFoundError(f"interpolation_ckpt does not exist: {interpolation_path}") + + return { + "keysync_repo_path": str(repo_path), + "models_dir": str(models_root), + "keyframes_ckpt": str(keyframes_path), + "interpolation_ckpt": str(interpolation_path), + "infer_script": str(infer_script), + } diff --git a/keysync_wrapper/runner.py b/keysync_wrapper/runner.py new file mode 100644 index 0000000..a84a55a --- /dev/null +++ b/keysync_wrapper/runner.py @@ -0,0 +1,244 @@ +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()) diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..76e179c --- /dev/null +++ b/nodes.py @@ -0,0 +1,155 @@ +from .keysync_wrapper.download_utils import download_or_check_models +from .keysync_wrapper.model_utils import load_keysync_model_config +from .keysync_wrapper.runner import run_keysync + + +DEFAULT_MODELS_DIR = "D:/ComfyUI/models/keysync" +DEFAULT_KEYSYNC_REPO_PATH = "D:/ComfyUI/custom_nodes/ComfyUI-Keysync-wrapper/keysync" + + +class KeySyncDownloadModels: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "download": ("BOOLEAN", {"default": False}), + } + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("models_dir",) + FUNCTION = "run" + CATEGORY = "video/lipsync" + + def run(self, download): + resolved_models_dir = download_or_check_models( + models_dir=DEFAULT_MODELS_DIR, + repo_id="toninio19/keysync", + download=download, + ) + return (resolved_models_dir,) + + +class KeySyncModelLoader: + @classmethod + def INPUT_TYPES(cls): + return {"required": {}} + + RETURN_TYPES = ("KEYSYNC_MODEL",) + RETURN_NAMES = ("keysync_model",) + FUNCTION = "load" + CATEGORY = "video/lipsync" + + def load(self): + model = load_keysync_model_config( + keysync_repo_path=DEFAULT_KEYSYNC_REPO_PATH, + models_dir=DEFAULT_MODELS_DIR, + keyframes_ckpt="checkpoints/keyframe_dub.pt", + interpolation_ckpt="checkpoints/interpolation_dub.pt", + ) + return (model,) + + +class KeySyncLipSync: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "keysync_model": ("KEYSYNC_MODEL",), + "images": ("IMAGE",), + "audio": ("AUDIO",), + "input_fps": ("INT", {"default": 25, "min": 1, "max": 120, "step": 1}), + "compute_until": ("INT", {"default": 30, "min": 1, "max": 300, "step": 1}), + } + } + + RETURN_TYPES = ("IMAGE", "AUDIO", "INT", "STRING") + RETURN_NAMES = ("images", "audio", "frame_rate", "output_video_path") + FUNCTION = "run" + CATEGORY = "video/lipsync" + + def run( + self, + keysync_model, + images, + audio, + input_fps, + compute_until, + ): + return run_keysync( + keysync_model=keysync_model, + images=images, + audio=audio, + input_fps=input_fps, + compute_until=compute_until, + 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, + ) + + +class KeySyncLipSyncAdvanced(KeySyncLipSync): + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "keysync_model": ("KEYSYNC_MODEL",), + "images": ("IMAGE",), + "audio": ("AUDIO",), + "input_fps": ("INT", {"default": 25, "min": 1, "max": 120, "step": 1}), + "compute_until": ("INT", {"default": 30, "min": 1, "max": 300, "step": 1}), + "fix_occlusion": ("BOOLEAN", {"default": False}), + "position": ("STRING", {"default": "None", "multiline": False}), + "start_frame": ("INT", {"default": 0, "min": 0, "max": 100000, "step": 1}), + } + } + + FUNCTION = "run_advanced" + CATEGORY = "video/lipsync" + + def run_advanced( + self, + keysync_model, + images, + audio, + input_fps, + compute_until, + fix_occlusion, + position, + start_frame, + ): + return run_keysync( + keysync_model=keysync_model, + images=images, + audio=audio, + input_fps=input_fps, + compute_until=compute_until, + fix_occlusion=fix_occlusion, + position=position, + start_frame=start_frame, + force_25fps_16k_audio=True, + trim_audio_to_frames=True, + keep_intermediate_outputs=True, + subprocess_timeout_sec=0, + use_subprocess=True, + ) + + +NODE_CLASS_MAPPINGS = { + "KeySyncDownloadModels": KeySyncDownloadModels, + "KeySyncModelLoader": KeySyncModelLoader, + "KeySyncLipSync": KeySyncLipSync, + "KeySyncLipSyncAdvanced": KeySyncLipSyncAdvanced, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "KeySyncDownloadModels": "KeySync Download Models", + "KeySyncModelLoader": "KeySync Model Loader", + "KeySyncLipSync": "KeySync LipSync", + "KeySyncLipSyncAdvanced": "KeySync LipSync Advanced", +} diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..7c0100e --- /dev/null +++ b/requirements.txt @@ -0,0 +1,4 @@ +opencv-python +moviepy +soundfile +huggingface_hub