src/webui.py

import json
import os
import shutil
import urllib.request
import zipfile
from argparse import ArgumentParser

import gradio as gr

from main import song_cover_pipeline

BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

mdxnet_models_dir = os.path.join(BASE_DIR, 'mdxnet_models')
rvc_models_dir = os.path.join(BASE_DIR, 'rvc_models')
output_dir = os.path.join(BASE_DIR, 'song_output')


def get_current_models(models_dir):
    models_list = os.listdir(models_dir)
    items_to_remove = ['hubert_base.pt', 'MODELS.txt', 'public_models.json', 'rmvpe.pt']
    return [item for item in models_list if item not in items_to_remove]


def update_models_list():
    models_l = get_current_models(rvc_models_dir)
    return gr.Dropdown.update(choices=models_l)


def load_public_models():
    models_table = []
    for model in public_models['voice_models']:
        if not model['name'] in voice_models:
            model = [model['name'], model['description'], model['credit'], model['url'], ', '.join(model['tags'])]
            models_table.append(model)

    tags = list(public_models['tags'].keys())
    return gr.DataFrame.update(value=models_table), gr.CheckboxGroup.update(choices=tags)


def extract_zip(extraction_folder, zip_name):
    os.makedirs(extraction_folder)
    with zipfile.ZipFile(zip_name, 'r') as zip_ref:
        zip_ref.extractall(extraction_folder)
    os.remove(zip_name)

    index_filepath, model_filepath = None, None
    for root, dirs, files in os.walk(extraction_folder):
        for name in files:
            if name.endswith('.index') and os.stat(os.path.join(root, name)).st_size > 1024 * 100:
                index_filepath = os.path.join(root, name)

            if name.endswith('.pth') and os.stat(os.path.join(root, name)).st_size > 1024 * 1024 * 40:
                model_filepath = os.path.join(root, name)

    if not model_filepath:
        raise gr.Error(f'No .pth model file was found in the extracted zip. Please check {extraction_folder}.')

    # move model and index file to extraction folder
    os.rename(model_filepath, os.path.join(extraction_folder, os.path.basename(model_filepath)))
    if index_filepath:
        os.rename(index_filepath, os.path.join(extraction_folder, os.path.basename(index_filepath)))

    # remove any unnecessary nested folders
    for filepath in os.listdir(extraction_folder):
        if os.path.isdir(os.path.join(extraction_folder, filepath)):
            shutil.rmtree(os.path.join(extraction_folder, filepath))


def download_online_model(url, dir_name, progress=gr.Progress()):
    try:
        progress(0, desc=f'[~] Downloading voice model with name {dir_name}...')
        zip_name = url.split('/')[-1]
        extraction_folder = os.path.join(rvc_models_dir, dir_name)
        if os.path.exists(extraction_folder):
            raise gr.Error(f'Voice model directory {dir_name} already exists! Choose a different name for your voice model.')

        if 'pixeldrain.com' in url:
            url = f'https://pixeldrain.com/api/file/{zip_name}'

        urllib.request.urlretrieve(url, zip_name)

        progress(0.5, desc='[~] Extracting zip...')
        extract_zip(extraction_folder, zip_name)
        return f'[+] {dir_name} Model successfully downloaded!'

    except Exception as e:
        raise gr.Error(str(e))


def upload_local_model(zip_path, dir_name, progress=gr.Progress()):
    try:
        extraction_folder = os.path.join(rvc_models_dir, dir_name)
        if os.path.exists(extraction_folder):
            raise gr.Error(f'Voice model directory {dir_name} already exists! Choose a different name for your voice model.')

        zip_name = zip_path.name
        progress(0.5, desc='[~] Extracting zip...')
        extract_zip(extraction_folder, zip_name)
        return f'[+] {dir_name} Model successfully uploaded!'

    except Exception as e:
        raise gr.Error(str(e))

def_link = "https://youtu.be/hT_nvWreIhg"

