src/webui.py

"""Module which defines the code for the "Manage models" tab."""

from collections.abc import Sequence
from functools import partial

import gradio as gr
import pandas as pd
import requests

# Function to search for RVC models on Hugging Face
def search_rvc_models(query):
    url = f"https://huggingface.co/api/models?search={query}&library=rvc"
    response = requests.get(url)
    if response.status_code == 200:
        models = response.json()
        # Create a DataFrame to store the results
        df = pd.DataFrame(models)
        # Filter the DataFrame to only include the desired columns
        df = df[["id", "likes", "downloads"]]
        # Add a new column for the download URL
        df["downloadUrl"] = "https://huggingface.co/" + df["id"]
        # Sort the DataFrame by downloads in descending order
        df = df.sort_values(by="downloads", ascending=False)
        return df
    else:
        return pd.DataFrame({"id": ["No models found"]})

from ultimate_rvc.core.manage.models import (
    delete_all_models,
    delete_models,
    download_model,
    filter_public_models_table,
    get_public_model_tags,
    get_saved_model_names,
    upload_model,
)
from ultimate_rvc.web.common import (
    PROGRESS_BAR,
    confirm_box_js,
    confirmation_harness,
    exception_harness,
    render_msg,
    update_dropdowns,
)
from ultimate_rvc.web.typing_extra import DropdownValue


def _update_models(
    num_components: int,
    value: DropdownValue = None,
    value_indices: Sequence[int] = [],
) -> gr.Dropdown | tuple[gr.Dropdown, ...]:
    """
    Update the choices of one or more dropdown components to the set of
    currently saved 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 : Sequence[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_saved_model_names, num_components, value, value_indices)


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

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

    Returns
    -------
    gr.Dataframe
        The filtered table rendered in a Gradio dataframe.

    """
    models_table = filter_public_models_table(tags, query)
    return gr.Dataframe(value=models_table)


def _autofill_model_name_and_url(
    public_models_table: pd.DataFrame,
    select_event: gr.SelectData,
) -> tuple[gr.Textbox, gr.Textbox]:
    """
    Autofill two textboxes with respectively the name and URL that is
    saved in the currently selected row of the public models table.

    Parameters
    ----------
    public_models_table : pd.DataFrame
        The public models table saved in a Pandas dataframe.
    select_event : gr.SelectData
        Event containing the index of the currently selected row in the
        public models table.

    Returns
    -------
    name : gr.Textbox
        The textbox containing the model name.

    url : gr.Textbox
        The textbox containing the model URL.

    Raises
    ------
    TypeError
        If the index in the provided event is not a sequence.

    """
    event_index = select_event.index
    if not isinstance(event_index, Sequence):
        err_msg = (
            f"Expected a sequence of indices but got {type(event_index)} from the"
            " provided event."
        )
        raise TypeError(err_msg)
    event_index = event_index[0]
    url = public_models_table.loc[event_index, "URL"]
    name = public_models_table.loc[event_index, "Name"]
    if isinstance(url, str) and isinstance(name, str):
        return gr.Textbox(value=name), gr.Textbox(value=url)
    err_msg = (
        "Expected model name and URL to be strings but got"
        f" {type(name)} and {type(url)} respectively."
    )
    raise TypeError(err_msg)


