src/main.py

import spaces
import torch
import argparse
import gc
import hashlib
import json
import os
import shlex
import subprocess
from contextlib import suppress
from urllib.parse import urlparse, parse_qs
import time
import shutil
import logging

import gradio as gr
import librosa
import numpy as np
import soundfile as sf
import sox
import yt_dlp
from pedalboard import (
    Pedalboard, Reverb, Compressor, HighpassFilter, LowpassFilter,
    NoiseGate, Limiter, Gain, HighShelfFilter, LowShelfFilter, PeakFilter,
    Distortion, Bitcrush, Chorus, Delay, Phaser, Clipping,
)
from pedalboard.io import AudioFile
from pydub import AudioSegment
import noisereduce as nr

from uvr import run_mdx
from rvc import Config, load_hubert, get_vc, rvc_infer

# Optional: import improved modules (caching, history, presets)
try:
    from modules import ModelCache, HistoryManager, PresetManager
    _HAS_MODULES = True
except Exception as _e:
    print(f"[main] improved modules not available: {_e}")
    _HAS_MODULES = False

logging.getLogger("httpx").setLevel(logging.WARNING)

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

# Directories for improved features
data_dir = os.path.join(BASE_DIR, 'data')
os.makedirs(data_dir, exist_ok=True)

# Singletons (lazy-init)
_history_manager = None
_preset_manager = None


def get_history_manager():
    global _history_manager
    if _history_manager is None and _HAS_MODULES:
        _history_manager = HistoryManager(os.path.join(data_dir, 'history.json'))
    return _history_manager


def get_preset_manager():
    global _preset_manager
    if _preset_manager is None and _HAS_MODULES:
        _preset_manager = PresetManager(os.path.join(data_dir, 'presets.json'))
    return _preset_manager


# Model cache (only used when not on ZeroGPU, since ZeroGPU reloads per-call)
_rvc_cache = None
def get_rvc_cache():
    global _rvc_cache
    if _rvc_cache is None and _HAS_MODULES and not IS_ZERO_GPU:
        _rvc_cache = ModelCache(max_entries=3, max_age_seconds=1200)
    return _rvc_cache


def clean_old_folders(base_path: str, max_age_seconds: int = 10800):
    if not os.path.isdir(base_path):
        print(f"Error: {base_path} is not a valid directory.")
        return

    now = time.time()

    for folder_name in os.listdir(base_path):
        folder_path = os.path.join(base_path, folder_name)
        if os.path.isdir(folder_path):
            last_modified = os.path.getmtime(folder_path)
            if now - last_modified > max_age_seconds:
                shutil.rmtree(folder_path)


def get_youtube_video_id(url, ignore_playlist=True):
    """
    Examples:
    http://youtu.be/SA2iWivDJiE
    http://www.youtube.com/watch?v=_oPAwA_Udwc&feature=feedu
    http://www.youtube.com/embed/SA2iWivDJiE
    http://www.youtube.com/v/SA2iWivDJiE?version=3&hl=en_US
    """
    if "m.youtube.com" in url:
        url = url.replace("m.youtube.com", "www.youtube.com")
    query = urlparse(url)
    if query.hostname == 'youtu.be':
        if query.path[1:] == 'watch':
            return query.query[2:]
        return query.path[1:]

    if query.hostname in {'www.youtube.com', 'youtube.com', 'music.youtube.com'}:
        if not ignore_playlist:
            # use case: get playlist id not current video in playlist
            with suppress(KeyError):
                return parse_qs(query.query)['list'][0]
        if query.path == '/watch':
            return parse_qs(query.query)['v'][0]
        if query.path[:7] == '/watch/':
            return query.path.split('/')[1]
        if query.path[:7] == '/embed/':
            return query.path.split('/')[2]
        if query.path[:3] == '/v/':
            return query.path.split('/')[2]

    # returns None for invalid YouTube url
    return None


def yt_download(link):
    if not link.strip():
        gr.Info("You need to provide a download link.")
        return None
    ydl_opts = {
        'format': 'bestaudio',
        'outtmpl': '%(title)s',
        'nocheckcertificate': True,
        'ignoreerrors': True,
        'no_warnings': True,
        'quiet': True,
        'extractaudio': True,
        'postprocessors': [{'key': 'FFmpegExtractAudio', 'preferredcodec': 'mp3'}],
    }
    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
        result = ydl.extract_info(link, download=True)
        download_path = ydl.prepare_filename(result, outtmpl='%(title)s.mp3')

    return download_path


def raise_exception(error_msg, is_webui):
    if is_webui:
        raise gr.Error(error_msg)
    else:
        raise Exception(error_msg)


