Codecooker/rvcapi

1import json2import os3os.system("pip install torchcrepe")4os.system("pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu")5import shutil6import urllib.request7import zipfile8from argparse import ArgumentParser9 10import gradio as gr11 12from main import song_cover_pipeline13 14BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))15 16mdxnet_models_dir = os.path.join(BASE_DIR, 'mdxnet_models')17rvc_models_dir = os.path.join(BASE_DIR, 'rvc_models')18output_dir = os.path.join(BASE_DIR, 'song_output')19 20 21def get_current_models(models_dir):22    models_list = os.listdir(models_dir)23    items_to_remove = ['hubert_base.pt', 'MODELS.txt', 'public_models.json', 'rmvpe.pt']24    return [item for item in models_list if item not in items_to_remove]25 26 27def update_models_list():28    models_l = get_current_models(rvc_models_dir)29    return gr.Dropdown.update(choices=models_l)30 31 32def load_public_models():33    models_table = []34    for model in public_models['voice_models']:35        if not model['name'] in voice_models:36            model = [model['name'], model['description'], model['credit'], model['url'], ', '.join(model['tags'])]37            models_table.append(model)38 39    tags = list(public_models['tags'].keys())40    return gr.DataFrame.update(value=models_table), gr.CheckboxGroup.update(choices=tags)41 42 43def extract_zip(extraction_folder, zip_name):44    os.makedirs(extraction_folder)45    with zipfile.ZipFile(zip_name, 'r') as zip_ref:46        zip_ref.extractall(extraction_folder)47    os.remove(zip_name)48 49    index_filepath, model_filepath = None, None50    for root, dirs, files in os.walk(extraction_folder):51        for name in files:52            if name.endswith('.index'):53                index_filepath = os.path.join(root, name)54 55            if name.endswith('.pth'):56                model_filepath = os.path.join(root, name)57 58    if not model_filepath:59        raise gr.Error(f'No .pth model file was found in the extracted zip. Please check {extraction_folder}.')60 61    # move model and index file to extraction folder62    os.rename(model_filepath, os.path.join(extraction_folder, os.path.basename(model_filepath)))63    if index_filepath:64        os.rename(index_filepath, os.path.join(extraction_folder, os.path.basename(index_filepath)))65 66    # remove any unnecessary nested folders67    for filepath in os.listdir(extraction_folder):68        if os.path.isdir(os.path.join(extraction_folder, filepath)):69            shutil.rmtree(os.path.join(extraction_folder, filepath))70 71 72def download_online_model(url, dir_name, progress=gr.Progress()):73    try:74        progress(0, desc=f'[~] Downloading voice model with name {dir_name}...')75        zip_name = url.split('/')[-1]76        extraction_folder = os.path.join(rvc_models_dir, dir_name)77        if os.path.exists(extraction_folder):78            raise gr.Error(f'Voice model directory {dir_name} already exists! Choose a different name for your voice model.')79 80        if 'pixeldrain.com' in url:81            url = f'https://pixeldrain.com/api/file/{zip_name}'82 83        urllib.request.urlretrieve(url, zip_name)84 85        progress(0.5, desc='[~] Extracting zip...')86        extract_zip(extraction_folder, zip_name)87        return f'[+] {dir_name} Model successfully downloaded!'88 89    except Exception as e:90        raise gr.Error(str(e))91 92 93def upload_local_model(zip_path, dir_name, progress=gr.Progress()):94    try:95        extraction_folder = os.path.join(rvc_models_dir, dir_name)96        if os.path.exists(extraction_folder):97            raise gr.Error(f'Voice model directory {dir_name} already exists! Choose a different name for your voice model.')98 99        zip_name = zip_path.name100        progress(0.5, desc='[~] Extracting zip...')101        extract_zip(extraction_folder, zip_name)102        return f'[+] {dir_name} Model successfully uploaded!'103 104    except Exception as e:105        raise gr.Error(str(e))106 107 108def filter_models(tags, query):109    models_table = []110 111    # no filter112    if len(tags) == 0 and len(query) == 0:113        for model in public_models['voice_models']:114            models_table.append([model['name'], model['description'], model['credit'], model['url'], model['tags']])115 116    # filter based on tags and query117    elif len(tags) > 0 and len(query) > 0:118        for model in public_models['voice_models']:119            if all(tag in model['tags'] for tag in tags):120                model_attributes = f"{model['name']} {model['description']} {model['credit']} {' '.join(model['tags'])}".lower()121                if query.lower() in model_attributes:122                    models_table.append([model['name'], model['description'], model['credit'], model['url'], model['tags']])123 124    # filter based on only tags125    elif len(tags) > 0:126        for model in public_models['voice_models']:127            if all(tag in model['tags'] for tag in tags):128                models_table.append([model['name'], model['description'], model['credit'], model['url'], model['tags']])129 130    # filter based on only query131    else:132        for model in public_models['voice_models']:133            model_attributes = f"{model['name']} {model['description']} {model['credit']} {' '.join(model['tags'])}".lower()134            if query.lower() in model_attributes:135                models_table.append([model['name'], model['description'], model['credit'], model['url'], model['tags']])136 137    return gr.DataFrame.update(value=models_table)138 139 140def pub_dl_autofill(pub_models, event: gr.SelectData):141    return gr.Text.update(value=pub_models.loc[event.index[0], 'URL']), gr.Text.update(value=pub_models.loc[event.index[0], 'Model Name'])142 143 144def swap_visibility():145    return gr.update(visible=True), gr.update(visible=False), gr.update(value=''), gr.update(value=None)146 147 148def process_file_upload(file):149    return file.name, gr.update(value=file.name)150 151 152if __name__ == '__main__':153    os.system("pip install torchcrepe")154    os.system("pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu")155    parser = ArgumentParser(description='Generate a AI cover song in the song_output/id directory.', add_help=True)156    parser.add_argument("--share", action="store_true", dest="share_enabled", default=False, help="Enable sharing")157    parser.add_argument("--listen", action="store_true", default=False, help="Make the WebUI reachable from your local network.")158    parser.add_argument('--listen-host', type=str, help='The hostname that the server will use.')159    parser.add_argument('--listen-port', type=int, help='The listening port that the server will use.')160    args = parser.parse_args()161 162    voice_models = get_current_models(rvc_models_dir)163    with open(os.path.join(rvc_models_dir, 'public_models.json'), encoding='utf8') as infile:164        public_models = json.load(infile)165 166    with gr.Blocks(title='AICoverGenWebUI') as app:167 168        gr.Label('AICoverGen WebUI created with ❤️', show_label=False)169 170        # main tab171        with gr.Tab("Generate"):172 173            with gr.Accordion('Main Options'):174                with gr.Row():175                    with gr.Column():176                        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')177                        ref_btn = gr.Button('Refresh Models 🔁', variant='primary')178 179                    with gr.Column() as yt_link_col:180                        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.')181                        show_file_upload_button = gr.Button('Upload file instead')182 183                    with gr.Column(visible=False) as file_upload_col:184                        local_file = gr.File(label='Audio file')185                        song_input_file = gr.UploadButton('Upload 📂', file_types=['audio'], variant='primary')186                        show_yt_link_button = gr.Button('Paste YouTube link/Path to local file instead')187                        song_input_file.upload(process_file_upload, inputs=[song_input_file], outputs=[local_file, song_input])188 189                    pitch = gr.Slider(-24, 24, value=0, step=1, label='Pitch Change', info='Pitch Change should be set to either -12, 0, or 12 (multiples of 12) to ensure the vocals are not out of tune')190                    show_file_upload_button.click(swap_visibility, outputs=[file_upload_col, yt_link_col, song_input, local_file])191                    show_yt_link_button.click(swap_visibility, outputs=[yt_link_col, file_upload_col, song_input, local_file])192 193            with gr.Accordion('Voice conversion options', open=False):194                with gr.Row():195                    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")196                    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')197                    rms_mix_rate = gr.Slider(0, 1, value=0.25, label='RMS mix rate', info="Control how much to use the original vocal's loudness (0) or a fixed loudness (1)")198                    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.')199                keep_files = gr.Checkbox(label='Keep intermediate files',200                                         info='Keep all audio files generated in the song_output/id directory, e.g. Isolated Vocals/Instrumentals. Leave unchecked to save space')201 202            with gr.Accordion('Audio mixing options', open=False):203                gr.Markdown('### Volume Change (decibels)')204                with gr.Row():205                    main_gain = gr.Slider(-20, 20, value=0, step=1, label='Main Vocals')206                    backup_gain = gr.Slider(-20, 20, value=0, step=1, label='Backup Vocals')207                    inst_gain = gr.Slider(-20, 20, value=0, step=1, label='Music')208 209                gr.Markdown('### Reverb Control on AI Vocals')210                with gr.Row():211                    reverb_rm_size = gr.Slider(0, 1, value=0.15, label='Room size', info='The larger the room, the longer the reverb time')212                    reverb_wet = gr.Slider(0, 1, value=0.2, label='Wetness level', info='Level of AI vocals with reverb')213                    reverb_dry = gr.Slider(0, 1, value=0.8, label='Dryness level', info='Level of AI vocals without reverb')214                    reverb_damping = gr.Slider(0, 1, value=0.7, label='Damping level', info='Absorption of high frequencies in the reverb')215 216            with gr.Row():217                clear_btn = gr.ClearButton(value='Clear', components=[song_input, rvc_model, keep_files, local_file])218                generate_btn = gr.Button("Generate", variant='primary')219                ai_cover = gr.Audio(label='AI Cover', show_share_button=False)220 221            ref_btn.click(update_models_list, None, outputs=rvc_model)222            is_webui = gr.Number(value=1, visible=False)223            generate_btn.click(song_cover_pipeline,224                               inputs=[song_input, rvc_model, pitch, keep_files, is_webui, main_gain, backup_gain,225                                       inst_gain, index_rate, filter_radius, rms_mix_rate, protect, reverb_rm_size,226                                       reverb_wet, reverb_dry, reverb_damping],227                               outputs=[ai_cover])228            clear_btn.click(lambda: [0, 0, 0, 0, 0.5, 3, 0.25, 0.33, 0.15, 0.2, 0.8, 0.7, None],229                            outputs=[pitch, main_gain, backup_gain, inst_gain, index_rate, filter_radius, rms_mix_rate,230                                     protect, reverb_rm_size, reverb_wet, reverb_dry, reverb_damping, ai_cover])231 232        # Download tab233        with gr.Tab('Download model'):234 235            with gr.Tab('From HuggingFace/Pixeldrain URL'):236                with gr.Row():237                    model_zip_link = gr.Text(label='Download link to model', info='Should be a zip file containing a .pth model file and an optional .index file.')238                    model_name = gr.Text(label='Name your model', info='Give your new model a unique name from your other voice models.')239 240                with gr.Row():241                    download_btn = gr.Button('Download 🌐', variant='primary', scale=19)242                    dl_output_message = gr.Text(label='Output Message', interactive=False, scale=20)243 244                download_btn.click(download_online_model, inputs=[model_zip_link, model_name], outputs=dl_output_message)245 246                gr.Markdown('## Input Examples')247                gr.Examples(248                    [249                        ['https://huggingface.co/phant0m4r/LiSA/resolve/main/LiSA.zip', 'Lisa'],250                        ['https://pixeldrain.com/u/3tJmABXA', 'Gura'],251                        ['https://huggingface.co/Kit-Lemonfoot/kitlemonfoot_rvc_models/resolve/main/AZKi%20(Hybrid).zip', 'Azki']252                    ],253                    [model_zip_link, model_name],254                    [],255                    download_online_model,256                )257 258            with gr.Tab('From Public Index'):259 260                gr.Markdown('## How to use')261                gr.Markdown('- Click Initialize public models table')262                gr.Markdown('- Filter models using tags or search bar')263                gr.Markdown('- Select a row to autofill the download link and model name')264                gr.Markdown('- Click Download')265 266                with gr.Row():267                    pub_zip_link = gr.Text(label='Download link to model')268                    pub_model_name = gr.Text(label='Model name')269 270                with gr.Row():271                    download_pub_btn = gr.Button('Download 🌐', variant='primary', scale=19)272                    pub_dl_output_message = gr.Text(label='Output Message', interactive=False, scale=20)273 274                filter_tags = gr.CheckboxGroup(value=[], label='Show voice models with tags', choices=[])275                search_query = gr.Text(label='Search')276                load_public_models_button = gr.Button(value='Initialize public models table', variant='primary')277 278                public_models_table = gr.DataFrame(value=[], headers=['Model Name', 'Description', 'Credit', 'URL', 'Tags'], label='Available Public Models', interactive=False)279                public_models_table.select(pub_dl_autofill, inputs=[public_models_table], outputs=[pub_zip_link, pub_model_name])280                load_public_models_button.click(load_public_models, outputs=[public_models_table, filter_tags])281                search_query.change(filter_models, inputs=[filter_tags, search_query], outputs=public_models_table)282                filter_tags.change(filter_models, inputs=[filter_tags, search_query], outputs=public_models_table)283                download_pub_btn.click(download_online_model, inputs=[pub_zip_link, pub_model_name], outputs=pub_dl_output_message)284 285        # Upload tab286        with gr.Tab('Upload model'):287            gr.Markdown('## Upload locally trained RVC v2 model and index file')288            gr.Markdown('- Find model file (weights folder) and optional index file (logs/[name] folder)')289            gr.Markdown('- Compress files into zip file')290            gr.Markdown('- Upload zip file and give unique name for voice')291            gr.Markdown('- Click Upload model')292 293            with gr.Row():294                with gr.Column():295                    zip_file = gr.File(label='Zip file')296 297                local_model_name = gr.Text(label='Model name')298 299            with gr.Row():300                model_upload_button = gr.Button('Upload model', variant='primary', scale=19)301                local_upload_output_message = gr.Text(label='Output Message', interactive=False, scale=20)302                model_upload_button.click(upload_local_model, inputs=[zip_file, local_model_name], outputs=local_upload_output_message)303 304    app.launch(305        share=args.share_enabled,306        enable_queue=True,307        server_name=None if not args.listen else (args.listen_host or '0.0.0.0'),308        server_port=args.listen_port,309    )310