src/webui.py

"""
Hyper-RVC WebUI Module
Fixed: Audio loading, model dropdown, and file upload issues
"""
import json
import os
import shutil
import urllib.request
import zipfile
from argparse import ArgumentParser
import spaces
import gradio as gr
import logging

def configure_logging_libs(debug=False):
    """Configure logging levels for noisy libraries."""
    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, batch_process_files, generate_preview

BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
IS_ZERO_GPU = os.getenv("SPACES_ZERO_GPU")

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):
    """Get list of available RVC models, excluding system files."""
    if not os.path.exists(models_dir):
        logging.warning(f"Models directory not found: {models_dir}")
        return []
    
    models_list = []
    try:
        models_list = os.listdir(models_dir)
    except Exception as e:
        logging.error(f"Error reading models directory: {e}")
        return []
    
    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 and os.path.isdir(os.path.join(models_dir, item))]


def update_models_list():
    """Refresh the model dropdown choices."""
    models_l = get_current_models(rvc_models_dir)
    if not models_l:
        gr.Warning("No models found in rvc_models folder. Please download a model first.")
    return gr.update(choices=models_l)


def load_public_models():
    """Load and filter public models from JSON."""
    models_table = []
    for model in public_models['voice_models']:
        if not model['name'] in voice_models:
            model_data = [model['name'], model['description'], model['credit'], model['url'], ', '.join(model['tags'])]
            models_table.append(model_data)

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


def extract_zip(extraction_folder, zip_name):
    """Extract ZIP file and locate model/index files."""
    os.makedirs(extraction_folder, exist_ok=True)
    
    # Check if zip file exists before extracting
    if not os.path.exists(zip_name):
        raise FileNotFoundError(f"ZIP file not found: {zip_name}")
    
    with zipfile.ZipFile(zip_name, 'r') as zip_ref:
        zip_ref.extractall(extraction_folder)
    
    # Don't remove zip immediately - keep it as backup
    # os.remove(zip_name)

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

            if name.endswith('.pth') and os.path.getsize(filepath) > 1024 * 1024 * 40:
                model_filepath = filepath

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

    # Move files to extraction folder root
    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)))

    # Clean up nested folders
    for filepath in os.listdir(extraction_folder):
        full_path = os.path.join(extraction_folder, filepath)
        if os.path.isdir(full_path):
            shutil.rmtree(full_path)
    
    # Now safe to remove zip
    if os.path.exists(zip_name):
        os.remove(zip_name)