def get_rvc_model(voice_model, is_webui):
    rvc_model_filename, rvc_index_filename = None, None
    model_dir = os.path.join(rvc_models_dir, voice_model)
    for file in os.listdir(model_dir):
        if os.path.isdir(file):
            for ff in os.listdir(file):
                ext = os.path.splitext(ff)[1]
                if ext == '.pth':
                    rvc_model_filename = ff
                if ext == '.index':
                    rvc_index_filename = ff
        ext = os.path.splitext(file)[1]
        if ext == '.pth':
            rvc_model_filename = file
        if ext == '.index':
            rvc_index_filename = file

    if rvc_model_filename is None:
        error_msg = f'No model file exists in {model_dir}.'
        raise_exception(error_msg, is_webui)

    return os.path.join(model_dir, rvc_model_filename), os.path.join(model_dir, rvc_index_filename) if rvc_index_filename else ''


def get_audio_paths(song_dir):
    orig_song_path = None
    instrumentals_path = None
    main_vocals_dereverb_path = None
    backup_vocals_path = None

    for file in os.listdir(song_dir):
        if file.endswith('_Instrumental.wav'):
            instrumentals_path = os.path.join(song_dir, file)
            orig_song_path = instrumentals_path.replace('_Instrumental', '')

        elif file.endswith('_Vocals_Main_DeReverb.wav'):
            main_vocals_dereverb_path = os.path.join(song_dir, file)

        elif file.endswith('_Vocals_Backup.wav'):
            backup_vocals_path = os.path.join(song_dir, file)

    return orig_song_path, instrumentals_path, main_vocals_dereverb_path, backup_vocals_path


def get_audio_with_suffix(song_dir, suffix="_mysuffix.wav"):
    target_path = None

    for file in os.listdir(song_dir):
        if file.endswith(suffix):
            target_path = os.path.join(song_dir, file)
            break

    return target_path


def convert_to_stereo(audio_path):
    # Use ffmpeg to probe + convert. This is more robust than librosa
    # because it doesn't depend on audioread backends.
    # First, convert to a standard stereo 44100Hz WAV regardless of input format.
    # This sidesteps mono/stereo detection issues and ensures the pipeline
    # always gets a clean stereo WAV.
    stereo_path = f'{os.path.splitext(audio_path)[0]}_stereo.wav'

    # Check if already a stereo WAV at 44100Hz — if so, skip conversion
    try:
        info = sf.info(audio_path)
        if info.channels == 2 and info.samplerate == 44100:
            return audio_path
    except Exception:
        pass

    # Use ffmpeg to convert to stereo WAV
    command = shlex.split(f'ffmpeg -y -loglevel error -i "{audio_path}" -ac 2 -ar 44100 -f wav "{stereo_path}"')
    result = subprocess.run(command, capture_output=True, text=True)
    if result.returncode != 0 or not os.path.exists(stereo_path):
        # Fallback: try librosa
        try:
            wave, sr = librosa.load(audio_path, mono=False, sr=44100)
            if type(wave[0]) != np.ndarray:
                # mono — duplicate to stereo
                wave = np.stack([wave, wave])
            sf.write(stereo_path, wave.T, 44100)
            return stereo_path
        except Exception as lib_err:
            raise Exception(
                f"Cannot convert audio to stereo. ffmpeg error: {result.stderr}. "
                f"librosa error: {lib_err}. "
                f"File: {audio_path}"
            )
    return stereo_path


def pitch_shift(audio_path, pitch_change):
    output_path = f'{os.path.splitext(audio_path)[0]}_p{pitch_change}.wav'
    if not os.path.exists(output_path):
        y, sr = sf.read(audio_path)
        tfm = sox.Transformer()
        tfm.pitch(pitch_change)
        y_shifted = tfm.build_array(input_array=y, sample_rate_in=sr)
        sf.write(output_path, y_shifted, sr)

    return output_path


def get_hash(filepath):
    with open(filepath, 'rb') as f:
        file_hash = hashlib.blake2b()
        while chunk := f.read(8192):
            file_hash.update(chunk)

    return file_hash.hexdigest()[:11]


def display_progress(message, percent, is_webui, progress=None):
    """Update the Gradio progress bar.

    Uses a smooth 0.0-1.0 scale with descriptive messages.
    Falls back to print() for CLI mode.
    """
    if is_webui and progress is not None:
        try:
            progress(float(percent), desc=message)
        except Exception:
            pass
    else:
        print(f"[{percent*100:.0f}%] {message}")