def render(
    model_delete: gr.Dropdown,
    model_1click: gr.Dropdown,
    model_multi: gr.Dropdown,
) -> None:
    """

    Render "Manage models" tab.

    Parameters
    ----------
    model_delete : gr.Dropdown
        Dropdown for selecting voice models to delete in the
        "Delete models" tab.
    model_1click : gr.Dropdown
        Dropdown for selecting a voice model to use in the
        "One-click generation" tab.
    model_multi : gr.Dropdown
        Dropdown for selecting a voice model to use in the
        "Multi-step generation" tab.

    """
    # Download tab

    dummy_checkbox = gr.Checkbox(visible=False)
    with gr.Tab("Download model"):
        with gr.Accordion("View public models table", open=False):
            gr.Markdown("")
            gr.Markdown("*HOW TO USE*")
            gr.Markdown(
                "- Filter voice models by selecting one or more tags and/or providing a"
                " search query.",
            )
            gr.Markdown(
                "- Select a row in the table to autofill the name and"
                " URL for the given voice model in the form fields below.",
            )
            gr.Markdown("")
            with gr.Row():
                search_query = gr.Textbox(label="Search query")
                tags = gr.CheckboxGroup(
                    value=[],
                    label="Tags",
                    choices=get_public_model_tags(),
                )
            with gr.Row():
                public_models_table = gr.Dataframe(
                    value=_filter_public_models_table,
                    inputs=[tags, search_query],
                    headers=["Name", "Description", "Tags", "Credit", "Added", "URL"],
                    label="Public models table",
                    interactive=False,
                )

        with gr.Row():
            model_url = gr.Textbox(
                label="Model URL",
                info=(
                    "Should point to a zip file containing a .pth model file and"
                    " optionally also an .index file."
                ),
            )
            model_name = gr.Textbox(
                label="Model name",
                info="Enter a unique name for the voice model.",
            )

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

        public_models_table.select(
            _autofill_model_name_and_url,
            inputs=public_models_table,
            outputs=[model_name, model_url],
            show_progress="hidden",
        )

        download_btn_click = download_btn.click(
            partial(
                exception_harness(download_model),
                progress_bar=PROGRESS_BAR,
            ),
            inputs=[model_url, model_name],
            outputs=download_msg,
        ).success(
            partial(
                render_msg,
                "[+] Succesfully downloaded voice model!",
            ),
            inputs=model_name,
            outputs=download_msg,
            show_progress="hidden",
        )

    # Upload tab
    with gr.Tab("Upload model"):
        with gr.Accordion("HOW TO USE"):
            gr.Markdown("")
            gr.Markdown(
                "1. Find the .pth file for a locally trained RVC model (e.g. in your"
                " local weights folder) and optionally also a corresponding .index file"
                " (e.g. in your logs/[name] folder)",
            )
            gr.Markdown(
                "2. Upload the files directly or save them to a folder, then compress"
                " that folder and upload the resulting .zip file",
            )
            gr.Markdown("3. Enter a unique name for the uploaded model")
            gr.Markdown("4. Click 'Upload'")

        with gr.Row():
            model_files = gr.File(
                label="Files",
                file_count="multiple",
                file_types=[".zip", ".pth", ".index"],
            )

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

        with gr.Row(equal_height=True):
            upload_btn = gr.Button("Upload", variant="primary", scale=19)
            upload_msg = gr.Textbox(
                label="Output message",
                interactive=False,
                scale=20,
            )
            upload_btn_click = upload_btn.click(
                partial(exception_harness(upload_model), progress_bar=PROGRESS_BAR),
                inputs=[model_files, local_model_name],
                outputs=upload_msg,
            ).success(
                partial(
                    render_msg,
                    "[+] Successfully uploaded voice model!",
                ),
                inputs=local_model_name,
                outputs=upload_msg,
                show_progress="hidden",
            )

    with gr.Tab("Delete models"):
        with gr.Row():
            with gr.Column():
                model_delete.render()
                delete_btn = gr.Button("Delete selected", variant="secondary")
                delete_all_btn = gr.Button("Delete all", variant="primary")
            with gr.Column():
                delete_msg = gr.Textbox(label="Output message", interactive=False)
        delete_btn_click = delete_btn.click(
            partial(confirmation_harness(delete_models), progress_bar=PROGRESS_BAR),
            inputs=[dummy_checkbox, model_delete],
            outputs=delete_msg,
            js=confirm_box_js(
                "Are you sure you want to delete the selected voice models?",
            ),
        ).success(
            partial(render_msg, "[-] Successfully deleted selected voice models!"),
            outputs=delete_msg,
            show_progress="hidden",
        )

        delete_all_btn_click = delete_all_btn.click(
            partial(
                confirmation_harness(delete_all_models),
                progress_bar=PROGRESS_BAR,
            ),
            inputs=dummy_checkbox,
            outputs=delete_msg,
            js=confirm_box_js("Are you sure you want to delete all voice models?"),
        ).success(
            partial(render_msg, "[-] Successfully deleted all voice models!"),
            outputs=delete_msg,
            show_progress="hidden",
        )
        
        
    with gr.Tab("Search models"): 
        # Textbox for user to enter search query
        query = gr.Textbox(label="Search for RVC models", placeholder="Enter your search query here")

            # Button to trigger the search
        search_button = gr.Button("Search")
    
         # Output for displaying the search results as a DataFrame
        results = gr.Dataframe(label="Search Results")
    
        # Event listener for the search button
        search_button.click(fn=search_rvc_models, inputs=query, outputs=results)

    for click_event in [
        download_btn_click,
        upload_btn_click,
        delete_btn_click,
        delete_all_btn_click,
    ]:
        click_event.success(
            partial(_update_models, 3, [], [2]),
            outputs=[model_1click, model_multi, model_delete],
            show_progress="hidden",
        )