def download_online_model(url, dir_name, progress=gr.Progress()):
    """
    Download voice model from URL with progress tracking.
    """
    try:
        progress(0, desc=f'[~] Downloading voice model: {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'Model "{dir_name}" already exists! Choose a different name.')

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

        if "," in url:
            urls = [u.strip() for u in url.split(",") if u.strip()]
            os.makedirs(extraction_folder, exist_ok=True)
            for i, u in enumerate(urls):
                u = u.replace("?download=true", "")
                file_name = u.split('/')[-1]
                file_path = os.path.join(extraction_folder, file_name)
                if not os.path.exists(file_path):
                    urllib.request.urlretrieve(u, file_path)
        else:
            urllib.request.urlretrieve(url, zip_name)

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

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


def upload_local_model(zip_path, dir_name, progress=gr.Progress()):
    """
    Upload and extract local voice model.
    FIXED: Copy file to permanent location before processing
    """
    try:
        extraction_folder = os.path.join(rvc_models_dir, dir_name)
        
        if os.path.exists(extraction_folder):
            raise gr.Error(f'Model "{dir_name}" already exists! Choose a different name.')

        # FIX: Get the actual file path and copy to safe location
        zip_name = getattr(zip_path, 'name', None)
        if not zip_name or not os.path.exists(zip_name):
            raise gr.Error("Uploaded file not found or invalid. Please re-upload.")
        
        # Copy to permanent location to prevent Gradio temp file deletion
        safe_zip_path = os.path.join(output_dir, f"upload_{dir_name}.zip")
        os.makedirs(output_dir, exist_ok=True)
        shutil.copy2(zip_name, safe_zip_path)
        
        progress(0.5, desc='[~] Extracting zip...')
        extract_zip(extraction_folder, safe_zip_path)
        
        return f'[+] Model "{dir_name}" uploaded successfully!'

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


def filter_models(tags, query):
    """Filter public models by tags and/or search query."""
    models_table = []

    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']])

    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']])

    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']])

    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.update(value=models_table)


def pub_dl_autofill(pub_models, event: gr.SelectData):
    """Autofill download form when selecting a row in the table."""
    return (
        gr.update(value=pub_models.loc[event.index[0], 'URL']), 
        gr.update(value=pub_models.loc[event.index[0], 'Model Name'])
    )


def swap_visibility():
    """Toggle visibility between YouTube link and file upload."""
    return (
        gr.update(visible=True), 
        gr.update(visible=False), 
        gr.update(value=''), 
        gr.update(value=None)
    )


def process_file_upload(file):
    """Process uploaded audio file."""
    if file is None:
        return None, ""
    return file.name, file.name


def show_hop_slider(pitch_detection_algo):
    """Show/hide crepe hop length slider based on algorithm selection."""
    if pitch_detection_algo in ['mangio-crepe', 'mangio-crepe-tiny', 'fcnf0']:
        return gr.update(visible=True)
    else:
        return gr.update(visible=False)


def get_hybrid_f0_combinations():
    """Generate available hybrid F0 method combinations."""
    combinations = [
        "hybrid[pm+crepe]",
        "hybrid[pm+harvest]",
        "hybrid[harvest+crepe]",
        "hybrid[rmvpe+crepe]",
        "hybrid[pm+fcnf0]",
        "hybrid[harvest+pyin]",
        "hybrid[crepe+reaper]",
        "hybrid[pm+harvest+crepe]",
        "hybrid[rmvpe+crepe+harvest]",
        "hybrid[pm+crepe+fcnf0]",
    ]
    return combinations


def get_audios_from_folder():
    """Get list of audio files from the audios folder."""
    audios_dir = os.path.join(BASE_DIR, 'audios')
    
    if not os.path.exists(audios_dir):
        os.makedirs(audios_dir)
        return []
    
    audio_extensions = ['.wav', '.mp3', '.flac', '.ogg', '.m4a', '.aac', '.wma']
    audio_files = []
    
    try:
        for file in sorted(os.listdir(audios_dir)):
            ext = os.path.splitext(file)[1].lower()
            if ext in audio_extensions:
                full_path = os.path.join(audios_dir, file)
                # Verify file exists and is readable
                if os.path.isfile(full_path) and os.access(full_path, os.R_OK):
                    audio_files.append(full_path)
    except Exception as e:
        logging.error(f"Error listing audio files: {e}")
    
    return audio_files


def refresh_audio_list():
    """Refresh the audio files dropdown."""
    files = get_audios_from_folder()
    if not files:
        gr.Info("No audio files found in /audios folder. Add some audio files!")
    return gr.update(choices=files)


def load_selected_audio(audio_path):
    """
    Load the selected audio file from the audios folder.
    FIXED: Properly validate and return audio path
    """
    print(f"[DEBUG] load_selected_audio called with: {audio_path}")
    
    if not audio_path:
        print("[DEBUG] No audio path provided")
        return None, ""
    
    # Validate path exists
    if not os.path.exists(audio_path):
        print(f"[DEBUG] Audio path does NOT exist: {audio_path}")
        # Try to find similar file
        dirname = os.path.dirname(audio_path)
        basename = os.path.basename(audio_path)
        if os.path.exists(dirname):
            for f in os.listdir(dirname):
                if basename.lower() in f.lower():
                    found_path = os.path.join(dirname, f)
                    print(f"[DEBUG] Found alternative: {found_path}")
                    return found_path, found_path
        return None, ""
    
    # Return the validated path
    print(f"[DEBUG] Returning VALID audio path: {audio_path}")
    return audio_path, audio_path


if __name__ == '__main__':
    
    # Initialize voice models list with error handling
    voice_models = get_current_models(rvc_models_dir)
    
    if not voice_models:
        print("[WARNING] No RVC models found on startup. Users need to download models first.")
    
    # Load public models database
    with open(os.path.join(rvc_models_dir, 'public_models.json'), encoding='utf8') as infile:
        public_models = json.load(infile)

    with gr.Blocks(title='Hyper RVC - AI Voice Conversion', fill_width=True, fill_height=False) as app:

        # main 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 "Hyper RVC --> rvc_models". After new models are added into this folder, click the refresh button',
                            allow_custom_value=False  # FIX: Prevent invalid values
                        )
                        ref_btn = gr.Button('Refresh Models 🔁', variant='primary')

                    with gr.Column(visible=False) 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.')
                        show_file_upload_button = gr.Button('Upload file instead')

                    with gr.Column(visible=True) as file_upload_col:
                        # Load from audios folder - PROMINENT POSITION
                        gr.Markdown('### 📁 Load from Audios Folder')
                        with gr.Row():
                            folder_audio_dropdown = gr.Dropdown(
                                choices=get_audios_from_folder(),
                                label='Select audio from /audios folder',
                                info='Place your audio files in the "audios" folder and click Refresh',
                                interactive=True,
                                allow_custom_value=False,
                                scale=4
                            )
                            refresh_audio_btn = gr.Button('🔄 Refresh List', variant='primary', scale=1)
                        
                        load_from_folder_btn = gr.Button('▶️ LOAD SELECTED AUDIO', variant='secondary')
                        
                        local_file = gr.Audio(label='Audio file (loaded here)', interactive=True, type="filepath")
                        
                        # Hidden state to store the actual audio path from any source
                        audio_path_state = gr.Textbox(value="", label="Audio Path State", visible=False)
                        
                        # Track source of audio to prevent overwriting
                        audio_source_state = gr.Textbox(value="", label="Audio Source", visible=False)  # "folder" or "upload"
                        
                        if not IS_ZERO_GPU:
                            with gr.Row():
                                with gr.Row(scale=2):
                                    url_media_gui = gr.Textbox(value="", label="Enter URL", placeholder="www.youtube.com/watch?v=g_9rPvbENUw", lines=1)
                                with gr.Row(scale=1):
                                    url_button_gui = gr.Button("Process URL", variant="secondary")
                            url_button_gui.click(yt_download, [url_media_gui], [local_file])
                        song_input_file = gr.UploadButton('Upload 📂', file_types=['audio'], variant='primary', visible=False)
                        show_yt_link_button = gr.Button('Paste YouTube link/Path to local file instead', visible=False)
                        song_input_file.upload(process_file_upload, inputs=[song_input_file], outputs=[local_file, song_input])
                        
                        # Connect folder audio buttons - updates both audio player and state
                        refresh_audio_btn.click(refresh_audio_list, outputs=folder_audio_dropdown)
                        
                        # Button click to load selected audio - FIXED: Also set source state
                        def load_from_folder_with_source(audio_path):
                            """Load audio from folder and mark source as 'folder'"""
                            result = load_selected_audio(audio_path)
                            if result[0]:  # If audio loaded successfully
                                return result[0], result[1], "folder"
                            return result[0], result[1], ""
                        
                        load_from_folder_btn.click(
                            load_from_folder_with_source, 
                            inputs=[folder_audio_dropdown], 
                            outputs=[local_file, audio_path_state, audio_source_state]
                        )
                        
                        # Also load on dropdown selection change
                        folder_audio_dropdown.change(
                            load_from_folder_with_source, 
                            inputs=[folder_audio_dropdown], 
                            outputs=[local_file, audio_path_state, audio_source_state]
                        )
                        
                        # When local_file changes (from upload/URL), update state ONLY if source is not "folder"
                        def update_audio_state_safe(x, current_source):
                            """Update audio state only if not loaded from folder"""
                            if current_source == "folder":
                                # Don't overwrite - keep the original folder path
                                return gr.skip()  # Skip update
                            
                            # Extract path from gr.Audio output
                            if x is None:
                                return "", "upload"
                            if isinstance(x, tuple):
                                return x[0] if x[0] else "", "upload"
                            return str(x) if x else "", "upload"
                        
                        local_file.change(
                            update_audio_state_safe, 
                            inputs=[local_file, audio_source_state], 
                            outputs=[audio_path_state, audio_source_state]
                        )

                    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+', 'rmvpe', 'fcnf0', 'pyin', 'reaper', 
                             'crepe', 'crepe-tiny', 'mangio-crepe', 'mangio-crepe-tiny',
                             'harvest', 'pm', 'dio'] + get_hybrid_f0_combinations(),
                            value='rmvpe+',
                            label='Pitch detection algorithm',
                            info='rmvpe+: Best overall | fcnf0: Stable & fast | pyin: Good for speech | reaper: Noisy audio | hybrid: Combine methods for accuracy'
                        )
                        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)
                with gr.Row():
                    with gr.Row():
                        steps = gr.Slider(minimum=1, maximum=3, label="Steps", value=1, step=1, interactive=True)
                    with gr.Row():
                        extra_denoise = gr.Checkbox(True, label='Denoise', info='Apply an additional noise reduction step to clean up the audio further.')
                        keep_files = gr.Checkbox((False if IS_ZERO_GPU else True), 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', interactive=(False if IS_ZERO_GPU else True))

            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", variant='primary')
            ai_cover = gr.File(label="AI Cover", interactive=False)
                
            
            ref_btn.click(update_models_list, None, outputs=rvc_model)
            is_webui = gr.Number(value=1, visible=False)
            
            # FIXED: Enhanced validation with better error messages
            def validated_pipeline(audio_path_state_val, local_file_val, rvc_model_val, pitch_val, keep_files_val, is_webui_val,
                                   main_gain_val, backup_gain_val, inst_gain_val, index_rate_val, 
                                   filter_radius_val, rms_mix_rate_val, f0_method_val, crepe_hop_length_val,
                                   protect_val, pitch_all_val, reverb_rm_size_val, reverb_wet_val, 
                                   reverb_dry_val, reverb_damping_val, output_format_val, extra_denoise_val, steps_val):
                
                print(f"[DEBUG] === PIPELINE START ===")
                print(f"[DEBUG] audio_path_state: {audio_path_state_val}")
                print(f"[DEBUG] local_file: {local_file_val}")
                print(f"[DEBUG] voice_model: {rvc_model_val}")
                
                # Validate model selection
                if not rvc_model_val:
                    raise gr.Error("Please select a voice model first!")
                
                # Determine audio input with priority logic
                audio_input = None
                
                # PRIORITY 1: Use audio_path_state (original path from audios folder) if valid
                if audio_path_state_val:
                    if os.path.exists(audio_path_state_val):
                        audio_input = audio_path_state_val
                        print(f"[DEBUG] Using FOLDER path (valid): {audio_input}")
                    else:
                        print(f"[WARNING] Folder path invalid: {audio_path_state_val}, trying other sources...")
                
                # PRIORITY 2: Use local_file (from upload/URL)
                if not audio_input and local_file_val:
                    if isinstance(local_file_val, tuple):
                        candidate_path = local_file_val[0]
                    elif isinstance(local_file_val, str):
                        candidate_path = local_file_val.strip()
                    else:
                        candidate_path = getattr(local_file_val, 'name', None)
                    
                    if candidate_path and os.path.exists(candidate_path):
                        audio_input = candidate_path
                        print(f"[DEBUG] Using UPLOAD/URL path (valid): {audio_input}")
                    elif candidate_path:
                        print(f"[WARNING] Upload path invalid: {candidate_path}")
                
                # FINAL CHECK: No valid audio found
                if not audio_input:
                    raise gr.Error("""
                    ❌ **No valid audio file found!**
                    
                    Please either:
                    1. Select an audio from the **Audios Folder** dropdown and click **LOAD**
                    2. **Upload** an audio file using the upload button
                    3. Paste a **YouTube URL** and process it
                    
                    Make sure the file actually exists!
                    """)
                
                # Final validation with detailed error message
                if not os.path.exists(audio_input):
                    raise gr.Error(f"""
                    ❌ **Audio file not found:** `{audio_input}`
                    
                    The file may have been deleted or moved. Please reload your audio.
                    """)
                
                print(f"[DEBUG] Final audio input: {audio_input}")
                print(f"[DEBUG] File size: {os.path.getsize(audio_input)} bytes")
                
                return song_cover_pipeline(
                    audio_input, rvc_model_val, pitch_val, keep_files_val, is_webui_val,
                    main_gain_val, backup_gain_val, inst_gain_val, index_rate_val, filter_radius_val,
                    rms_mix_rate_val, f0_method_val, crepe_hop_length_val, protect_val, pitch_all_val,
                    reverb_rm_size_val, reverb_wet_val, reverb_dry_val, reverb_damping_val,
                    output_format_val, extra_denoise_val, steps_val
                )
            
            generate_btn.click(validated_pipeline,
                               inputs=[audio_path_state, local_file, 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],
                               outputs=[ai_cover])
            
            # Fixed clear button to also reset audio_source_state
            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, "", ""],
                            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, 
                                     audio_path_state, audio_source_state])

        # Download 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/phantom4r/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 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)

        # Batch Processing Tab
        with gr.Tab('Batch Processing'):
            gr.Markdown('## Process Multiple Audio Files at Once')
            gr.Markdown('- Upload multiple audio files or select a folder')
            gr.Markdown('- Apply same voice conversion settings to all files')
            gr.Markdown('- Results will be saved to song_output/batch_output/')
            
            with gr.Row():
                with gr.Column():
                    batch_input_files = gr.File(
                        label='Audio files (multiple)', 
                        file_count='multiple',
                        file_types=['audio', '.wav', '.mp3', '.flac', '.ogg']
                    )
                    
                with gr.Column():
                    batch_model = gr.Dropdown(voice_models, label='Voice Model')
                    batch_pitch = gr.Slider(-12, 12, value=0, step=1, label='Pitch Change (semitones)')
            
            with gr.Row():
                with gr.Column():
                    batch_f0_method = gr.Dropdown(
                        ['rmvpe+', 'rmvpe', 'fcnf0', 'pyin', 'reaper'] + get_hybrid_f0_combinations(),
                        value='rmvpe+',
                        label='Pitch detection algorithm'
                    )
                    batch_index_rate = gr.Slider(0, 1, value=0.5, label='Index Rate')
                
                with gr.Column():
                    batch_protect = gr.Slider(0, 0.5, value=0.33, label='Protect rate')
                    batch_format = gr.Dropdown(['wav', 'mp3'], value='mp3', label='Output format')
            
            with gr.Row():
                batch_process_btn = gr.Button('Start Batch Processing 🚀', variant='primary')
                batch_progress = gr.Textbox(label='Progress', interactive=False)
                batch_results = gr.Files(label='Output files')

            # Audio Preview Section
            gr.Markdown('---')
            gr.Markdown('## 🔊 Audio Preview')
            gr.Markdown('Preview a short clip of your voice conversion before processing the full audio.')
            
            with gr.Row():
                with gr.Column():
                    preview_audio = gr.Audio(label='Input audio for preview', type='filepath')
                    preview_duration = gr.Slider(5, 30, value=10, step=5, label='Preview duration (seconds)')
                
                with gr.Column():
                    preview_btn = gr.Button('Generate Preview ▶️', variant='secondary')
                    preview_output = gr.Audio(label='Preview output', interactive=False)

            # Connect batch processing button
            batch_process_btn.click(
                batch_process_files,
                inputs=[batch_input_files, batch_model, batch_pitch, batch_f0_method, 
                        batch_index_rate, batch_protect, batch_format],
                outputs=[batch_progress, batch_results]
            )

            # Connect preview button
            preview_btn.click(
                generate_preview,
                inputs=[preview_audio, preview_duration],
                outputs=preview_output
            )

    app.launch()