def filter_models(tags, query):
    models_table = []

    # no filter
    if len(tags) == 0 and len(query) == 0:
        for model in public_models['voice_models']:
            models_table.append([model['name'], model['description'], model['credit'], model['url'], model['tags']])

    # filter based on tags and query
    elif len(tags) > 0 and len(query) > 0:
        for model in public_models['voice_models']:
            if all(tag in model['tags'] for tag in tags):
                model_attributes = f"{model['name']} {model['description']} {model['credit']} {' '.join(model['tags'])}".lower()
                if query.lower() in model_attributes:
                    models_table.append([model['name'], model['description'], model['credit'], model['url'], model['tags']])

    # filter based on only tags
    elif len(tags) > 0:
        for model in public_models['voice_models']:
            if all(tag in model['tags'] for tag in tags):
                models_table.append([model['name'], model['description'], model['credit'], model['url'], model['tags']])

    # filter based on only query
    else:
        for model in public_models['voice_models']:
            model_attributes = f"{model['name']} {model['description']} {model['credit']} {' '.join(model['tags'])}".lower()
            if query.lower() in model_attributes:
                models_table.append([model['name'], model['description'], model['credit'], model['url'], model['tags']])

    return gr.DataFrame.update(value=models_table)


def pub_dl_autofill(pub_models, event: gr.SelectData):
    return gr.Text.update(value=pub_models.loc[event.index[0], 'URL']), gr.Text.update(value=pub_models.loc[event.index[0], 'Model Name'])


def swap_visibility():
    return gr.update(visible=True), gr.update(visible=False), gr.update(value=''), gr.update(value=None)


def process_file_upload(file):
    return file.name, gr.update(value=file.name)


def show_hop_slider(pitch_detection_algo):
    if pitch_detection_algo == 'mangio-crepe':
        return gr.update(visible=True)
    else:
        return gr.update(visible=False)


