src/webui.py

import json
import os
import shutil
import urllib.request
import zipfile
from argparse import ArgumentParser
import gradio as gr
import logging

def configure_logging_libs(debug=False):
    modules = [
        "numba",
        "httpx",
        "markdown_it",
        "fairseq",
        "faiss",
    ]
    try:
        for module in modules:
            logging.getLogger(module).setLevel(logging.WARNING)
        os.environ['TF_CPP_MIN_LOG_LEVEL'] = "3" if not debug else "1"
    except Exception as error:
        pass

configure_logging_libs()

from main import song_cover_pipeline, yt_download
from main import get_history_manager, get_preset_manager, get_rvc_cache, _HAS_MODULES

from ui.config import (
    BASE_DIR, IS_ZERO_GPU, rvc_models_dir, AUDIO_FOLDERS_DIR,
    AUDIO_EXTENSIONS, SUPPORTED_AUDIO_FORMATS, ROOT_FOLDER_LABEL,
)
from ui.handlers import (
    get_current_models, update_models_list, load_public_models,
    extract_zip, download_online_model, upload_local_model,
    filter_models, pub_dl_autofill, show_hop_slider,
    get_audio_files_from_folder, folder_name_to_path, get_audio_folders,
    refresh_audio_folders, select_audio_folder, set_selected_audio,
    sync_file_to_path, prepare_song_input,
    create_audio_folder, upload_to_audio_folder,
    hide_audio_folder_selector, show_audio_folder_selector, show_folder_stats,
    preset_load_selected, preset_refresh_list, preset_save_current,
    preset_delete_selected, preset_reset_builtin, preset_describe,
    history_refresh, history_play_selected, history_clear, history_get_stats,
    cache_stats_refresh, cache_clear, show_audio_info, show_disk_usage,
    set_globals as set_handler_globals,
)

# MAIN APPLICATION
# ============================================

