src/ultimate_rvc/core/manage/models.py

"""Module which defines functions to manage voice models."""

import re
import shutil
import urllib.request
import zipfile
from _collections_abc import Sequence
from pathlib import Path

import gradio as gr

from ultimate_rvc.common import RVC_MODELS_DIR
from ultimate_rvc.core.common import (
    FLAG_FILE,
    copy_files_to_new_dir,
    display_progress,
    json_load,
    validate_url,
)
from ultimate_rvc.core.exceptions import (
    Entity,
    Location,
    NotFoundError,
    NotProvidedError,
    UIMessage,
    UploadFormatError,
    UploadLimitError,
    VoiceModelExistsError,
    VoiceModelNotFoundError,
)
from ultimate_rvc.core.typing_extra import (
    ModelMetaData,
    ModelMetaDataList,
    ModelMetaDataPredicate,
    ModelMetaDataTable,
    ModelTagName,
)
from ultimate_rvc.typing_extra import StrPath

PUBLIC_MODELS_JSON = json_load(Path(__file__).parent / "public_models.json")
PUBLIC_MODELS_TABLE = ModelMetaDataTable.model_validate(PUBLIC_MODELS_JSON)


def get_saved_model_names() -> list[str]:
    """
    Get the names of all saved voice models.

    Returns
    -------
    list[str]
        A list of names of all saved voice models.

    """
    model_paths = RVC_MODELS_DIR.iterdir()
    names_to_remove = ["hubert_base.pt", "rmvpe.pt", FLAG_FILE.name]
    return sorted([
        model_path.name
        for model_path in model_paths
        if model_path.name not in names_to_remove
    ])


def load_public_models_table(
    predicates: Sequence[ModelMetaDataPredicate],
) -> ModelMetaDataList:
    """
    Load table containing metadata of public voice models, optionally
    filtered by a set of predicates.

    Parameters
    ----------
    predicates : Sequence[ModelMetaDataPredicate]
        Predicates to filter the metadata table by.

    Returns
    -------
    ModelMetaDataList
        List containing metadata for each public voice model that
        satisfies the given predicates.

    """
    return [
        [
            model.name,
            model.description,
            model.tags,
            model.credit,
            model.added,
            model.url,
        ]
        for model in PUBLIC_MODELS_TABLE.models
        if all(predicate(model) for predicate in predicates)
    ]


def get_public_model_tags() -> list[ModelTagName]:
    """
    get the names of all valid public voice model tags.

    Returns
    -------
    list[str]
        A list of names of all valid public voice model tags.

    """
    return [tag.name for tag in PUBLIC_MODELS_TABLE.tags]


def filter_public_models_table(
    tags: Sequence[str],
    query: str,
) -> ModelMetaDataList:
    """
    Filter table containing metadata of public voice models by tags and
    a search query.


    The search query is matched against the name, description, tags,
    credit,and added date of each entry in the metadata table. Case
    insensitive search is performed. If the search query is empty, the
    metadata table is filtered only bythe given tags.

    Parameters
    ----------
    tags : Sequence[str]
        Tags to filter the metadata table by.
    query : str
        Search query to filter the metadata table by.

    Returns
    -------
    ModelMetaDataList
        List containing metadata for each public voice model that
        match the given tags and search query.

    """

    def _tags_predicate(model: ModelMetaData) -> bool:
        return all(tag in model.tags for tag in tags)

    def _query_predicate(model: ModelMetaData) -> bool:
        return (
            query.lower()
            in (
                f"{model.name} {model.description} {' '.join(model.tags)} "
                f"{model.credit} {model.added}"
            ).lower()
            if query
            else True
        )

    filter_fns = [_tags_predicate, _query_predicate]

    return load_public_models_table(filter_fns)