def preprocess_song(song_input, mdx_model_params, song_id, is_webui, input_type, progress, keep_orig, orig_song_path):

    song_output_dir = os.path.join(output_dir, song_id)

    display_progress('Separating vocals from instrumental...', 0.10, is_webui, progress)
    vocals_path, instrumentals_path = run_mdx(mdx_model_params, song_output_dir, os.path.join(mdxnet_models_dir, 'UVR-MDX-NET-Voc_FT.onnx'), orig_song_path, denoise=True, keep_orig=keep_orig)

    display_progress('Separating main vocals from backup vocals...', 0.20, is_webui, progress)
    backup_vocals_path, main_vocals_path = run_mdx(mdx_model_params, song_output_dir, os.path.join(mdxnet_models_dir, 'UVR_MDXNET_KARA_2.onnx'), vocals_path, suffix='Backup', invert_suffix='Main', denoise=True)

    display_progress('Removing reverb from vocals...', 0.30, is_webui, progress)
    _, main_vocals_dereverb_path = run_mdx(mdx_model_params, song_output_dir, os.path.join(mdxnet_models_dir, 'Reverb_HQ_By_FoxJoy.onnx'), main_vocals_path, invert_suffix='DeReverb', exclude_main=True, denoise=True)

    return orig_song_path, vocals_path, instrumentals_path, main_vocals_path, backup_vocals_path, main_vocals_dereverb_path


def get_audio_file(song_input, is_webui, input_type, progress):
    keep_orig = False
    if input_type == 'yt':
        display_progress('Downloading song from YouTube...', 0.02, is_webui, progress)
        song_link = song_input.split('&')[0]
        orig_song_path = yt_download(song_link)
    elif input_type == 'local':
        orig_song_path = song_input
        keep_orig = True
    else:
        orig_song_path = None
    return keep_orig, orig_song_path

device = "cuda:0" if torch.cuda.is_available() else "cpu"
compute_half = True if torch.cuda.is_available() else False
config = Config(device, compute_half)
hubert_model = load_hubert("cuda", config.is_half, os.path.join(rvc_models_dir, 'hubert_base.pt'))
print(device, "half>>", config.is_half)


def voice_change(voice_model, vocals_path, output_path, pitch_change, f0_method, index_rate, filter_radius, rms_mix_rate, protect, crepe_hop_length, is_webui, steps):
    rvc_model_path, rvc_index_path = get_rvc_model(voice_model, is_webui)

    device = "cuda:0" if torch.cuda.is_available() else "cpu"
    compute_half = True if torch.cuda.is_available() else False
    config = Config(device, compute_half)

    # Try to use cached model (avoids re-loading the .pth on every call).
    # We cache only (cpt, version, net_g, tgt_sr, vc) keyed by model path + half precision flag.
    cache = get_rvc_cache()
    cache_key = f"vc::{rvc_model_path}::{config.is_half}"
    cached = cache.get(cache_key) if cache else None
    cache_hit = False
    if cached is not None:
        cpt, version, net_g, tgt_sr, vc = cached
        cache_hit = True
    else:
        cpt, version, net_g, tgt_sr, vc = get_vc(device, config.is_half, config, rvc_model_path)
        if cache is not None:
            cache.put(cache_key, (cpt, version, net_g, tgt_sr, vc))

    # convert main vocals
    global hubert_model
    rvc_infer(rvc_index_path, index_rate, vocals_path, output_path, pitch_change, f0_method, cpt, version, net_g, filter_radius, tgt_sr, rms_mix_rate, protect, crepe_hop_length, vc, hubert_model, steps)

    # Do NOT delete cpt here when cached - the cache owns it.
    # When not cached, free it now to release memory.
    if not cache_hit:
        del cpt
    gc.collect()


