Test.py

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 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
from pedalboard.io import AudioFile
from pydub import AudioSegment
from audio_separator.separator import Separator
from rvc import Config, load_hubert, get_vc, rvc_infer

# Base directories
BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
mdxnet_models_dir = os.path.join(BASE_DIR, 'mdxnet_models')
rvc_models_dir = os.path.join(BASE_DIR, 'rvc_models')
output_dir = os.path.join(BASE_DIR, 'song_output')


def get_youtube_video_id(url, ignore_playlist=True):
    """
    Extract the YouTube video ID from various URL formats.
    
    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
    """
    parsed_url = urlparse(url)
    hostname = parsed_url.hostname or ''
    path = parsed_url.path

    if hostname.lower() == 'youtu.be':
        return path.lstrip('/')

    if hostname.lower() in {'www.youtube.com', 'youtube.com', 'music.youtube.com'}:
        if not ignore_playlist:
            with suppress(KeyError):
                return parse_qs(parsed_url.query)['list'][0]
        if parsed_url.path == '/watch':
            return parse_qs(parsed_url.query).get('v', [None])[0]
        if parsed_url.path.startswith('/watch/'):
            return parsed_url.path.split('/')[1]
        if parsed_url.path.startswith('/embed/'):
            return parsed_url.path.split('/')[2]
        if parsed_url.path.startswith('/v/'):
            return parsed_url.path.split('/')[2]

    return None


def yt_download(link):
    """
    Download the audio from a YouTube link as an mp3 file.
    """
    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 display_progress(message, percent, is_webui, progress=None):
    """
    Display progress either via the provided progress callback or by printing.
    """
    if is_webui and progress is not None:
        progress(percent, desc=message)
    else:
        print(message)


def raise_exception(error_msg, is_webui):
    """
    Raise an exception. If running in a web UI, use gr.Error.
    """
    if is_webui:
        raise gr.Error(error_msg)
    else:
        raise Exception(error_msg)


def get_rvc_model(voice_model, is_webui):
    """
    Search the specified RVC model directory for the model (.pth) and index (.index) files.
    """
    rvc_model_filename, rvc_index_filename = None, None
    model_dir = os.path.join(rvc_models_dir, voice_model)
    if not os.path.exists(model_dir):
        raise_exception(f'Model directory {model_dir} does not exist.', is_webui)
    for file in os.listdir(model_dir):
        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)

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


def separation_uvr(filename, output):
    """
    Run the separation steps using different pre-trained models.
    Returns a tuple of four file paths:
      - vocals_no_reverb: The vocals after initial de-echo/de-reverb (used as intermediate vocals)
      - instrumental_path: The separated instrumental audio
      - main_vocals_dereverb: The lead vocals after final de-reverb processing
      - backup_vocals: The backup vocals extracted in the final stage
    """
    separator = Separator(output_dir=output)
    base_name = os.path.splitext(os.path.basename(filename))[0]

    instrumental_path = os.path.join(output, f'{base_name}_Instrumental.wav')
    initial_vocals = os.path.join(output, f'{base_name}_Vocals.wav')
    vocals_no_reverb = os.path.join(output, f'{base_name}_Vocals (No Reverb).wav')
    vocals_reverb = os.path.join(output, f'{base_name}_Vocals (Reverb).wav')
    main_vocals_dereverb = os.path.join(output, f'{base_name}_Vocals_Main_DeReverb.wav')
    backup_vocals = os.path.join(output, f'{base_name}_Vocals_Backup.wav')

    separator.load_model(model_filename='model_bs_roformer_ep_317_sdr_12.9755.ckpt')
    voc_inst = separator.separate(filename)
    os.rename(os.path.join(output, voc_inst[0]), instrumental_path)
    os.rename(os.path.join(output, voc_inst[1]), initial_vocals)

    separator.load_model(model_filename='UVR-DeEcho-DeReverb.pth')
    voc_no_reverb = separator.separate(initial_vocals)
    os.rename(os.path.join(output, voc_no_reverb[0]), vocals_no_reverb)
    os.rename(os.path.join(output, voc_no_reverb[1]), vocals_reverb)

    separator.load_model(model_filename='mel_band_roformer_karaoke_aufr33_viperx_sdr_10.1956.ckpt')
    voc_split = separator.separate(vocals_no_reverb)
    os.rename(os.path.join(output, voc_split[0]), backup_vocals)
    os.rename(os.path.join(output, voc_split[1]), main_vocals_dereverb)

    if os.path.exists(vocals_reverb):
        os.remove(vocals_reverb)

    return vocals_no_reverb, instrumental_path, main_vocals_dereverb, backup_vocals


