Implement KeySync ComfyUI wrapper
This commit is contained in:
magicwang1111
@@ -0,0 +1 @@
|
||||
"""Utilities for the ComfyUI KeySync wrapper."""
|
||||
@@ -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))
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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))
|
||||
@@ -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),
|
||||
}
|
||||
@@ -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())
|
||||
Reference in New Issue
Block a user