def add_audio_effects(audio_path, reverb_rm_size, reverb_wet, reverb_dry, reverb_damping,
                      highpass_freq=80.0, lowpass_freq=0.0,
                      low_shelf_gain=0.0, low_shelf_freq=200.0,
                      mid_peak_gain=0.0, mid_peak_freq=1000.0, mid_peak_q=1.0,
                      high_shelf_gain=0.0, high_shelf_freq=5000.0,
                      noise_gate_threshold=-55.0, noise_gate_ratio=10.0,
                      compressor_threshold=-15.0, compressor_ratio=4.0,
                      limiter_ceiling=-1.0, final_gain=0.0,
                      soundgoodizer=0.0,
                      distortion_drive=0.0,
                      bitcrush_bits=0.0,
                      chorus_depth=0.0, chorus_rate=0.0, chorus_mix=0.0,
                      delay_seconds=0.0, delay_feedback=0.0, delay_mix=0.0,
                      phaser_depth=0.0, phaser_rate=0.0, phaser_mix=0.0,
                      clipping_threshold=0.0):
    """Apply audio effects chain to vocals.

    Extended version with full EQ, dynamics, soundgoodizer, saturation,
    modulation (chorus/phaser), delay, and bitcrush controls.
    All gain/freq params have safe defaults that match the original behavior.
    """
    output_path = f'{os.path.splitext(audio_path)[0]}_mixed.wav'

    # Build the pedalboard chain
    chain = []

    # High-pass filter (remove rumble) - only if freq > 0
    if highpass_freq and float(highpass_freq) > 0:
        chain.append(HighpassFilter(cutoff_frequency_hz=float(highpass_freq)))

    # Low-pass filter (remove harshness) - only if freq > 0
    if lowpass_freq and float(lowpass_freq) > 0:
        chain.append(LowpassFilter(cutoff_frequency_hz=float(lowpass_freq)))

    # Low shelf EQ (bass)
    if low_shelf_gain and float(low_shelf_gain) != 0:
        chain.append(LowShelfFilter(cutoff_frequency_hz=float(low_shelf_freq), gain_db=float(low_shelf_gain)))

    # Mid peak EQ
    if mid_peak_gain and float(mid_peak_gain) != 0:
        chain.append(PeakFilter(cutoff_frequency_hz=float(mid_peak_freq), gain_db=float(mid_peak_gain), q=float(mid_peak_q)))

    # High shelf EQ (treble)
    if high_shelf_gain and float(high_shelf_gain) != 0:
        chain.append(HighShelfFilter(cutoff_frequency_hz=float(high_shelf_freq), gain_db=float(high_shelf_gain)))

    # Noise gate
    if noise_gate_threshold and float(noise_gate_threshold) < 0:
        chain.append(NoiseGate(threshold_db=float(noise_gate_threshold), ratio=float(noise_gate_ratio), attack_ms=5.0, release_ms=150.0))

    # Compressor
    if compressor_ratio and float(compressor_ratio) > 1:
        chain.append(Compressor(threshold_db=float(compressor_threshold), ratio=float(compressor_ratio), attack_ms=5.0, release_ms=100.0))

    # Soundgoodizer - a "make it sound good" one-knob effect.
    # Implemented as a warm multiband-style chain: gentle compression + saturation + makeup gain.
    # Amount 0 = off, 1 = max.
    if soundgoodizer and float(soundgoodizer) > 0:
        amt = float(soundgoodizer)
        # Warm compressor (glue)
        chain.append(Compressor(threshold_db=-18 + (1 - amt) * 6, ratio=2.0 + amt * 2.0, attack_ms=10.0, release_ms=120.0))
        # Subtle saturation via Distortion with low drive
        chain.append(Distortion(drive_db=amt * 4.0))
        # Makeup gain
        chain.append(Gain(gain_db=amt * 2.0))

    # Distortion / saturation
    if distortion_drive and float(distortion_drive) > 0:
        chain.append(Distortion(drive_db=float(distortion_drive)))

    # Bitcrush (lo-fi effect) - bit_depth 16 = no effect, lower = more crushed
    if bitcrush_bits and float(bitcrush_bits) > 0 and float(bitcrush_bits) < 16:
        chain.append(Bitcrush(bit_depth=float(bitcrush_bits)))

    # Clipping (hard limit)
    if clipping_threshold and float(clipping_threshold) < 0:
        chain.append(Clipping(threshold_db=float(clipping_threshold)))

    # Chorus (modulation)
    if chorus_mix and float(chorus_mix) > 0:
        chain.append(Chorus(
            rate_hz=float(chorus_rate) if chorus_rate else 0.5,
            depth=float(chorus_depth) if chorus_depth else 0.25,
            centre_delay_ms=7.0,
            feedback=0.0,
            mix=float(chorus_mix),
        ))

    # Phaser (modulation)
    if phaser_mix and float(phaser_mix) > 0:
        chain.append(Phaser(
            rate_hz=float(phaser_rate) if phaser_rate else 0.5,
            depth=float(phaser_depth) if phaser_depth else 0.5,
            centre_frequency_hz=1300.0,
            feedback=0.0,
            mix=float(phaser_mix),
        ))

    # Delay (echo)
    if delay_mix and float(delay_mix) > 0 and delay_seconds and float(delay_seconds) > 0:
        chain.append(Delay(
            delay_seconds=float(delay_seconds),
            feedback=float(delay_feedback) if delay_feedback else 0.0,
            mix=float(delay_mix),
        ))

    # Reverb
    chain.append(Reverb(room_size=float(reverb_rm_size), dry_level=float(reverb_dry), wet_level=float(reverb_wet), damping=float(reverb_damping)))

    # Limiter (peak ceiling)
    if limiter_ceiling is not None and float(limiter_ceiling) < 0:
        chain.append(Limiter(threshold_db=float(limiter_ceiling), release_ms=100.0))

    # Final gain
    if final_gain and float(final_gain) != 0:
        chain.append(Gain(gain_db=float(final_gain)))

    board = Pedalboard(chain)

    with AudioFile(audio_path) as f:
        with AudioFile(output_path, 'w', f.samplerate, f.num_channels) as o:
            while f.tell() < f.frames:
                chunk = f.read(int(f.samplerate))
                effected = board(chunk, f.samplerate, reset=False)
                o.write(effected)

    return output_path


def combine_audio(audio_paths, output_path, main_gain, backup_gain, inst_gain, output_format):
    main_vocal_audio = AudioSegment.from_wav(audio_paths[0]) - 4 + main_gain
    backup_vocal_audio = AudioSegment.from_wav(audio_paths[1]) - 6 + backup_gain
    instrumental_audio = AudioSegment.from_wav(audio_paths[2]) - 7 + inst_gain
    main_vocal_audio.overlay(backup_vocal_audio).overlay(instrumental_audio).export(output_path, format=output_format)