def get_audio_paths(song_dir):
    """
    Search the given directory for expected audio files.
    Returns:
      orig_song_path, instrumentals_path, main_vocals_dereverb_path, backup_vocals_path
    """
    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 convert_to_stereo(audio_path):
    """
    Convert the given audio file to stereo (2 channels) if it is mono.
    """
    wave, sr = librosa.load(audio_path, mono=False, sr=44100)
    if wave.ndim == 1:
        stereo_path = f'{os.path.splitext(audio_path)[0]}_stereo.wav'
        command = shlex.split(f'ffmpeg -y -loglevel error -i "{audio_path}" -ac 2 -f wav "{stereo_path}"')
        subprocess.run(command, check=True)
        return stereo_path
    return audio_path


def pitch_shift(audio_path, pitch_change):
    """
    Shift the pitch of the audio by the specified amount.
    """
    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):
    """
    Calculate a short BLAKE2b hash for the given file.
    """
    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 preprocess_song(song_input, song_id, is_webui, input_type, progress):
    """
    Preprocess the input song:
      - Download if YouTube URL.
      - Convert to stereo.
      - Separate vocals and instrumentals.
    Returns a tuple with six values matching the expected unpacking in the pipeline.
    """
    if input_type == 'yt':
        display_progress('[~] Downloading song...', 0, 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
    else:
        orig_song_path = None

    song_output_dir = os.path.join(output_dir, song_id)
    if not os.path.exists(song_output_dir):
        os.makedirs(song_output_dir)

    orig_song_path = convert_to_stereo(orig_song_path)

    display_progress('[~] Separating Vocals from Instrumental...', 0.1, is_webui, progress)
    vocals_no_reverb, instrumental_path, main_vocals_dereverb, backup_vocals = separation_uvr(orig_song_path, song_output_dir)
    return orig_song_path, vocals_no_reverb, instrumental_path, main_vocals_dereverb, backup_vocals, main_vocals_dereverb


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):
    """
    Convert the input vocals using the specified RVC model.
    """
    rvc_model_path, rvc_index_path = get_rvc_model(voice_model, is_webui)
    device = 'cuda:0'
    config = Config(device, True)
    hubert_model = load_hubert(embedder_model="contentvec", embedder_model_custom=None)
    cpt, version, net_g, tgt_sr, vc = get_vc(device, config.is_half, config, rvc_model_path)

    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
    )
    del hubert_model, cpt
    gc.collect()


