src/frontend/tabs/manage_models.py

"""
This module contains the code for the "Manage models" tab.
"""

from typings.extra import DropdownValue

from functools import partial

import gradio as gr
import pandas as pd

from backend.manage_voice_models import (
    delete_all_models,
    delete_models,
    download_online_model,
    filter_public_models_table,
    get_current_models,
    load_public_model_tags,
    load_public_models_table,
    upload_local_model,
)

from frontend.common import (
    PROGRESS_BAR,
    confirm_box_js,
    confirmation_harness,
    exception_harness,
    identity,
    update_dropdowns,
)


def _update_model_lists(
    num_components: int, value: DropdownValue = None, value_indices: list[int] = []
) -> gr.Dropdown | tuple[gr.Dropdown, ...]:
    """
    Updates the choices of one or more dropdown
    components to the current set of voice models.

    Optionally updates the default value of one or more of these components.

    Parameters
    ----------
    num_components : int
        Number of dropdown components to update.
    value : DropdownValue, optional
        New value for dropdown components.
    value_indices : list[int], default=[]
        Indices of dropdown components to update the value for.

    Returns
    -------
    gr.Dropdown | tuple[gr.Dropdown, ...]
        Updated dropdown component or components.
    """
    return update_dropdowns(get_current_models, num_components, value, value_indices)


def _filter_public_models_table_harness(
    tags: list[str], query: str, progress_bar: gr.Progress
) -> gr.Dataframe:
    """
    Filter the public models table based on tags and search query.

    Parameters
    ----------
    tags : list[str]
        Tags to filter the table by.
    query : str
        Search query to filter the table by.
    progress_bar : gr.Progress
        Progress bar to display progress.

    Returns
    -------
    gr.Dataframe
        The filtered public models table rendered in a Gradio dataframe.
    """
    models_table = filter_public_models_table(tags, query, progress_bar)
    return gr.Dataframe(value=models_table)


def _pub_dl_autofill(
    pub_models: pd.DataFrame, event: gr.SelectData
) -> tuple[gr.Textbox, gr.Textbox]:
    """
    Autofill download link and model name based on selected row in public models table.

    Parameters
    ----------
    pub_models : pd.DataFrame
        Public models table.
    event : gr.SelectData
        Event containing the selected row.

    Returns
    -------
    download_link : gr.Textbox
        Autofilled download link.
    model_name : gr.Textbox
        Autofilled model name.
    """
    event_index = event.index[0]
    url_str = pub_models.loc[event_index, "URL"]
    model_str = pub_models.loc[event_index, "Model Name"]

    return gr.Textbox(value=url_str), gr.Textbox(value=model_str)


