38 lines
1.6 KiB
Python
38 lines
1.6 KiB
Python
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),
|
|
}
|