def build_app(args):
    """Construct the Gradio Blocks application. Returns the app object (not yet launched)."""
    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)

    # Set globals so handler functions in ui.handlers can access them
    set_handler_globals(public_models, voice_models)

    with gr.Blocks(title='AICoverGen WebUI', fill_width=True, fill_height=False) as app:

        # No top-level markdown headers — keep it clean.

        # ============================================
        # GENERATE TAB
        # ============================================
        with gr.Tab("Generate"):

            # --- Voice Model & Pitch (no accordion, always visible) ---
            with gr.Row():
                _default_model = 'Silver' if 'Silver' in voice_models else (voice_models[0] if voice_models else None)
                rvc_model = gr.Dropdown(voice_models, value=_default_model, label='Voice Models', scale=3)
                ref_btn = gr.Button('Refresh Models', variant='primary', scale=1)
            with gr.Row():
                pitch = gr.Slider(-3, 3, value=0, step=1, label='Pitch Change (Vocals ONLY)')
                pitch_all = gr.Slider(-12, 12, value=0, step=1, label='Overall Pitch Change')

            # --- Input source ---
            with gr.Accordion('Input Source', open=True):
                # Hidden textbox that reliably carries the song path to the pipeline.
                # gr.File doesn't pass arbitrary filesystem paths correctly when set
                # programmatically (e.g. from the audio folder selector), so we use
                # this textbox as the actual input to song_cover_pipeline.
                song_input_path = gr.Textbox(visible=False, value="")

                with gr.Row():
                    with gr.Column(visible=True) as file_upload_col:
                        audio_file_types = [ext for ext in AUDIO_EXTENSIONS if ext.startswith('.')]
                        local_file = gr.File(label='Audio file', interactive=True, type="filepath", file_types=audio_file_types, height=150)
                        audio_info_box = gr.Textbox(label="File info", interactive=False, lines=5, max_lines=6)
                        # Sync the uploaded file path into the hidden textbox
                        local_file.change(show_audio_info, inputs=[local_file], outputs=[audio_info_box])
                        local_file.change(sync_file_to_path, inputs=[local_file], outputs=[song_input_path])

                        if not IS_ZERO_GPU:
                            with gr.Row():
                                url_media_gui = gr.Textbox(value="", label="Or enter YouTube URL", placeholder="https://www.youtube.com/watch?v=...", lines=1, scale=3)
                                url_button_gui = gr.Button("Download from URL", variant="secondary", scale=1)
                            url_button_gui.click(yt_download, [url_media_gui], [local_file])

                        show_folder_selector_btn = gr.Button('Select from Audio Folders', variant='secondary')

                    with gr.Column(visible=False) as audio_folder_col:
                        with gr.Row():
                            _init_folders = get_audio_folders()
                            _init_value = _init_folders[0] if _init_folders else None
                            folder_dropdown = gr.Dropdown(
                                choices=_init_folders,
                                value=_init_value,
                                label='Select Audio Folder',
                                interactive=True,
                                scale=3,
                            )
                            refresh_folders_btn = gr.Button('Refresh', variant='secondary', scale=1)

                        audio_files_list = gr.Dropdown([], label='Select Audio File', interactive=True, multiselect=False)
                        folder_status = gr.Textbox(label='Status', interactive=False, value="Select a folder to see available audio files")
                        with gr.Row():
                            select_audio_btn = gr.Button('Select Audio File', variant='primary')
                            back_to_upload_btn = gr.Button('Back to File Upload', variant='secondary')

                        refresh_folders_btn.click(refresh_audio_folders, outputs=folder_dropdown)
                        folder_dropdown.change(select_audio_folder, inputs=folder_dropdown, outputs=[audio_files_list, folder_status])
                        # set_selected_audio returns TWO values: gr.update for local_file, path string for song_input_path
                        select_audio_btn.click(
                            set_selected_audio,
                            inputs=[audio_files_list, folder_dropdown],
                            outputs=[local_file, song_input_path],
                        ).then(
                            show_audio_info,
                            inputs=[song_input_path],
                            outputs=[audio_info_box],
                        )
                        show_folder_selector_btn.click(
                            hide_audio_folder_selector,
                            outputs=[file_upload_col, audio_folder_col],
                        ).then(
                            select_audio_folder,
                            inputs=folder_dropdown,
                            outputs=[audio_files_list, folder_status],
                        )
                        back_to_upload_btn.click(show_audio_folder_selector, outputs=[file_upload_col, audio_folder_col])

            # --- All AI options in one accordion ---
            with gr.Accordion('AI Options', open=False):

                # 1. Voice conversion options
                with gr.Accordion('Voice conversion options', open=False):
                    with gr.Row():
                        index_rate = gr.Slider(0, 1, value=0.5, label='Index Rate')
                        filter_radius = gr.Slider(0, 7, value=3, step=1, label='Filter radius')
                        rms_mix_rate = gr.Slider(0, 1, value=0.25, label='RMS mix rate')
                        protect = gr.Slider(0, 0.5, value=0.33, label='Protect rate')
                        with gr.Column():
                            f0_method = gr.Dropdown(['rmvpe+', 'rmvpe', 'mangio-crepe'], value='rmvpe+', label='Pitch detection algorithm')
                            crepe_hop_length = gr.Slider(32, 320, value=128, step=1, visible=False, label='Crepe hop length')
                            f0_method.change(show_hop_slider, inputs=f0_method, outputs=crepe_hop_length)

                    with gr.Row():
                        steps = gr.Slider(minimum=1, maximum=3, label="Steps", value=1, step=1, interactive=True)
                        extra_denoise = gr.Checkbox(True, label='Denoise')
                        keep_files = gr.Checkbox((False if IS_ZERO_GPU else True), label='Keep intermediate files', interactive=(False if IS_ZERO_GPU else True))

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

                    with gr.Row():
                        reverb_rm_size = gr.Slider(0, 1, value=0.15, label='Reverb Room size')
                        reverb_wet = gr.Slider(0, 1, value=0.2, label='Reverb Wetness')
                        reverb_dry = gr.Slider(0, 1, value=0.8, label='Reverb Dryness')
                        reverb_damping = gr.Slider(0, 1, value=0.7, label='Reverb Damping')

                    output_format = gr.Dropdown(['mp3', 'wav'], value='mp3', label='Output file type')

                # 3. Advanced mixing options
                with gr.Accordion('Advanced mixing options', open=False):
                    with gr.Row():
                        soundgoodizer = gr.Slider(0, 1, value=0, step=0.05, label='Soundgoodizer (0=off, 1=max)')

                    with gr.Row():
                        highpass_freq = gr.Slider(0, 500, value=80, step=10, label='High-pass filter (Hz)')
                        lowpass_freq = gr.Slider(0, 20000, value=0, step=100, label='Low-pass filter (Hz, 0=off)')

                    with gr.Row():
                        low_shelf_gain = gr.Slider(-15, 15, value=0, step=0.5, label='Low shelf gain (dB)')
                        low_shelf_freq = gr.Slider(50, 500, value=200, step=10, label='Low shelf freq (Hz)')
                        mid_peak_gain = gr.Slider(-15, 15, value=0, step=0.5, label='Mid peak gain (dB)')
                        mid_peak_freq = gr.Slider(200, 5000, value=1000, step=50, label='Mid peak freq (Hz)')
                        mid_peak_q = gr.Slider(0.1, 5, value=1.0, step=0.1, label='Mid peak Q')
                        high_shelf_gain = gr.Slider(-15, 15, value=0, step=0.5, label='High shelf gain (dB)')
                        high_shelf_freq = gr.Slider(2000, 12000, value=5000, step=100, label='High shelf freq (Hz)')

                    with gr.Row():
                        noise_gate_threshold = gr.Slider(-100, 0, value=-55, step=1, label='Noise gate threshold (dB)')
                        noise_gate_ratio = gr.Slider(1, 20, value=10, step=0.5, label='Noise gate ratio')
                        compressor_threshold = gr.Slider(-60, 0, value=-15, step=1, label='Compressor threshold (dB)')
                        compressor_ratio = gr.Slider(1, 20, value=4, step=0.5, label='Compressor ratio')
                        limiter_ceiling = gr.Slider(-10, 0, value=-1, step=0.1, label='Limiter ceiling (dB)')
                        final_gain = gr.Slider(-12, 12, value=0, step=0.5, label='Final gain (dB)')

                    with gr.Row():
                        distortion_drive = gr.Slider(0, 30, value=0, step=0.5, label='Distortion drive (dB, 0=off)')
                        bitcrush_bits = gr.Slider(0, 16, value=0, step=0.5, label='Bitcrush bits (0=off, 16=none)')
                        clipping_threshold = gr.Slider(-24, 0, value=0, step=0.5, label='Clipping threshold (dB, 0=off)')

                    with gr.Row():
                        chorus_mix = gr.Slider(0, 1, value=0, step=0.05, label='Chorus mix (0=off)')
                        chorus_depth = gr.Slider(0, 1, value=0.25, step=0.05, label='Chorus depth')
                        chorus_rate = gr.Slider(0, 10, value=0.5, step=0.1, label='Chorus rate (Hz)')

                    with gr.Row():
                        phaser_mix = gr.Slider(0, 1, value=0, step=0.05, label='Phaser mix (0=off)')
                        phaser_depth = gr.Slider(0, 1, value=0.5, step=0.05, label='Phaser depth')
                        phaser_rate = gr.Slider(0, 10, value=0.5, step=0.1, label='Phaser rate (Hz)')

                    with gr.Row():
                        delay_seconds = gr.Slider(0, 2, value=0, step=0.05, label='Delay time (s, 0=off)')
                        delay_feedback = gr.Slider(0, 0.95, value=0.3, step=0.05, label='Delay feedback')
                        delay_mix = gr.Slider(0, 1, value=0, step=0.05, label='Delay mix (0=off)')

                # 4. Backing vocal inference
                with gr.Accordion('Backing Vocal Inference', open=False):
                    backup_vocal_infer = gr.Checkbox(False, label='Infer backing vocals with AI voice')
                    backup_vocal_pitch = gr.Slider(-12, 12, value=0, step=1, label='Backing vocal harmony pitch (semitones, 0=unison, +3/-3=harmony)')

                # 5. Presets
                with gr.Accordion('Presets', open=False):
                    if _HAS_MODULES and get_preset_manager() is not None:
                        preset_names_init = get_preset_manager().list_names()
                    else:
                        preset_names_init = []
                    with gr.Row():
                        preset_dropdown = gr.Dropdown(preset_names_init, label='Load Preset', interactive=True, scale=3)
                        preset_load_btn = gr.Button('Load', variant='primary', scale=1)
                        preset_refresh_btn = gr.Button('Refresh', variant='secondary', scale=1)
                    preset_description_box = gr.Textbox(label='Preset description', interactive=False, lines=2)
                    preset_dropdown.change(preset_describe, inputs=[preset_dropdown], outputs=[preset_description_box])
                    preset_refresh_btn.click(preset_refresh_list, outputs=[preset_dropdown, preset_description_box])
                    preset_load_btn.click(
                        preset_load_selected,
                        inputs=[preset_dropdown],
                        outputs=[
                            rvc_model, pitch, pitch_all, index_rate, filter_radius,
                            rms_mix_rate, protect, f0_method, crepe_hop_length, steps,
                            main_gain, backup_gain, inst_gain, reverb_rm_size, reverb_wet,
                            reverb_dry, reverb_damping, output_format, extra_denoise,
                            preset_description_box,
                        ],
                    )

            # --- Generate button + output ---
            with gr.Row():
                clear_btn = gr.ClearButton(value='Clear', components=[local_file, rvc_model])
                generate_btn = gr.Button("Generate", variant='primary')
                # Always use gr.Audio so the output is playable inline.
                # type='filepath' lets us receive the file path as a string.
                ai_cover = gr.Audio(label='AI Cover', type='filepath')

            ref_btn.click(update_models_list, None, outputs=rvc_model)
            is_webui = gr.Number(value=1, visible=False)
            # Generate button: first resolve the song path via prepare_song_input,
            # then call the pipeline with the resolved path.
            # prepare_song_input raises gr.Error if inputs are missing, which
            # stops the chain and shows a clean error to the user.
            # Store the pipeline chain so we can add a .then() auto-refresh
            # of the history tab after all tabs are defined.
            _pipeline_chain = generate_btn.click(
                prepare_song_input,
                inputs=[local_file, song_input_path, rvc_model],
                outputs=[song_input_path],
            ).then(
                song_cover_pipeline,
                inputs=[song_input_path, 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, extra_denoise, steps,
                        highpass_freq, lowpass_freq,
                        low_shelf_gain, low_shelf_freq,
                        mid_peak_gain, mid_peak_freq, mid_peak_q,
                        high_shelf_gain, high_shelf_freq,
                        noise_gate_threshold, noise_gate_ratio,
                        compressor_threshold, compressor_ratio,
                        limiter_ceiling, final_gain,
                        soundgoodizer,
                        distortion_drive,
                        bitcrush_bits,
                        chorus_depth, chorus_rate, chorus_mix,
                        delay_seconds, delay_feedback, delay_mix,
                        phaser_depth, phaser_rate, phaser_mix,
                        clipping_threshold,
                        backup_vocal_infer, backup_vocal_pitch],
                outputs=[ai_cover],
            )
            # Clear button resets all controls including advanced mixing + backing vocals
            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, True, 1, "",
                         80, 0, 0, 200, 0, 1000, 1.0, 0, 5000, -55, 10, -15, 4, -1, 0,
                         0, 0, 0, 0.25, 0.5, 0, 0, 0.3, 0, 0.5, 0.5, 0, 0,
                         False, 0],
                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, extra_denoise, steps, song_input_path,
                         highpass_freq, lowpass_freq, low_shelf_gain, low_shelf_freq,
                         mid_peak_gain, mid_peak_freq, mid_peak_q, high_shelf_gain, high_shelf_freq,
                         noise_gate_threshold, noise_gate_ratio, compressor_threshold, compressor_ratio,
                         limiter_ceiling, final_gain,
                         soundgoodizer,
                         distortion_drive,
                         bitcrush_bits,
                         chorus_depth, chorus_rate, chorus_mix,
                         delay_seconds, delay_feedback, delay_mix,
                         phaser_depth, phaser_rate, phaser_mix,
                         clipping_threshold,
                         backup_vocal_infer, backup_vocal_pitch],
            )

        # ============================================
        # DOWNLOAD MODEL TAB
        # ============================================
        with gr.Tab('Download model'):

            with gr.Tab('From HuggingFace/Pixeldrain URL'):
                with gr.Row():
                    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.')
                    model_name = gr.Text(label='Name your model', info='Give your new model a unique name from your other voice models.')

                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('## Input Examples')
                gr.Examples(
                    [
                        ['https://huggingface.co/MrDawg/ToothBrushing/resolve/main/ToothBrushing.zip?download=true', 'ToothBrushing'],
                        ['https://huggingface.co/sail-rvc/Aldeano_Minecraft__RVC_V2_-_500_Epochs_/resolve/main/model.pth?download=true, https://huggingface.co/sail-rvc/Aldeano_Minecraft__RVC_V2_-_500_Epochs_/resolve/main/model.index?download=true', 'Minecraft_Villager'],
                        ['https://huggingface.co/phant0m4r/LiSA/resolve/main/LiSA.zip', 'Lisa'],
                        ['https://pixeldrain.com/u/3tJmABXA', 'Gura'],
                        ['https://huggingface.co/Kit-Lemonfoot/kitlemonfoot_rvc_models/resolve/main/AZKi%20(Hybrid).zip', 'Azki']
                    ],
                    [model_zip_link, model_name],
                    [],
                    download_online_model,
                    cache_examples=False,
                )

            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 MODEL 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)

        # ============================================
        # PRESET MANAGER TAB (NEW)
        # ============================================
        with gr.Tab('Preset Manager'):
            gr.Markdown("## Preset Manager")
            gr.Markdown("Save your current settings (voice model + pitch + reverb + gains + etc.) as a named preset, then load it later in one click from the Generate tab.")

            if _HAS_MODULES and get_preset_manager() is not None:
                pm_init_names = get_preset_manager().list_names()
            else:
                pm_init_names = []

            with gr.Accordion("Load / Delete Preset", open=True):
                with gr.Row():
                    pm_dropdown = gr.Dropdown(pm_init_names, label='Presets', interactive=True, scale=3)
                    pm_refresh_btn = gr.Button('Refresh List', variant='secondary', scale=1)
                    pm_delete_btn = gr.Button('Delete Selected', variant='stop', scale=1)
                    pm_reset_btn = gr.Button('Restore Built-ins', variant='secondary', scale=1)
                pm_description = gr.Textbox(label='Description', interactive=False, lines=3)
                pm_status = gr.Textbox(label='Status', interactive=False, lines=1)
                pm_dropdown.change(preset_describe, inputs=[pm_dropdown], outputs=[pm_description])
                pm_refresh_btn.click(preset_refresh_list, outputs=[pm_dropdown, pm_description])
                pm_delete_btn.click(preset_delete_selected, inputs=[pm_dropdown], outputs=[pm_status, pm_dropdown])
                pm_reset_btn.click(preset_reset_builtin, outputs=[pm_status, pm_dropdown])

            with gr.Accordion("Save Current Settings as Preset", open=False):
                gr.Markdown("Fill in the values below to save a new preset. (You can also use the Load Preset button in the Generate tab after creating it here.)")
                save_name = gr.Textbox(label="Preset name", placeholder="e.g. My Smooth Female Voice")
                save_desc = gr.Textbox(label="Description", lines=2, placeholder="What does this preset do?")
                with gr.Row():
                    save_voice_model = gr.Textbox(label="Voice model name (must match a folder in rvc_models)")
                with gr.Row():
                    save_pitch = gr.Slider(-3, 3, value=0, step=1, label='Pitch (vocals)')
                    save_pitch_all = gr.Slider(-12, 12, value=0, step=1, label='Pitch (all)')
                with gr.Row():
                    save_index_rate = gr.Slider(0, 1, value=0.5, label='Index rate')
                    save_filter_radius = gr.Slider(0, 7, value=3, step=1, label='Filter radius')
                    save_rms = gr.Slider(0, 1, value=0.25, label='RMS mix rate')
                    save_protect = gr.Slider(0, 0.5, value=0.33, label='Protect')
                with gr.Row():
                    save_f0 = gr.Dropdown(['rmvpe+', 'rmvpe', 'mangio-crepe'], value='rmvpe+', label='F0 method')
                    save_crepe_hop = gr.Slider(32, 320, value=128, step=1, label='Crepe hop')
                    save_steps = gr.Slider(1, 3, value=1, step=1, label='Steps')
                with gr.Row():
                    save_main_gain = gr.Slider(-20, 20, value=0, step=1, label='Main gain (dB)')
                    save_backup_gain = gr.Slider(-20, 20, value=0, step=1, label='Backup gain (dB)')
                    save_inst_gain = gr.Slider(-20, 20, value=0, step=1, label='Inst gain (dB)')
                with gr.Row():
                    save_rev_size = gr.Slider(0, 1, value=0.15, label='Reverb room size')
                    save_rev_wet = gr.Slider(0, 1, value=0.2, label='Reverb wet')
                    save_rev_dry = gr.Slider(0, 1, value=0.8, label='Reverb dry')
                    save_rev_damp = gr.Slider(0, 1, value=0.7, label='Reverb damping')
                with gr.Row():
                    save_output_format = gr.Dropdown(['mp3', 'wav'], value='mp3', label='Output format')
                    save_extra_denoise = gr.Checkbox(True, label='Extra denoise')
                save_btn = gr.Button('Save Preset', variant='primary')
                save_status = gr.Textbox(label='Status', interactive=False, lines=1)
                save_btn.click(
                    preset_save_current,
                    inputs=[
                        save_name, save_desc, save_voice_model, save_pitch, save_pitch_all,
                        save_index_rate, save_filter_radius, save_rms, save_protect,
                        save_f0, save_crepe_hop, save_steps,
                        save_main_gain, save_backup_gain, save_inst_gain,
                        save_rev_size, save_rev_wet, save_rev_dry, save_rev_damp,
                        save_output_format, save_extra_denoise,
                    ],
                    outputs=[save_status, pm_dropdown],
                )

        # ============================================
        # HISTORY TAB (NEW)
        # ============================================
        with gr.Tab('History'):
            with gr.Row():
                history_refresh_btn = gr.Button('Refresh', variant='primary')
                history_stats_btn = gr.Button('Show Stats', variant='secondary')
                history_clear_btn = gr.Button('Clear History', variant='stop')
                history_clear_files_cb = gr.Checkbox(False, label='Also delete cover files')

            history_summary = gr.Textbox(label='Summary', interactive=False, lines=1)
            history_table = gr.DataFrame(
                value=[],
                headers=['When', 'Song', 'Voice Model', 'Pitch', 'Pitch All', 'Duration (s)', 'Format', 'File Exists', 'Path'],
                label='History (newest first) - click a row to play it',
                interactive=False,
                wrap=True,
            )
            history_player = gr.Audio(label='Selected cover (click a row above to play)', interactive=False, type='filepath')
            history_status = gr.Textbox(label='Status', interactive=False, lines=2)

            # history_refresh now returns 3 values: (table, summary, status)
            history_refresh_btn.click(
                history_refresh,
                outputs=[history_table, history_summary, history_status],
            )
            history_stats_btn.click(history_get_stats, outputs=[history_status])
            # history_clear now returns 4 values: (status, table, summary, player)
            history_clear_btn.click(
                history_clear,
                inputs=[history_clear_files_cb],
                outputs=[history_status, history_table, history_summary, history_player],
            )
            # Click a row to play it
            history_table.select(
                history_play_selected,
                inputs=[history_table],
                outputs=[history_player],
            )

        # ============================================
        # SYSTEM / CACHE STATS TAB (NEW)
        # ============================================
        with gr.Tab('System'):
            gr.Markdown("## System & Cache")
            gr.Markdown("Monitor model cache hit rate and clear it to free memory.")

            with gr.Row():
                cache_refresh_btn = gr.Button('Refresh Cache Stats', variant='primary')
                cache_clear_btn = gr.Button('Clear Cache', variant='stop')

            cache_stats_box = gr.Textbox(label='Model Cache', interactive=False, lines=12, max_lines=15, value="Click Refresh to view.")

            cache_refresh_btn.click(cache_stats_refresh, outputs=[cache_stats_box])
            cache_clear_btn.click(cache_clear, outputs=[cache_stats_box])

            gr.Markdown("---")
            gr.Markdown("### Disk usage")
            disk_usage_box = gr.Textbox(label='Disk usage', interactive=False, lines=8, value="Click button below to compute.")
            disk_refresh_btn = gr.Button('Compute Disk Usage', variant='secondary')
            disk_refresh_btn.click(show_disk_usage, outputs=[disk_usage_box])

        # ============================================
        # AUDIO FOLDERS MANAGEMENT TAB
        # ============================================
        with gr.Tab('Audio Folders'):
            gr.Markdown("## Manage Audio Folders")
            gr.Markdown(f"Audio library location: `{AUDIO_FOLDERS_DIR}`")
            gr.Markdown(f"The first entry **{ROOT_FOLDER_LABEL}** represents files placed directly in this folder. You can also create subfolders to organize your library.")

            with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown("### Current Audio Folders")
                    _af_init_folders = get_audio_folders()
                    _af_init_value = _af_init_folders[0] if _af_init_folders else None
                    current_folders = gr.Dropdown(
                        choices=_af_init_folders,
                        value=_af_init_value,
                        label='Available Folders',
                        interactive=True,
                    )
                    refresh_folders_display = gr.Button('Refresh Folders List', variant='secondary')
                    refresh_folders_display.click(refresh_audio_folders, outputs=current_folders)

                    gr.Markdown("### Create New Folder")
                    new_folder_name = gr.Textbox(label='New folder name', placeholder='e.g. My Songs')
                    create_folder_btn = gr.Button('Create Folder', variant='primary')
                    create_folder_status = gr.Textbox(label='Status', interactive=False, lines=1)
                    create_folder_btn.click(
                        create_audio_folder,
                        inputs=[new_folder_name],
                        outputs=[create_folder_status, current_folders],
                    )

                    gr.Markdown("### Upload Audio to Selected Folder")
                    upload_target_folder = gr.Dropdown(
                        choices=_af_init_folders,
                        value=_af_init_value,
                        label='Upload to folder',
                        interactive=True,
                    )
                    upload_files_input = gr.File(
                        label='Select audio file(s) to upload',
                        file_count='multiple',
                        file_types=[ext.lstrip('.') for ext in AUDIO_EXTENSIONS if ext.startswith('.')],
                    )
                    upload_files_btn = gr.Button('Upload to Folder', variant='primary')
                    upload_status = gr.Textbox(label='Upload status', interactive=False, lines=2)
                    upload_files_btn.click(
                        upload_to_audio_folder,
                        inputs=[upload_files_input, upload_target_folder],
                        outputs=[upload_status, upload_target_folder],
                    ).then(
                        # Also refresh the main folder list after upload
                        refresh_audio_folders,
                        outputs=current_folders,
                    )

                    gr.Markdown("### File Statistics")
                    folder_stats = gr.Textbox(label='Statistics', interactive=False, lines=20, max_lines=25)
                    get_stats_btn = gr.Button('Show Statistics', variant='secondary')
                    get_stats_btn.click(show_folder_stats, outputs=folder_stats)

                with gr.Column(scale=2):
                    gr.Markdown("### How to Use Audio Folders")
                    gr.Markdown(f"""
                    1. **Quick start:** Upload audio directly via the **Upload Audio to Selected Folder** section on the left (the **{ROOT_FOLDER_LABEL}** entry puts files at the root).
                    2. (Optional) Create subfolders to organize songs by artist, genre, etc.
                    3. Go to the **Generate** tab.
                    4. Click **"Select from Audio Folders"**.
                    5. Choose a folder from the dropdown (the audio files list will refresh automatically).
                    6. Pick an audio file from the second dropdown.
                    7. Click **"Select Audio File"** to load it into the input.
                    8. Proceed with generating your AI cover.
                    """)

                    gr.Markdown("### Supported Audio Formats")
                    gr.Markdown(f"**Common:** MP3, MP2, WAV, FLAC, M4A, AAC, OGG, WMA, ALAC, AIFF, OPUS, AMR  \n**Additional:** 3GP, 3G2, AC3, EAC3, APE, AU, SND, CAF, DTS, FLV, GSM, MKA, MLP, MPC, PCM, QCP, RA, RM, RAM, SHN, SPX, TTA, VOC, VOX, WV")

                    gr.Markdown("### Tips")
                    gr.Markdown("""
                    - Files uploaded to **Main Library** go directly into the audio_folders root.
                    - Subfolders inside your audio folders are scanned recursively.
                    - The audio files load into the same input as file uploads on the Generate tab.
                    - Convert unsupported formats to MP3, WAV, or FLAC for best compatibility.
                    - Use Refresh after manually adding files via the filesystem.
                    """)

        # ============================================
        # AUTO-REFRESH HISTORY AFTER GENERATION
        # ============================================
        # After a cover is generated, auto-refresh the history tab so the
        # new entry appears without the user needing to click Refresh.
        # Chained off the pipeline .then() so it runs only on success.
        _pipeline_chain.then(
            history_refresh,
            outputs=[history_table, history_summary, history_status],
        )

    # ============================================
    # RETURN APP (caller is responsible for launching)
    # ============================================
    return app


# ============================================
# ENTRY POINT
# ============================================

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("--builtin-player",  action="store_true", default=False, help="Use the builtin audio player")
    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.')
    parser.add_argument('--theme', type=str, default="NoCrypt/miku", help='Set the theme (default: NoCrypt/miku)')
    parser.add_argument("--ssr", action="store_true", help="Enable SSR (Server-Side Rendering)")
    args = parser.parse_args()

    app = build_app(args)
    app.launch(
        share=args.share_enabled,
        debug=args.share_enabled,
        show_error=True,
        server_name=None if not args.listen else (args.listen_host or '0.0.0.0'),
        server_port=args.listen_port,
        ssr_mode=args.ssr,
    )