pipeline/storage.py

"""
Model storage module: persist voice reference files to HuggingFace Dataset repo.
"""

import os
import logging
from datetime import datetime

logger = logging.getLogger(__name__)

MODELS_REPO_ID = None
LOCAL_MODELS_DIR = "/tmp/rvc_models"


def init_storage(repo_id):
    """Initialize storage with the HF dataset repo ID."""
    global MODELS_REPO_ID
    MODELS_REPO_ID = repo_id
    os.makedirs(LOCAL_MODELS_DIR, exist_ok=True)
    logger.info("Storage initialized with repo: {}".format(repo_id))


def upload_model(model_name, pth_path, index_path=None, big_npy_path=None, reference_path=None):
    """Upload model files to HF dataset repo."""
    if not MODELS_REPO_ID:
        logger.warning("No HF repo configured. Model saved locally only.")
        return False

    try:
        from huggingface_hub import HfApi
        api = HfApi()

        # Upload .pth marker
        if pth_path and os.path.exists(pth_path):
            api.upload_file(
                path_or_fileobj=pth_path,
                path_in_repo="models/{}/{}.pth".format(model_name, model_name),
                repo_id=MODELS_REPO_ID,
                repo_type="dataset",
            )
            logger.info("Uploaded {}.pth to HF".format(model_name))

        # Upload reference audio
        if reference_path and os.path.exists(reference_path):
            ref_filename = os.path.basename(reference_path)
            api.upload_file(
                path_or_fileobj=reference_path,
                path_in_repo="models/{}/{}".format(model_name, ref_filename),
                repo_id=MODELS_REPO_ID,
                repo_type="dataset",
            )
            logger.info("Uploaded {} to HF".format(ref_filename))

        # Upload .index file if exists (backward compat)
        if index_path and os.path.exists(index_path):
            api.upload_file(
                path_or_fileobj=index_path,
                path_in_repo="models/{}/{}.index".format(model_name, model_name),
                repo_id=MODELS_REPO_ID,
                repo_type="dataset",
            )

        # Upload metadata
        metadata = {
            "name": model_name,
            "created": datetime.now().isoformat(),
            "engine": "seed-vc",
        }
        import json
        import tempfile

        with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
            json.dump(metadata, f)
            meta_path = f.name

        try:
            api.upload_file(
                path_or_fileobj=meta_path,
                path_in_repo="models/{}/metadata.json".format(model_name),
                repo_id=MODELS_REPO_ID,
                repo_type="dataset",
            )
        finally:
            os.unlink(meta_path)

        return True
    except Exception as e:
        logger.error("Failed to upload model: {}".format(e))
        return False


def download_model(model_name):
    """Download model from HF dataset repo. Returns (pth_path, reference_path)."""
    if not MODELS_REPO_ID:
        return _get_local_model(model_name)

    try:
        from huggingface_hub import hf_hub_download

        local_dir = os.path.join(LOCAL_MODELS_DIR, model_name)
        os.makedirs(local_dir, exist_ok=True)

        pth_path = hf_hub_download(
            repo_id=MODELS_REPO_ID,
            repo_type="dataset",
            filename="models/{}/{}.pth".format(model_name, model_name),
            local_dir=local_dir,
        )

        # Try to download reference audio
        ref_path = None
        try:
            ref_path = hf_hub_download(
                repo_id=MODELS_REPO_ID,
                repo_type="dataset",
                filename="models/{}/{}_ref.wav".format(model_name, model_name),
                local_dir=local_dir,
            )
        except Exception:
            # Try .index for backward compat with old RVC models
            try:
                ref_path = hf_hub_download(
                    repo_id=MODELS_REPO_ID,
                    repo_type="dataset",
                    filename="models/{}/{}.index".format(model_name, model_name),
                    local_dir=local_dir,
                )
            except Exception:
                pass

        return pth_path, ref_path
    except Exception as e:
        logger.error("Failed to download model from HF: {}".format(e))
        return _get_local_model(model_name)


def _get_local_model(model_name):
    """Get model from local storage."""
    local_dir = os.path.join(LOCAL_MODELS_DIR, model_name)
    pth_path = os.path.join(local_dir, "{}.pth".format(model_name))
    ref_path = os.path.join(local_dir, "{}_ref.wav".format(model_name))

    if os.path.exists(pth_path):
        return pth_path, ref_path if os.path.exists(ref_path) else None
    return None, None


def get_reference_path(model_name):
    """Get the reference audio path for a model."""
    local_dir = os.path.join(LOCAL_MODELS_DIR, model_name)
    ref_path = os.path.join(local_dir, "{}_ref.wav".format(model_name))
    if os.path.exists(ref_path):
        return ref_path
    # Search in subdirectories (HF download structure)
    for root, dirs, files in os.walk(local_dir):
        for f in files:
            if f.endswith("_ref.wav"):
                return os.path.join(root, f)
    return None


def list_models():
    """List all available models."""
    models = set()

    if MODELS_REPO_ID:
        try:
            from huggingface_hub import HfApi
            api = HfApi()
            files = api.list_repo_files(MODELS_REPO_ID, repo_type="dataset")
            for f in files:
                if f.startswith("models/") and f.endswith(".pth"):
                    parts = f.split("/")
                    if len(parts) >= 3:
                        models.add(parts[1])
        except Exception as e:
            logger.error("Failed to list models from HF: {}".format(e))

    if os.path.exists(LOCAL_MODELS_DIR):
        for name in os.listdir(LOCAL_MODELS_DIR):
            model_dir = os.path.join(LOCAL_MODELS_DIR, name)
            if os.path.isdir(model_dir):
                pth = os.path.join(model_dir, "{}.pth".format(name))
                if os.path.exists(pth):
                    models.add(name)

    return sorted(models)


def delete_model(model_name):
    """Delete a model from HF repo and local storage."""
    if MODELS_REPO_ID:
        try:
            from huggingface_hub import HfApi
            api = HfApi()
            files = api.list_repo_files(MODELS_REPO_ID, repo_type="dataset")
            for f in files:
                if f.startswith("models/{}/".format(model_name)):
                    api.delete_file(f, MODELS_REPO_ID, repo_type="dataset")
            logger.info("Deleted {} from HF repo".format(model_name))
        except Exception as e:
            logger.error("Failed to delete from HF: {}".format(e))

    import shutil
    local_dir = os.path.join(LOCAL_MODELS_DIR, model_name)
    if os.path.exists(local_dir):
        shutil.rmtree(local_dir)
        logger.info("Deleted {} from local storage".format(model_name))

    return True