def render(
    dummy_deletion_checkbox: gr.Checkbox,
    delete_confirmation: gr.State,
    rvc_models_to_delete: gr.Dropdown,
    rvc_model_1click: gr.Dropdown,
    rvc_model_multi: gr.Dropdown,
) -> None:
    """
    Render "Manage models" tab.

    Parameters
    ----------
    dummy_deletion_checkbox : gr.Checkbox
        Dummy component needed for deletion confirmation in the
        "Manage audio" tab and the "Manage models" tab.
    delete_confirmation : gr.State
        Component storing deletion confirmation status in the
        "Manage audio" tab and the "Manage models" tab.
    rvc_models_to_delete : gr.Dropdown
        Dropdown for selecting models to delete in the
        "Manage models" tab.
    rvc_model_1click : gr.Dropdown
        Dropdown for selecting models in the "One-click generation" tab.
    rvc_model_multi : gr.Dropdown
        Dropdown for selecting models in the "Multi-step generation" tab.
    """

    # Download tab
    with gr.Tab("Download model"):

        with gr.Accordion("View public models table", open=False):

            gr.Markdown("")
            gr.Markdown("HOW TO USE")
            gr.Markdown("- Filter models using tags or search bar")
            gr.Markdown("- Select a row to autofill the download link and model name")

            filter_tags = gr.CheckboxGroup(
                value=[],
                label="Show voice models with tags",
                choices=load_public_model_tags(),
            )
            search_query = gr.Textbox(label="Search")

            public_models_table = gr.DataFrame(
                value=load_public_models_table([]),
                headers=["Model Name", "Description", "Tags", "Credit", "Added", "URL"],
                label="Available Public Models",
                interactive=False,
            )

        with gr.Row():
            model_zip_link = gr.Textbox(
                label="Download link to model",
                info=(
                    "Should point to a zip file containing a .pth model file and an"
                    " optional .index file."
                ),
            )
            model_name = gr.Textbox(
                label="Model name", info="Enter a unique name for the model."
            )

        with gr.Row():
            download_btn = gr.Button("Download 🌐", variant="primary", scale=19)
            dl_output_message = gr.Textbox(
                label="Output message", interactive=False, scale=20
            )

        download_button_click = download_btn.click(
            partial(
                exception_harness(download_online_model), progress_bar=PROGRESS_BAR
            ),
            inputs=[model_zip_link, model_name],
            outputs=dl_output_message,
        )

        public_models_table.select(
            _pub_dl_autofill,
            inputs=public_models_table,
            outputs=[model_zip_link, model_name],
            show_progress="hidden",
        )
        search_query.change(
            partial(
                exception_harness(_filter_public_models_table_harness),
                progress_bar=PROGRESS_BAR,
            ),
            inputs=[filter_tags, search_query],
            outputs=public_models_table,
            show_progress="hidden",
        )
        filter_tags.select(
            partial(
                exception_harness(_filter_public_models_table_harness),
                progress_bar=PROGRESS_BAR,
            ),
            inputs=[filter_tags, search_query],
            outputs=public_models_table,
            show_progress="hidden",
        )

    # Upload tab
    with gr.Tab("Upload model"):
        with gr.Accordion("HOW TO USE"):
            gr.Markdown(
                "- Find locally trained RVC v2 model file (weights folder) and optional"
                " index file (logs/[name] folder)"
            )
            gr.Markdown(
                "- Upload model file and optional index file directly or compress into"
                " a zip file and upload that"
            )
            gr.Markdown("- Enter a unique name for the model")
            gr.Markdown("- Click 'Upload model'")

        with gr.Row():
            with gr.Column():
                model_files = gr.File(label="Files", file_count="multiple")

            local_model_name = gr.Textbox(label="Model name")

        with gr.Row():
            model_upload_button = gr.Button("Upload model", variant="primary", scale=19)
            local_upload_output_message = gr.Textbox(
                label="Output message", interactive=False, scale=20
            )
            model_upload_button_click = model_upload_button.click(
                partial(
                    exception_harness(upload_local_model), progress_bar=PROGRESS_BAR
                ),
                inputs=[model_files, local_model_name],
                outputs=local_upload_output_message,
            )

    with gr.Tab("Delete models"):
        with gr.Row():
            with gr.Column():
                rvc_models_to_delete.render()
            with gr.Column():
                rvc_models_deleted_message = gr.Textbox(
                    label="Output message", interactive=False
                )

        with gr.Row():
            with gr.Column():
                delete_models_button = gr.Button(
                    "Delete selected models", variant="secondary"
                )
                delete_all_models_button = gr.Button(
                    "Delete all models", variant="primary"
                )
            with gr.Column():
                pass
        delete_models_button_click = delete_models_button.click(
            # NOTE not sure why, but in order for subsequent event listener
            # to trigger, changes coming from the js code
            # have to be routed through an identity function which takes as
            # input some dummy component of type bool.
            identity,
            inputs=dummy_deletion_checkbox,
            outputs=delete_confirmation,
            js=confirm_box_js("Are you sure you want to delete the selected models?"),
            show_progress="hidden",
        ).then(
            partial(confirmation_harness(delete_models), progress_bar=PROGRESS_BAR),
            inputs=[delete_confirmation, rvc_models_to_delete],
            outputs=rvc_models_deleted_message,
        )

        delete_all_models_btn_click = delete_all_models_button.click(
            identity,
            inputs=dummy_deletion_checkbox,
            outputs=delete_confirmation,
            js=confirm_box_js("Are you sure you want to delete all models?"),
            show_progress="hidden",
        ).then(
            partial(confirmation_harness(delete_all_models), progress_bar=PROGRESS_BAR),
            inputs=delete_confirmation,
            outputs=rvc_models_deleted_message,
        )

    for click_event in [
        download_button_click,
        model_upload_button_click,
        delete_models_button_click,
        delete_all_models_btn_click,
    ]:
        click_event.success(
            partial(_update_model_lists, 3, [], [2]),
            outputs=[rvc_model_1click, rvc_model_multi, rvc_models_to_delete],
            show_progress="hidden",
        )