if __name__ == '__main__':
    parser = ArgumentParser(description='Generate a AI cover song in the song_output/id directory.', add_help=True)
    parser.add_argument("--share", action="store_true", dest="share_enabled", default=False, help="Enable sharing")
    parser.add_argument("--listen", action="store_true", default=False, help="Make the WebUI reachable from your local network.")
    parser.add_argument('--listen-host', type=str, help='The hostname that the server will use.')
    parser.add_argument('--listen-port', type=int, help='The listening port that the server will use.')
    args = parser.parse_args()

    voice_models = get_current_models(rvc_models_dir)
    with open(os.path.join(rvc_models_dir, 'public_models.json'), encoding='utf8') as infile:
        public_models = json.load(infile)

    with gr.Blocks(title='oItsMineZ\'s AI Cover WebUI', theme=gr.themes.Base(font=[gr.themes.GoogleFont("Noto Sans Thai"), "sans-serif"])) as app:

        gr.Label('oItsMineZ\'s RVC v2 AI Cover WebUI', show_label=False)
        gr.Markdown(
            "<div align='center'>\n\n"+
            "RVC v2 Model"+
            "[![oItsMineZ's RVC Model](https://img.shields.io/badge/%F0%9F%A4%97_Hugging_Face-_oItsMineZ's%20RVC%20%20Model-yellow?style=for-the-badge&logoColor=yellow)](https://huggingface.co/oItsMineZ/oItsMineZ-RVC-Model)\n\n"+
            "ติดตาม oItsMineZ"+
            "[![oItsMineZ on YouTube](https://img.shields.io/badge/YouTube-FF0000?style=for-the-badge&logo=youtube&logoColor=white)](https://www.youtube.com/@oItsMineZ?sub_confirmation=1)"+
            "</div>"
        )
        
        # Main Tab
        with gr.Tab("📢 Info"):
            
            gr.Markdown("## 📌แนะนำให้โคลน Space นี้ไว้ในบัญชีของคุณ เพื่อการใช้งานที่ดียิ่งขึ้น (ต้องสมัครบัญชี Hugging Face ก่อน)")
            gr.Markdown("[![Duplicate this Space](https://huggingface.co/datasets/huggingface/badges/raw/main/duplicate-this-space-sm-dark.svg)](https://huggingface.co/spaces/oItsMineZ/RVC-v2-AI-Cover-WebUI?duplicate=true)\n\n") 
            
            gr.HTML("<b><h2> 📄ข้อควรรู้ </h2></b>")
            gr.Markdown("- RVC v2 (Retrieval Based Voice Conversion v2) เป็น AI Voice Model ที่ปรับปรุงมาจาก VITS ที่ทำให้เทรนโมเดลได้ง่ายขึ้น และคุณภาพของเสียงดีขึ้น")
            gr.Markdown("- WebUI นี้ใช้สำหรับเฉพาะ **AI Cover** เพลงเท่านั้น! ถ้าอยากใช้เฉพาะเสียงพูดให้ใช้ [**ตัวนี้แทน**](https://huggingface.co/spaces/oItsMineZ/RVC-v2-WebUI)")
            gr.Markdown("- บางเพลงอาจ**ใช้เวลานานมากๆ** ขึ้นอยู่กับความยาวหรือขนาดไฟล์เพลง (จากที่ผมลองเพลงของ [**OneRepublic - Counting Stars**](https://youtu.be/hT_nvWreIhg) ใช้เวลา 55 นาที 😱)")
            gr.Markdown("- ถ้าใช้ในโทรศัพท์ **ห้าม**ออกจากหน้า Web ขณะเว็บกำลังดำเนินการอยู่ เพราะทำให้ไฟล์หายระหว่างขั้นตอนได้")
            gr.Markdown("- ถ้าคุณพร้อมที่จะทำเพลง AI Cover แล้ว ให้คลิกแท็บ 🎵 Generate ได้เลย!")

            gr.HTML("<b><h2> ✨ฟีเจอร์ </h2></b>")
            gr.Markdown("- นำเพลงจาก YouTube มา Cover ได้ทันที เพียงแค่ก็อบลิงก์มาวาง")
            gr.Markdown("- อัปโหลดไฟล์เพลงของตัวเองได้เลย โดยไม่ต้องลบทำนองออกก่อน")
            gr.Markdown("- มี UVR5 (Ultimate Vocal Remover v5) โดยช่วยแยกเสียงร้องกับทำนองออกจากเพลง")
            gr.Markdown("- สามารถดาวน์โหลด Model อื่นๆ ได้ที่แท็บ ⬇️ Download Model [**(เว็บสำหรับหา Model เพิ่มเติม)**](https://voice-models.com)")
            gr.Markdown("- ที่สำคัญ **อย่าลืม** *Refresh Model* ทุกครั้งเมื่อโหลด Model ใหม่เข้ามา")

            gr.HTML("<b><h2> 📋รายชื่อ Model </h2></b>")
            gr.Markdown("- อาจารย์แดง (DaengGuitar) - 500 Epochs")
            gr.Markdown("- เต้ (TAEEXZENFIRE) - 500 Epochs")
            gr.Markdown("- ท่านศาสดา - 50 Epochs")
            gr.Markdown("- Model ใหม่เร็วๆ นี้ 🤫")

            gr.HTML("<b><h2> 🌐WebUI อื่นๆ </h2></b>")
            gr.Markdown("- เฉพาะ Vocal (เสียงพูดปกติ)")
            gr.Markdown("[![Hugging Face Spaces](https://img.shields.io/badge/%F0%9F%8E%99%EF%B8%8F_Space-_oItsMineZ's%20RVC%20v2%20WebUI-orange?style=for-the-badge)](https://huggingface.co/spaces/oItsMineZ/RVC-v2-WebUI)")

            gr.HTML("<b><h2> 📱เวอร์ชั่นอื่นๆ </h2></b>")
            gr.Markdown("- Google Colab (No WebUI)")
            gr.Markdown("[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/oItsMineZ/RVC-v2-AICover-Colab/blob/main/oItsMineZ-rvc-v2-AICover-Colab.ipynb)")

            gr.HTML("<b><h2> ❤️ขอขอบคุณ </h2></b>")
            gr.Markdown("- [**@SociallyIneptWeeb**](https://github.com/SociallyIneptWeeb) for [***AICoverGen***](https://github.com/SociallyIneptWeeb/AICoverGen)")
            
        # Generate Tab
        with gr.Tab("🎵 Generate"):

            with gr.Accordion('Main Options'):
                with gr.Row():
                    with gr.Column():
                        rvc_model = gr.Dropdown(voice_models, label='Voice Models', info='Models folder "AICoverGen --> rvc_models". After new models are added into this folder, click the refresh button')
                        ref_btn = gr.Button('Refresh Models 🔁', variant='primary')

                    with gr.Column() as yt_link_col:
                        song_input = gr.Text(label='Song input', info='Link to a song on YouTube or full path to a local file. For file upload, click the button below. Example: https://youtu.be/dQw4w9WgXcQ', value=def_link)
                        show_file_upload_button = gr.Button('Upload file instead 📤')

                    with gr.Column(visible=False) as file_upload_col:
                        local_file = gr.File(label='Audio file')
                        song_input_file = gr.UploadButton('Upload 📂', file_types=['audio'], variant='primary')
                        show_yt_link_button = gr.Button('Paste YouTube link/Path to local file instead 🔗')
                        song_input_file.upload(process_file_upload, inputs=[song_input_file], outputs=[local_file, song_input])

                    with gr.Column():
                        pitch = gr.Slider(-3, 3, value=0, step=1, label='Pitch Change (Vocals ONLY)', info='Generally, use 1 for male to female conversions and -1 for vice-versa. (Octaves)')
                        pitch_all = gr.Slider(-12, 12, value=0, step=1, label='Overall Pitch Change', info='Changes pitch/key of vocals and instrumentals together. Altering this slightly reduces sound quality. (Semitones)')
                    show_file_upload_button.click(swap_visibility, outputs=[file_upload_col, yt_link_col, song_input, local_file])
                    show_yt_link_button.click(swap_visibility, outputs=[yt_link_col, file_upload_col, song_input, local_file])

            with gr.Accordion('Voice conversion options', open=False):
                with gr.Row():
                    index_rate = gr.Slider(0, 1, value=0.5, label='Index Rate', info="Controls how much of the AI voice's accent to keep in the vocals")
                    filter_radius = gr.Slider(0, 7, value=3, step=1, label='Filter radius', info='If >=3: apply median filtering median filtering to the harvested pitch results. Can reduce breathiness')
                    rms_mix_rate = gr.Slider(0, 1, value=0.25, label='RMS mix rate', info="Control how much to mimic the original vocal's loudness (0) or a fixed loudness (1)")
                    protect = gr.Slider(0, 0.5, value=0.33, label='Protect rate', info='Protect voiceless consonants and breath sounds. Set to 0.5 to disable.')
                    with gr.Column():
                        f0_method = gr.Dropdown(['rmvpe', 'mangio-crepe'], value='rmvpe', label='Pitch detection algorithm', info='Best option is rmvpe (clarity in vocals), then mangio-crepe (smoother vocals)')
                        crepe_hop_length = gr.Slider(32, 320, value=128, step=1, visible=False, label='Crepe hop length', info='Lower values leads to longer conversions and higher risk of voice cracks, but better pitch accuracy.')
                        f0_method.change(show_hop_slider, inputs=f0_method, outputs=crepe_hop_length)
                keep_files = gr.Checkbox(label='Keep intermediate files', info='Keep all audio files generated in the song_output/id directory, e.g. Isolated Vocals/Instrumentals. Leave unchecked to save space')

            with gr.Accordion('Audio mixing options', open=False):
                gr.Markdown('### Volume Change (decibels)')
                with gr.Row():
                    main_gain = gr.Slider(-20, 20, value=0, step=1, label='Main Vocals')
                    backup_gain = gr.Slider(-20, 20, value=0, step=1, label='Backup Vocals')
                    inst_gain = gr.Slider(-20, 20, value=0, step=1, label='Music')

                gr.Markdown('### Reverb Control on AI Vocals')
                with gr.Row():
                    reverb_rm_size = gr.Slider(0, 1, value=0.15, label='Room size', info='The larger the room, the longer the reverb time')
                    reverb_wet = gr.Slider(0, 1, value=0.2, label='Wetness level', info='Level of AI vocals with reverb')
                    reverb_dry = gr.Slider(0, 1, value=0.8, label='Dryness level', info='Level of AI vocals without reverb')
                    reverb_damping = gr.Slider(0, 1, value=0.7, label='Damping level', info='Absorption of high frequencies in the reverb')

                gr.Markdown('### Audio Output Format')
                output_format = gr.Dropdown(['mp3', 'wav'], value='mp3', label='Output file type', info='mp3: small file size, decent quality. wav: Large file size, best quality')

            with gr.Row():
                clear_btn = gr.ClearButton(value='Clear', components=[song_input, rvc_model, keep_files, local_file])
                generate_btn = gr.Button("✨ Generate Song ✨", variant='primary')
                ai_cover = gr.Audio(label='AI Cover', show_share_button=False)

            ref_btn.click(update_models_list, None, outputs=rvc_model)
            is_webui = gr.Number(value=1, visible=False)
            generate_btn.click(song_cover_pipeline,
                               inputs=[song_input, rvc_model, pitch, keep_files, is_webui, main_gain, backup_gain,
                                       inst_gain, index_rate, filter_radius, rms_mix_rate, f0_method, crepe_hop_length,
                                       protect, pitch_all, reverb_rm_size, reverb_wet, reverb_dry, reverb_damping,
                                       output_format],
                               outputs=[ai_cover])
            clear_btn.click(lambda: [0, 0, 0, 0, 0.5, 3, 0.25, 0.33, 'rmvpe', 128, 0, 0.15, 0.2, 0.8, 0.7, 'mp3', None],
                            outputs=[pitch, main_gain, backup_gain, inst_gain, index_rate, filter_radius, rms_mix_rate,
                                     protect, f0_method, crepe_hop_length, pitch_all, reverb_rm_size, reverb_wet,
                                     reverb_dry, reverb_damping, output_format, ai_cover])

        # Download Tab
        with gr.Tab('⬇️ Download Model'):

            with gr.Tab('From HuggingFace/Pixeldrain URL'):
                with gr.Row():
                    model_name = gr.Text(label='Model Name', info='Give your new model a unique name from your other voice models.')
                    model_zip_link = gr.Text(label='Link', info='Should be a zip file containing a .pth model file and an optional .index file.')

                with gr.Row():
                    download_btn = gr.Button('Download 🌐', variant='primary', scale=19)
                    dl_output_message = gr.Text(label='Output Message', interactive=False, scale=20)

                download_btn.click(download_online_model, inputs=[model_zip_link, model_name], outputs=dl_output_message)

                gr.Markdown('## oItsMineZ\'s Model List')
                gr.Examples(
                    [
                        ['DaengGuitar', 'https://huggingface.co/oItsMineZ/oItsMineZ-RVC-Model/resolve/main/DaengGuitar/DaengGuitar.zip'],
                        ['TAEEXZENFIRE', 'https://huggingface.co/oItsMineZ/oItsMineZ-RVC-Model/resolve/main/TAEEXZENFIRE/TAEEXZENFIRE.zip'],
                        ['ท่านศาสดา', 'https://huggingface.co/oItsMineZ/oItsMineZ-RVC-Model/resolve/main/Sadsada/Sadsada.zip']
                    ],
                    [model_name, model_zip_link],
                    [],
                    download_online_model,
                )

            with gr.Tab('From Public Index'):

                gr.Markdown('## How to use')
                gr.Markdown('- Click Initialize public models table')
                gr.Markdown('- Filter models using tags or search bar')
                gr.Markdown('- Select a row to autofill the download link and model name')
                gr.Markdown('- Click Download')

                with gr.Row():
                    pub_zip_link = gr.Text(label='Download link to model')
                    pub_model_name = gr.Text(label='Model name')

                with gr.Row():
                    download_pub_btn = gr.Button('Download 🌐', variant='primary', scale=19)
                    pub_dl_output_message = gr.Text(label='Output Message', interactive=False, scale=20)

                filter_tags = gr.CheckboxGroup(value=[], label='Show voice models with tags', choices=[])
                search_query = gr.Text(label='Search')
                load_public_models_button = gr.Button(value='Initialize public models table', variant='primary')

                public_models_table = gr.DataFrame(value=[], headers=['Model Name', 'Description', 'Credit', 'URL', 'Tags'], label='Available Public Models', interactive=False)
                public_models_table.select(pub_dl_autofill, inputs=[public_models_table], outputs=[pub_zip_link, pub_model_name])
                load_public_models_button.click(load_public_models, outputs=[public_models_table, filter_tags])
                search_query.change(filter_models, inputs=[filter_tags, search_query], outputs=public_models_table)
                filter_tags.change(filter_models, inputs=[filter_tags, search_query], outputs=public_models_table)
                download_pub_btn.click(download_online_model, inputs=[pub_zip_link, pub_model_name], outputs=pub_dl_output_message)

        # Upload tab
        with gr.Tab('⬆️ Upload Model'):
            gr.Markdown('## Upload locally trained RVC v2 model and index file')
            gr.Markdown('- Find model file (weights folder) and optional index file (logs/[name] folder)')
            gr.Markdown('- Compress files into zip file')
            gr.Markdown('- Upload zip file and give unique name for voice')
            gr.Markdown('- Click Upload model')

            with gr.Row():
                with gr.Column():
                    zip_file = gr.File(label='Zip file')

                local_model_name = gr.Text(label='Model name')

            with gr.Row():
                model_upload_button = gr.Button('Upload model', variant='primary', scale=19)
                local_upload_output_message = gr.Text(label='Output Message', interactive=False, scale=20)
                model_upload_button.click(upload_local_model, inputs=[zip_file, local_model_name], outputs=local_upload_output_message)

    app.launch(
        share=args.share_enabled,
        enable_queue=True,
        server_name=None if not args.listen else (args.listen_host or '0.0.0.0'),
        server_port=args.listen_port,
    )