def add_audio_effects(audio_path, reverb_rm_size, reverb_wet, reverb_dry, reverb_damping):
    """
    Apply a chain of audio effects (highpass, compression, reverb) to the input audio.
    """
    output_path = f'{os.path.splitext(audio_path)[0]}_mixed.wav'
    board = Pedalboard([
        HighpassFilter(),
        Compressor(ratio=4, threshold_db=-15),
        Reverb(room_size=reverb_rm_size, dry_level=reverb_dry, wet_level=reverb_wet, damping=reverb_damping)
    ])

    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):
    """
    Combine main vocals, backup vocals, and instrumental audio into a final mix.
    """
    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
    final_audio = main_vocal_audio.overlay(backup_vocal_audio).overlay(instrumental_audio)
    final_audio.export(output_path, format=output_format)


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',
                        progress=gr.Progress()):
    """
    Main pipeline that orchestrates the AI cover song generation.
    """
    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 Pipeline...', 0, is_webui, progress)

        if urlparse(song_input).scheme == 'https':
            input_type = 'yt'
            song_id = get_youtube_video_id(song_input)
            if song_id is None:
                raise_exception('Invalid YouTube url.', is_webui)
        else:
            input_type = 'local'
            song_input = song_input.strip('\"')
            if os.path.exists(song_input):
                song_id = get_hash(song_input)
            else:
                raise_exception(f'{song_input} does not exist.', is_webui)

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

        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, song_id, is_webui, input_type, progress
            )
        else:
            vocals_path, main_vocals_path = None, None
            paths = get_audio_paths(song_dir)
            if any(path is None for path in paths) or keep_files:
                (orig_song_path, vocals_path, instrumentals_path,
                 main_vocals_path, backup_vocals_path, main_vocals_dereverb_path) = preprocess_song(
                    song_input, song_id, is_webui, input_type, progress
                )
            else:
                orig_song_path, instrumentals_path, main_vocals_dereverb_path, backup_vocals_path = paths
                main_vocals_path = main_vocals_dereverb_path

        pitch_change += pitch_change_all

        base_song_name = os.path.splitext(os.path.basename(orig_song_path))[0]
        algo_suffix = f"_{crepe_hop_length}" if f0_method == "mangio-crepe" else ""
        ai_vocals_path = os.path.join(
            song_dir,
            f'{base_song_name}_lead_{voice_model}_p{pitch_change}_i{index_rate}_fr{filter_radius}_'
            f'rms{rms_mix_rate}_pro{protect}_{f0_method}{algo_suffix}.wav'
        )
        ai_backing_path = os.path.join(
            song_dir,
            f'{base_song_name}_backing_{voice_model}_p{pitch_change}_i{index_rate}_fr{filter_radius}_'
            f'rms{rms_mix_rate}_pro{protect}_{f0_method}{algo_suffix}.wav'
        )
        ai_cover_path = os.path.join(song_dir, f'{base_song_name} ({voice_model} Ver).{output_format}')
        ai_cover_backing_path = os.path.join(song_dir, f'{base_song_name} ({voice_model} Ver With Backing).{output_format}')

        if not os.path.exists(ai_vocals_path):
            display_progress('[~] Converting lead voice using RVC...', 0.5, 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)

            display_progress('[~] Converting backing voice using RVC...', 0.65, is_webui, progress)
            voice_change(voice_model, backup_vocals_path, ai_backing_path, pitch_change,
                         f0_method, index_rate, filter_radius, rms_mix_rate, protect, crepe_hop_length, is_webui)

        display_progress('[~] Applying audio effects to Vocals...', 0.8, is_webui, progress)
        ai_vocals_mixed_path = add_audio_effects(ai_vocals_path, reverb_rm_size, reverb_wet, reverb_dry, reverb_damping)
        ai_backing_mixed_path = add_audio_effects(ai_backing_path, reverb_rm_size, reverb_wet, reverb_dry, reverb_damping)

        if pitch_change_all != 0:
            display_progress('[~] Applying overall pitch change', 0.85, is_webui, progress)
            instrumentals_path = pitch_shift(instrumentals_path, pitch_change_all)
            backup_vocals_path = pitch_shift(backup_vocals_path, pitch_change_all)

        display_progress('[~] Combining AI Vocals and Instrumentals...', 0.9, 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)
        combine_audio([ai_vocals_mixed_path, ai_backing_mixed_path, instrumentals_path],
                      ai_cover_backing_path, main_gain, backup_gain, inst_gain, output_format)

        if not keep_files:
            display_progress('[~] Removing intermediate audio files...', 0.95, is_webui, progress)
            intermediate_files = [vocals_path, main_vocals_path, ai_vocals_mixed_path, ai_backing_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):
                    os.remove(file)

        return ai_cover_path, ai_cover_backing_path

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


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description='AICoverGen: Mod.',
        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. 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.')
    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 dB and 3 to increase by 3 dB')
    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_dir = os.path.join(rvc_models_dir, args.rvc_dirname)
    if not os.path.exists(rvc_dir):
        raise Exception(f'The folder {rvc_dir} does not exist.')

    cover_path, cover_with_backing = song_cover_pipeline(
        args.song_input, args.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}')