51 lines
1.6 KiB
Python
51 lines
1.6 KiB
Python
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))
|