download_models.py

#!/usr/bin/env python3
"""
download_models.py

Downloads all required models to /tmp/models (more space available on HF Spaces)
and creates a symlink at ./models → /tmp/models so main.py can find them
without any changes.

Run this ONCE before starting the server:
    python download_models.py
"""

import os
import sys
import logging
from pathlib import Path
from huggingface_hub import snapshot_download

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
)
logger = logging.getLogger(__name__)

# ── Config ────────────────────────────────────────────────────────────────────
TMP_MODELS_DIR = Path("/tmp/models")
SYMLINK_PATH   = Path("./models")          # what main.py reads from

# Map of  local_subdir → HuggingFace repo id
MODELS = {
    "urdu-speecht5-finetuned":  "ahmedjaved812/urdu-tts-phonemes-speecht5-finetuned",
    "urdu-tts-processor":       "ahmedjaved812/urdu-phonemes-speecht5-processor",
    "speecht5-hifigan":         "microsoft/speecht5_hifigan",
    "spkrec-xvect-voxceleb":    "speechbrain/spkrec-xvect-voxceleb",
}


def download_models():
    TMP_MODELS_DIR.mkdir(parents=True, exist_ok=True)
    logger.info(f"Downloading models to {TMP_MODELS_DIR}")

    for subdir, repo_id in MODELS.items():
        dest = TMP_MODELS_DIR / subdir

        if dest.exists() and any(dest.iterdir()):
            logger.info(f"✅ Already exists, skipping: {subdir}")
            continue

        logger.info(f"⬇️  Downloading {repo_id} → {dest}")
        try:
            snapshot_download(
                repo_id=repo_id,
                local_dir=str(dest),
                ignore_patterns=["*.msgpack", "flax_model*", "tf_model*", "*.h5"],
            )
            logger.info(f"✅ Done: {subdir}")
        except Exception as e:
            logger.error(f"❌ Failed to download {repo_id}: {e}")
            sys.exit(1)


def create_symlink():
    # Resolve to absolute path so the symlink works from any working directory
    target  = TMP_MODELS_DIR.resolve()
    symlink = SYMLINK_PATH.resolve()

    if symlink.is_symlink():
        if symlink.resolve() == target:
            logger.info(f"✅ Symlink already correct: {symlink} → {target}")
            return
        else:
            logger.info(f"🔁 Removing stale symlink: {symlink}")
            symlink.unlink()

    elif symlink.exists():
        # A real directory called 'models' exists — don't overwrite it
        logger.error(
            f"❌ '{symlink}' already exists as a real directory, not a symlink. "
            f"Remove it manually and re-run this script."
        )
        sys.exit(1)

    os.symlink(target, symlink)
    logger.info(f"✅ Symlink created: {symlink} → {target}")


def copy_embeddings():
    """
    Copy the embeddings/ directory from the app folder into /tmp/models/embeddings
    so it lives alongside the models and is accessible via the ./models symlink.
    """
    src = Path("./embeddings").resolve()
    dst = TMP_MODELS_DIR / "embeddings"

    if not src.exists():
        logger.error(f"❌ embeddings/ directory not found at {src}. Add normal_embedding.pt to it.")
        sys.exit(1)

    if dst.exists():
        logger.info(f"✅ Embeddings already copied, skipping: {dst}")
        return

    import shutil
    shutil.copytree(str(src), str(dst))
    logger.info(f"✅ Embeddings copied: {src} → {dst}")


if __name__ == "__main__":
    download_models()
    copy_embeddings()
    create_symlink()
    logger.info("🎉 All models ready. You can now start the server.")