def _extract_model(
    zip_file: StrPath,
    extraction_dir: StrPath,
    remove_incomplete: bool = True,
    remove_zip: bool = False,
) -> None:
    """
    Extract a zipped voice model to a directory.

    Parameters
    ----------
    zip_file : StrPath
        The path to a zip file containing the voice model to extract.
    extraction_dir : StrPath
        The path to the directory to extract the voice model to.

    remove_incomplete : bool, default=True
        Whether to remove the extraction directory if the extraction
        process fails.
    remove_zip : bool, default=False
        Whether to remove the zip file once the extraction process is
        complete.

    Raises
    ------
    NotFoundError
        If no model file is found in the extracted zip file.

    """
    extraction_path = Path(extraction_dir)
    zip_path = Path(zip_file)
    extraction_completed = False
    try:
        extraction_path.mkdir(parents=True)
        with zipfile.ZipFile(zip_path, "r") as zip_ref:
            zip_ref.extractall(extraction_path)
        file_path_map = {
            ext: Path(root, name)
            for root, _, files in extraction_path.walk()
            for name in files
            for ext in [".index", ".pth"]
            if Path(name).suffix == ext
            and Path(root, name).stat().st_size
            > 1024 * (100 if ext == ".index" else 1024 * 40)
        }
        if ".pth" not in file_path_map:
            raise NotFoundError(
                entity=Entity.MODEL_FILE,
                location=Location.EXTRACTED_ZIP_FILE,
                is_path=False,
            )

        # move model and index file to root of the extraction directory
        for file_path in file_path_map.values():
            file_path.rename(extraction_path / file_path.name)

        # remove any sub-directories within the extraction directory
        for path in extraction_path.iterdir():
            if path.is_dir():
                shutil.rmtree(path)
        extraction_completed = True
    finally:
        if not extraction_completed and remove_incomplete and extraction_path.is_dir():
            shutil.rmtree(extraction_path)
        if remove_zip and zip_path.exists():
            zip_path.unlink()


def download_model(
    url: str,
    name: str,
    progress_bar: gr.Progress | None = None,
    percentages: tuple[float, float] = (0.0, 0.5),
) -> None:
    """
    Download a zipped voice model.

    Parameters
    ----------
    url : str
        An URL pointing to a location where the zipped voice model can
        be downloaded from.
    name : str
        The name to give to the downloaded voice model.
    progress_bar : gr.Progress, optional
        Gradio progress bar to update.
    percentages : tuple[float, float], default=(0.0, 0.5)
        Percentages to display in the progress bar.

    Raises
    ------
    NotProvidedError
        If no URL or name is provided.
    VoiceModelExistsError
        If a voice model with the provided name already exists.

    """
    if not url:
        raise NotProvidedError(entity=Entity.URL)
    if not name:
        raise NotProvidedError(entity=Entity.MODEL_NAME)
    extraction_path = RVC_MODELS_DIR / name
    if extraction_path.exists():
        raise VoiceModelExistsError(name)

    validate_url(url)
    zip_name = url.split("/")[-1].split("?")[0]

    # NOTE in case huggingface link is a direct link rather
    # than a resolve link then convert it to a resolve link
    url = re.sub(
        r"https://huggingface.co/([^/]+)/([^/]+)/blob/(.*)",
        r"https://huggingface.co/\1/\2/resolve/\3",
        url,
    )
    if "pixeldrain.com" in url:
        url = f"https://pixeldrain.com/api/file/{zip_name}"

    display_progress(
        "[~] Downloading voice model ...",
        percentages[0],
        progress_bar,
    )
    urllib.request.urlretrieve(url, zip_name)  # noqa: S310

    display_progress("[~] Extracting zip file...", percentages[1], progress_bar)
    _extract_model(zip_name, extraction_path, remove_zip=True)


