|
|
| from __future__ import annotations |
|
|
| import sys |
| import subprocess |
| from pathlib import Path |
|
|
| from huggingface_hub import snapshot_download |
|
|
|
|
| ROOT = Path(__file__).resolve().parent |
|
|
| ARC2FACE_CODE = ROOT / "Arc2Face" |
| MODEL_DIR = ROOT / "models" |
|
|
| GITHUB_REPO = "https://github.com/foivospar/Arc2Face.git" |
| HF_REPO = "FoivosPar/Arc2Face" |
|
|
| REQUIRED_MODEL_FILES = [ |
| "encoder/config.json", |
| "encoder/pytorch_model.bin", |
| "arc2face/config.json", |
| "arc2face/diffusion_pytorch_model.safetensors", |
| ] |
|
|
|
|
| def setup_arc2face(): |
|
|
| |
| |
| |
|
|
| package_init = ARC2FACE_CODE / "arc2face" / "__init__.py" |
|
|
| if not package_init.exists(): |
|
|
| if ARC2FACE_CODE.exists(): |
| raise RuntimeError( |
| f"Arc2Face directory exists but is incomplete: " |
| f"{ARC2FACE_CODE}. Remove it and restart the Space." |
| ) |
|
|
| print("Downloading Arc2Face GitHub repository...") |
|
|
| subprocess.run( |
| [ |
| "git", |
| "clone", |
| "--depth", |
| "1", |
| GITHUB_REPO, |
| str(ARC2FACE_CODE), |
| ], |
| check=True, |
| ) |
|
|
| |
| repo_path = str(ARC2FACE_CODE) |
|
|
| if repo_path not in sys.path: |
| sys.path.insert(0, repo_path) |
|
|
| if not package_init.exists(): |
| raise FileNotFoundError( |
| f"Arc2Face Python package not found: {package_init}" |
| ) |
|
|
| |
| |
| |
|
|
| MODEL_DIR.mkdir(parents=True, exist_ok=True) |
|
|
| missing_files = [ |
| filename |
| for filename in REQUIRED_MODEL_FILES |
| if not (MODEL_DIR / filename).exists() |
| ] |
|
|
| if missing_files: |
|
|
| print("Downloading Arc2Face pretrained weights...") |
|
|
| snapshot_download( |
| repo_id=HF_REPO, |
| local_dir=str(MODEL_DIR), |
| allow_patterns=REQUIRED_MODEL_FILES, |
| ) |
|
|
| |
| |
| |
|
|
| missing_files = [ |
| filename |
| for filename in REQUIRED_MODEL_FILES |
| if not (MODEL_DIR / filename).is_file() |
| ] |
|
|
| if missing_files: |
| raise FileNotFoundError( |
| "Arc2Face model download incomplete. Missing:\n" |
| + "\n".join(missing_files) |
| ) |
|
|
| print("Arc2Face code and pretrained weights are ready.") |
|
|
| return ARC2FACE_CODE, MODEL_DIR |
|
|
|
|
| if __name__ == "__main__": |
| setup_arc2face() |
|
|