@spaces.GPU(duration=65)
def process_song(
    song_dir, song_input, mdx_model_params, song_id, is_webui, input_type, progress,
    keep_files, pitch_change, pitch_change_all, voice_model, index_rate, filter_radius,
    rms_mix_rate, protect, f0_method, crepe_hop_length, output_format, keep_orig, orig_song_path, steps
):

    if not os.path.exists(song_dir):
        os.makedirs(song_dir)
        orig_song_path, vocals_path, instrumentals_path, main_vocals_path, backup_vocals_path, main_vocals_dereverb_path = preprocess_song(song_input, mdx_model_params, song_id, is_webui, input_type, progress, keep_orig, orig_song_path)
    else:
        vocals_path, main_vocals_path = None, None
        paths = get_audio_paths(song_dir)

        # if any of the audio files aren't available or keep intermediate files, rerun preprocess
        if any(path is None for path in paths):
            orig_song_path, vocals_path, instrumentals_path, main_vocals_path, backup_vocals_path, main_vocals_dereverb_path = preprocess_song(song_input, mdx_model_params, song_id, is_webui, input_type, progress, keep_orig, orig_song_path)
        else:
            orig_song_path, instrumentals_path, main_vocals_dereverb_path, backup_vocals_path = paths

    ins_path = get_audio_with_suffix(song_dir, "_Voiceless.wav")
    if not ins_path:
        display_progress('Extracting voiceless instrumental track...', 0.40, is_webui, progress)
        instrumentals_path, _ = run_mdx(
            mdx_model_params,
            song_dir,
            os.path.join(mdxnet_models_dir, "UVR-MDX-NET-Inst_HQ_4.onnx"),
            instrumentals_path,
            exclude_inversion=True,
            suffix="Voiceless",
            denoise=False,
            keep_orig=True,
            base_device=("cuda" if IS_ZERO_GPU else "")
        )
        ins_path = get_audio_with_suffix(song_dir, "_Voiceless.wav")
    else:
        instrumentals_path = ins_path

    pitch_change = pitch_change * 12 + pitch_change_all
    ai_vocals_path = os.path.join(song_dir, f'{os.path.splitext(os.path.basename(orig_song_path))[0]}_{voice_model}_p{pitch_change}_i{index_rate}_fr{filter_radius}_rms{rms_mix_rate}_pro{protect}_{f0_method}{"" if f0_method != "mangio-crepe" else f"_{crepe_hop_length}"}_s{steps}.wav')
    ai_cover_path = os.path.join(song_dir, f'{os.path.splitext(os.path.basename(orig_song_path))[0]} ({voice_model} Ver).{output_format}')

    if not os.path.exists(ai_vocals_path):
        display_progress('Converting voice with RVC AI model...', 0.50, is_webui, progress)
        voice_change(voice_model, main_vocals_dereverb_path, ai_vocals_path, pitch_change, f0_method, index_rate, filter_radius, rms_mix_rate, protect, crepe_hop_length, is_webui, steps)

    return ai_vocals_path, ai_cover_path, instrumentals_path, backup_vocals_path, vocals_path, main_vocals_path, ins_path


def apply_noisereduce(audio_list, type_output="wav"):
    # https://github.com/sa-if/Audio-Denoiser
    print("Noise reduce")

    result = []
    for audio_path in audio_list:
        out_path = f"{os.path.splitext(audio_path)[0]}_nr.{type_output}"

        try:
            # Load audio file
            audio = AudioSegment.from_file(audio_path)

            # Convert to numpy with shape (channels, samples) for stereo support
            samples = np.array(audio.get_array_of_samples()).astype(np.float32)
            if audio.channels > 1:
                samples = samples.reshape(-1, audio.channels).T  # (channels, samples)

            # Reduce noise (handles 1D and 2D)
            reduced_noise = nr.reduce_noise(samples, sr=audio.frame_rate, prop_decrease=0.6)

            # Convert back to int16 interleaved
            reduced_noise = np.clip(reduced_noise, -32768, 32767).astype(np.int16)
            if audio.channels > 1:
                reduced_noise = reduced_noise.T.reshape(-1)

            reduced_audio = AudioSegment(
                reduced_noise.tobytes(),
                frame_rate=audio.frame_rate,
                sample_width=audio.sample_width,
                channels=audio.channels,
            )

            reduced_audio.export(out_path, format=type_output)
            result.append(out_path)

        except Exception as e:
            print(f"Error noisereduce: {str(e)}")
            result.append(audio_path)

    return result