def upload_model(
    files: Sequence[StrPath],
    name: str,
    progress_bar: gr.Progress | None = None,
    percentage: float = 0.5,
) -> None:
    """
    Upload a voice model from either a zip file or a .pth file and an
    optional index file.

    Parameters
    ----------
    files : Sequence[StrPath]
        Paths to the files to upload.
    name : str
        The name to give to the uploaded voice model.
    progress_bar : gr.Progress, optional
        Gradio progress bar to update.
    percentage : float, default=0.5
        Percentage to display in the progress bar.

    Raises
    ------
    NotProvidedError
        If no file paths or name are provided.
    VoiceModelExistsError
        If a voice model with the provided name already
        exists.
    UploadFormatError
        If a single uploaded file is not a .pth file or a .zip file.
        If two uploaded files are not a .pth file and an .index file.
    UploadLimitError
        If more than two file paths are provided.

    """
    if not files:
        raise NotProvidedError(entity=Entity.FILES, ui_msg=UIMessage.NO_UPLOADED_FILES)
    if not name:
        raise NotProvidedError(entity=Entity.MODEL_NAME)
    model_dir_path = RVC_MODELS_DIR / name
    if model_dir_path.exists():
        raise VoiceModelExistsError(name)
    sorted_file_paths = sorted([Path(f) for f in files], key=lambda f: f.suffix)
    match sorted_file_paths:
        case [file_path]:
            if file_path.suffix == ".pth":
                display_progress("[~] Copying .pth file ...", percentage, progress_bar)
                copy_files_to_new_dir([file_path], model_dir_path)
            # NOTE a .pth file is actually itself a zip file
            elif zipfile.is_zipfile(file_path):
                display_progress("[~] Extracting zip file...", percentage, progress_bar)
                _extract_model(file_path, model_dir_path)
            else:
                raise UploadFormatError(
                    entity=Entity.FILES,
                    formats=[".pth", ".zip"],
                    multiple=False,
                )
        case [index_path, pth_path]:
            if index_path.suffix == ".index" and pth_path.suffix == ".pth":
                display_progress(
                    "[~] Copying .pth file and index file ...",
                    percentage,
                    progress_bar,
                )
                copy_files_to_new_dir([index_path, pth_path], model_dir_path)
            else:
                raise UploadFormatError(
                    entity=Entity.FILES,
                    formats=[".pth", ".index"],
                    multiple=True,
                )
        case _:
            raise UploadLimitError(entity=Entity.FILES, limit="two")


def delete_models(
    names: Sequence[str],
    progress_bar: gr.Progress | None = None,
    percentage: float = 0.5,
) -> None:
    """
    Delete one or more voice models.

    Parameters
    ----------
    names : Sequence[str]
        Names of the voice models to delete.
    progress_bar : gr.Progress, optional
        Gradio progress bar to update.
    percentage : float, default=0.5
        Percentage to display in the progress bar.

    Raises
    ------
    NotProvidedError
        If no names are provided.
    VoiceModelNotFoundError
        If a voice model with a provided name does not exist.

    """
    if not names:
        raise NotProvidedError(
            entity=Entity.MODEL_NAMES,
            ui_msg=UIMessage.NO_VOICE_MODELS,
        )
    display_progress(
        "[~] Deleting voice models ...",
        percentage,
        progress_bar,
    )
    for name in names:
        model_dir_path = RVC_MODELS_DIR / name
        if not model_dir_path.is_dir():
            raise VoiceModelNotFoundError(name)
        shutil.rmtree(model_dir_path)


def delete_all_models(
    progress_bar: gr.Progress | None = None,
    percentage: float = 0.5,
) -> None:
    """
    Delete all voice models.

    Parameters
    ----------
    progress_bar : gr.Progress, optional
        Gradio progress bar to update.
    percentage : float, default=0.5
        Percentage to display in the progress bar.

    """
    all_model_names = get_saved_model_names()
    display_progress("[~] Deleting all voice models ...", percentage, progress_bar)
    for model_name in all_model_names:
        model_dir_path = RVC_MODELS_DIR / model_name
        if model_dir_path.is_dir():
            shutil.rmtree(model_dir_path)