def song_cover_pipeline(song_input, voice_model, pitch_change, keep_files,
                        is_webui=0, main_gain=0, backup_gain=0, inst_gain=0, index_rate=0.5, filter_radius=3,
                        rms_mix_rate=0.25, f0_method='rmvpe', crepe_hop_length=128, protect=0.33, pitch_change_all=0,
                        reverb_rm_size=0.15, reverb_wet=0.2, reverb_dry=0.8, reverb_damping=0.7, output_format='mp3',
                        extra_denoise=False, steps=1,
                        highpass_freq=80.0, lowpass_freq=0.0,
                        low_shelf_gain=0.0, low_shelf_freq=200.0,
                        mid_peak_gain=0.0, mid_peak_freq=1000.0, mid_peak_q=1.0,
                        high_shelf_gain=0.0, high_shelf_freq=5000.0,
                        noise_gate_threshold=-55.0, noise_gate_ratio=10.0,
                        compressor_threshold=-15.0, compressor_ratio=4.0,
                        limiter_ceiling=-1.0, final_gain=0.0,
                        soundgoodizer=0.0,
                        distortion_drive=0.0,
                        bitcrush_bits=0.0,
                        chorus_depth=0.0, chorus_rate=0.0, chorus_mix=0.0,
                        delay_seconds=0.0, delay_feedback=0.0, delay_mix=0.0,
                        phaser_depth=0.0, phaser_rate=0.0, phaser_mix=0.0,
                        clipping_threshold=0.0,
                        backup_vocal_infer=False, backup_vocal_pitch=0,
                        progress=gr.Progress()):
    if not keep_files or IS_ZERO_GPU:
        clean_old_folders("./song_output", 14400)

    if IS_ZERO_GPU:
        clean_old_folders("./rvc_models", 10800)

    try:
        if not song_input or not voice_model:
            raise_exception('Ensure that the song input field and voice model field is filled.', is_webui)

        display_progress('Starting AI cover generation...', 0.00, is_webui, progress)

        with open(os.path.join(mdxnet_models_dir, 'model_data.json')) as infile:
            mdx_model_params = json.load(infile)

        # if youtube url
        if urlparse(song_input).scheme == 'https':
            input_type = 'yt'
            song_id = get_youtube_video_id(song_input)
            if song_id is None:
                error_msg = 'Invalid YouTube url.'
                raise_exception(error_msg, is_webui)

        # local audio file
        else:
            input_type = 'local'
            song_input = song_input.strip('\"')
            if os.path.exists(song_input):
                song_id = get_hash(song_input)
            else:
                error_msg = f'{song_input} does not exist.'
                song_id = None
                raise_exception(error_msg, is_webui)

        song_dir = os.path.join(output_dir, song_id)

        keep_orig, orig_song_path = get_audio_file(song_input, is_webui, input_type, progress)
        orig_song_path = convert_to_stereo(orig_song_path)

        start = time.time()

        (
            ai_vocals_path,
            ai_cover_path,
            instrumentals_path,
            backup_vocals_path,
            vocals_path,
            main_vocals_path,
            ins_path
        ) = process_song(
            song_dir,
            song_input,
            mdx_model_params,
            song_id,
            is_webui,
            input_type,
            progress,
            keep_files,
            pitch_change,
            pitch_change_all,
            voice_model,
            index_rate,
            filter_radius,
            rms_mix_rate,
            protect,
            f0_method,
            crepe_hop_length,
            output_format,
            keep_orig,
            orig_song_path,
            steps,
        )

        end = time.time()
        print(f"Execution time: {end - start:.4f} seconds")
        with sf.SoundFile(ai_vocals_path) as f:
            duration__ = len(f) / f.samplerate
        print(f"Audio duration: {duration__:.2f} seconds")

        display_progress('Applying audio effects to AI vocals...', 0.65, is_webui, progress)

        nr_path = ai_vocals_path  # get_audio_with_suffix(song_dir, "_nr.wav")
        if extra_denoise:
            ai_vocals_path = apply_noisereduce([ai_vocals_path])[0]

        ai_vocals_mixed_path = add_audio_effects(
            ai_vocals_path, reverb_rm_size, reverb_wet, reverb_dry, reverb_damping,
            highpass_freq=highpass_freq, lowpass_freq=lowpass_freq,
            low_shelf_gain=low_shelf_gain, low_shelf_freq=low_shelf_freq,
            mid_peak_gain=mid_peak_gain, mid_peak_freq=mid_peak_freq, mid_peak_q=mid_peak_q,
            high_shelf_gain=high_shelf_gain, high_shelf_freq=high_shelf_freq,
            noise_gate_threshold=noise_gate_threshold, noise_gate_ratio=noise_gate_ratio,
            compressor_threshold=compressor_threshold, compressor_ratio=compressor_ratio,
            limiter_ceiling=limiter_ceiling, final_gain=final_gain,
            soundgoodizer=soundgoodizer,
            distortion_drive=distortion_drive,
            bitcrush_bits=bitcrush_bits,
            chorus_depth=chorus_depth, chorus_rate=chorus_rate, chorus_mix=chorus_mix,
            delay_seconds=delay_seconds, delay_feedback=delay_feedback, delay_mix=delay_mix,
            phaser_depth=phaser_depth, phaser_rate=phaser_rate, phaser_mix=phaser_mix,
            clipping_threshold=clipping_threshold,
        )
        
        if pitch_change_all != 0:
            display_progress('Applying overall pitch change to instrumentals...', 0.75, is_webui, progress)
            instrumentals_path = pitch_shift(instrumentals_path, pitch_change_all)
            backup_vocals_path = pitch_shift(backup_vocals_path, pitch_change_all)

        # Backing vocal inference: run RVC on backup vocals with optional harmony pitch
        ai_backup_vocals_mixed_path = None
        if backup_vocal_infer and backup_vocals_path:
            display_progress('Converting backing vocals with RVC...', 0.80, is_webui, progress)
            backup_pitch = pitch_change + (backup_vocal_pitch * 12)
            ai_backup_vocals_path = os.path.join(song_dir, f'{os.path.splitext(os.path.basename(orig_song_path))[0]}_{voice_model}_backup_p{backup_pitch}.wav')
            if not os.path.exists(ai_backup_vocals_path):
                voice_change(voice_model, backup_vocals_path, ai_backup_vocals_path, backup_pitch, f0_method, index_rate, filter_radius, rms_mix_rate, protect, crepe_hop_length, is_webui, steps)
            # Apply same audio effects to backup vocals (slightly reduced reverb)
            ai_backup_vocals_mixed_path = add_audio_effects(
                ai_backup_vocals_path, reverb_rm_size * 0.7, reverb_wet * 0.5, reverb_dry, reverb_damping,
                highpass_freq=highpass_freq, lowpass_freq=lowpass_freq,
                low_shelf_gain=low_shelf_gain, low_shelf_freq=low_shelf_freq,
                mid_peak_gain=mid_peak_gain, mid_peak_freq=mid_peak_freq, mid_peak_q=mid_peak_q,
                high_shelf_gain=high_shelf_gain, high_shelf_freq=high_shelf_freq,
                noise_gate_threshold=noise_gate_threshold, noise_gate_ratio=noise_gate_ratio,
                compressor_threshold=compressor_threshold, compressor_ratio=compressor_ratio,
                limiter_ceiling=limiter_ceiling, final_gain=final_gain,
                soundgoodizer=soundgoodizer,
                distortion_drive=distortion_drive,
                bitcrush_bits=bitcrush_bits,
                chorus_depth=chorus_depth, chorus_rate=chorus_rate, chorus_mix=chorus_mix,
                delay_seconds=delay_seconds, delay_feedback=delay_feedback, delay_mix=delay_mix,
                phaser_depth=phaser_depth, phaser_rate=phaser_rate, phaser_mix=phaser_mix,
                clipping_threshold=clipping_threshold,
            )
            # Use the AI backup vocals instead of original
            backup_vocals_path = ai_backup_vocals_mixed_path

        display_progress('Combining AI vocals and instrumentals...', 0.90, is_webui, progress)
        combine_audio([ai_vocals_mixed_path, backup_vocals_path, instrumentals_path], ai_cover_path, main_gain, backup_gain, inst_gain, output_format)

        if not keep_files:
            display_progress('Cleaning up intermediate files...', 0.95, is_webui, progress)
            intermediate_files = [vocals_path, main_vocals_path, ai_vocals_mixed_path, nr_path]
            if ai_backup_vocals_mixed_path:
                intermediate_files.append(ai_backup_vocals_mixed_path)
            if pitch_change_all != 0:
                intermediate_files += [instrumentals_path, backup_vocals_path]
            for file in intermediate_files:
                if file and os.path.exists(file) and file != ins_path:
                    os.remove(file)

        # Record history entry (if modules are available)
        try:
            hist = get_history_manager()
            if hist is not None:
                song_name = os.path.splitext(os.path.basename(orig_song_path))[0]
                hist.add(
                    song_name=song_name,
                    voice_model=voice_model,
                    pitch=pitch_change,
                    pitch_all=pitch_change_all,
                    ai_cover_path=ai_cover_path,
                    duration_seconds=duration__,
                    output_format=output_format,
                    settings={
                        "index_rate": index_rate,
                        "filter_radius": filter_radius,
                        "rms_mix_rate": rms_mix_rate,
                        "protect": protect,
                        "f0_method": f0_method,
                        "crepe_hop_length": crepe_hop_length,
                        "steps": steps,
                        "main_gain": main_gain,
                        "backup_gain": backup_gain,
                        "inst_gain": inst_gain,
                        "reverb_rm_size": reverb_rm_size,
                        "reverb_wet": reverb_wet,
                        "reverb_dry": reverb_dry,
                        "reverb_damping": reverb_damping,
                        "extra_denoise": extra_denoise,
                        "processing_time_seconds": round(end - start, 2),
                        "backup_vocal_infer": backup_vocal_infer,
                        "backup_vocal_pitch": backup_vocal_pitch,
                    },
                )
        except Exception as hist_err:
            print(f"[history] could not record entry: {hist_err}")

        display_progress('AI cover generated successfully!', 1.0, is_webui, progress)
        return ai_cover_path

    except Exception as e:
        raise_exception(str(e), is_webui)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Generate a AI cover song in the song_output/id directory.', add_help=True)
    parser.add_argument('-i', '--song-input', type=str, required=True, help='Link to a YouTube video or the filepath to a local mp3/wav file to create an AI cover of')
    parser.add_argument('-dir', '--rvc-dirname', type=str, required=True, help='Name of the folder in the rvc_models directory containing the RVC model file and optional index file to use')
    parser.add_argument('-p', '--pitch-change', type=int, required=True, help='Change the pitch of AI Vocals only. Generally, use 1 for male to female and -1 for vice-versa. (Octaves)')
    parser.add_argument('-k', '--keep-files', action=argparse.BooleanOptionalAction, help='Whether to keep all intermediate audio files generated in the song_output/id directory, e.g. Isolated Vocals/Instrumentals')
    parser.add_argument('-ir', '--index-rate', type=float, default=0.5, help='A decimal number e.g. 0.5, used to reduce/resolve the timbre leakage problem. If set to 1, more biased towards the timbre quality of the training dataset')
    parser.add_argument('-fr', '--filter-radius', type=int, default=3, help='A number between 0 and 7. If >=3: apply median filtering to the harvested pitch results. The value represents the filter radius and can reduce breathiness.')
    parser.add_argument('-rms', '--rms-mix-rate', type=float, default=0.25, help="A decimal number e.g. 0.25. Control how much to use the original vocal's loudness (0) or a fixed loudness (1).")
    parser.add_argument('-palgo', '--pitch-detection-algo', type=str, default='rmvpe', help='Best option is rmvpe (clarity in vocals), then mangio-crepe (smoother vocals).')
    parser.add_argument('-hop', '--crepe-hop-length', type=int, default=128, help='If pitch detection algo is mangio-crepe, controls how often it checks for pitch changes in milliseconds. The higher the value, the faster the conversion and less risk of voice cracks, but there is less pitch accuracy. Recommended: 128.')
    parser.add_argument('-pro', '--protect', type=float, default=0.33, help='A decimal number e.g. 0.33. Protect voiceless consonants and breath sounds to prevent artifacts such as tearing in electronic music. Set to 0.5 to disable. Decrease the value to increase protection, but it may reduce indexing accuracy.')
    parser.add_argument('-mv', '--main-vol', type=int, default=0, help='Volume change for AI main vocals in decibels. Use -3 to decrease by 3 decibels and 3 to increase by 3 decibels')
    parser.add_argument('-bv', '--backup-vol', type=int, default=0, help='Volume change for backup vocals in decibels')
    parser.add_argument('-iv', '--inst-vol', type=int, default=0, help='Volume change for instrumentals in decibels')
    parser.add_argument('-pall', '--pitch-change-all', type=int, default=0, help='Change the pitch/key of vocals and instrumentals. Changing this slightly reduces sound quality')
    parser.add_argument('-rsize', '--reverb-size', type=float, default=0.15, help='Reverb room size between 0 and 1')
    parser.add_argument('-rwet', '--reverb-wetness', type=float, default=0.2, help='Reverb wet level between 0 and 1')
    parser.add_argument('-rdry', '--reverb-dryness', type=float, default=0.8, help='Reverb dry level between 0 and 1')
    parser.add_argument('-rdamp', '--reverb-damping', type=float, default=0.7, help='Reverb damping between 0 and 1')
    parser.add_argument('-oformat', '--output-format', type=str, default='mp3', help='Output format of audio file. mp3 for smaller file size, wav for best quality')
    args = parser.parse_args()

    rvc_dirname = args.rvc_dirname
    if not os.path.exists(os.path.join(rvc_models_dir, rvc_dirname)):
        raise Exception(f'The folder {os.path.join(rvc_models_dir, rvc_dirname)} does not exist.')

    cover_path = song_cover_pipeline(args.song_input, rvc_dirname, args.pitch_change, args.keep_files,
                                     main_gain=args.main_vol, backup_gain=args.backup_vol, inst_gain=args.inst_vol,
                                     index_rate=args.index_rate, filter_radius=args.filter_radius,
                                     rms_mix_rate=args.rms_mix_rate, f0_method=args.pitch_detection_algo,
                                     crepe_hop_length=args.crepe_hop_length, protect=args.protect,
                                     pitch_change_all=args.pitch_change_all,
                                     reverb_rm_size=args.reverb_size, reverb_wet=args.reverb_wetness,
                                     reverb_dry=args.reverb_dryness, reverb_damping=args.reverb_damping,
                                     output_format=args.output_format)
    print(f'[+] Cover generated at {cover_path}')