Update main.py

main.py

@@ -21,7 +21,7 @@ from pydub import AudioSegment
 from mdx import run_mdx
 from rvc import Config, load_hubert, get_vc, rvc_infer
 
-BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+BASE_DIR = 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')
@@ -30,11 +30,7 @@ output_dir = os.path.join(BASE_DIR, 'song_output')
 
 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
+    Extracts the video ID from a YouTube URL.
     """
     query = urlparse(url)
     if query.hostname == 'youtu.be':
@@ -44,7 +40,6 @@ def get_youtube_video_id(url, ignore_playlist=True):
 
     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':
@@ -55,12 +50,17 @@ def get_youtube_video_id(url, ignore_playlist=True):
             return query.path.split('/')[2]
         if query.path[:3] == '/v/':
             return query.path.split('/')[2]
+        if query.path[:8] == '/shorts/':
+            return query.path.split('/')[2]
 
-    # returns None for invalid YouTube url
     return None
 
 
+
 def yt_download(link):
+    """
+    Downloads the best audio format from a YouTube link.
+    """
     ydl_opts = {
         'format': 'bestaudio',
         'outtmpl': '%(title)s',
@@ -125,7 +125,6 @@ def get_audio_paths(song_dir):
 def convert_to_stereo(audio_path):
     wave, sr = librosa.load(audio_path, mono=False, sr=44100)
 
-    # check if mono
     if type(wave[0]) != np.ndarray:
         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}"')
@@ -197,7 +196,6 @@ def voice_change(voice_model, vocals_path, output_path, pitch_change, f0_method,
     hubert_model = load_hubert(device, config.is_half, os.path.join(rvc_models_dir, 'hubert_base.pt'))
     cpt, version, net_g, tgt_sr, vc = get_vc(device, config.is_half, config, rvc_model_path)
 
-    # convert main vocals
     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()
@@ -206,18 +204,16 @@ def voice_change(voice_model, vocals_path, output_path, pitch_change, f0_method,
 def add_audio_effects(audio_path, reverb_rm_size, reverb_wet, reverb_dry, reverb_damping):
     output_path = f'{os.path.splitext(audio_path)[0]}_mixed.wav'
 
-    # Initialize audio effects plugins
     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:
-            # Read one second of audio at a time, until the file is empty:
             while f.tell() < f.frames:
                 chunk = f.read(int(f.samplerate))
                 effected = board(chunk, f.samplerate, reset=False)
@@ -226,138 +222,91 @@ def add_audio_effects(audio_path, reverb_rm_size, reverb_wet, reverb_dry, reverb
     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)
+def merge_audios(audio_paths, output_path):
+    combined = AudioSegment.from_file(audio_paths[0])
+    for path in audio_paths[1:]:
+        combined = combined.overlay(AudioSegment.from_file(path))
+    combined.export(output_path, format='wav')
 
 
-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()):
+def process_and_save_song(song_input, input_type, voice_model, pitch_change, f0_method, index_rate, filter_radius, rms_mix_rate, protect, crepe_hop_length, reverb_rm_size, reverb_wet, reverb_dry, reverb_damping, progress, is_webui=False):
+    song_id = get_hash(song_input)
+    if not os.path.exists(output_dir):
+        os.makedirs(output_dir)
+
+    song_output_dir = os.path.join(output_dir, song_id)
+    if not os.path.exists(song_output_dir):
+        os.makedirs(song_output_dir)
+
+    if input_type == 'yt' and not get_youtube_video_id(song_input):
+        raise_exception('[!] Invalid YouTube link.', is_webui)
+
+    mdx_model_params = {
+        'demucs_model_path': os.path.join(mdxnet_models_dir, 'models_demucs.h5'),
+        'mdx_model_path': os.path.join(mdxnet_models_dir, 'models_mdx.h5'),
+        'output_path': output_dir,
+        'noise_protect': 0.33,
+        'voc_model_path': os.path.join(mdxnet_models_dir, 'models_vocal.h5')
+    }
+
     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)
-
-        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)
-
-        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)
-
-        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) or keep_files:
-                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)
-            else:
-                orig_song_path, instrumentals_path, main_vocals_dereverb_path, backup_vocals_path = paths
-
-        pitch_change = pitch_change + pitch_change_all
-        ai_vocals_path = os.path.join(song_dir, f'{os.path.splitext(os.path.basename(orig_song_path))[0]}_lead_{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}"}.wav')
-        ai_backing_path = os.path.join(song_dir, f'{os.path.splitext(os.path.basename(orig_song_path))[0]}_backing_{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}"}.wav')
-
-        ai_cover_path = os.path.join(song_dir, f'normal {os.path.splitext(os.path.basename(orig_song_path))[0]} ({voice_model} Ver).{output_format}')
-        ai_cover_backing_path = os.path.join(song_dir, f'with backing {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 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
+        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)
+
+        display_progress('[~] Changing Main Vocals to Target Voice...', 0.4, is_webui, progress)
+        pitch_shifted_main_vocals_path = pitch_shift(main_vocals_dereverb_path, pitch_change)
+        output_vocals_path = os.path.join(song_output_dir, 'main_vocals_changed.wav')
+        voice_change(voice_model, pitch_shifted_main_vocals_path, output_vocals_path, pitch_change, f0_method, index_rate, filter_radius, rms_mix_rate, protect, crepe_hop_length, is_webui)
+
+        display_progress('[~] Adding Audio Effects...', 0.5, is_webui, progress)
+        final_output_vocals_path = add_audio_effects(output_vocals_path, reverb_rm_size, reverb_wet, reverb_dry, reverb_damping)
+
+        display_progress('[~] Merging Vocal and Instrumental Tracks...', 0.6, is_webui, progress)
+        final_output_path = os.path.join(output_dir, f'{os.path.basename(orig_song_path)}_{voice_model}_vocal_conversion.wav')
+        merge_audios([final_output_vocals_path, instrumentals_path], final_output_path)
 
+        display_progress('[~] Done!', 1.0, is_webui, progress)
     except Exception as e:
-        raise_exception(str(e), is_webui)
+        raise_exception(f'[!] Processing failed: {str(e)}', is_webui)
+    finally:
+        with suppress(FileNotFoundError):
+            os.remove(orig_song_path)
+
+    return final_output_path
 
 
 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')
+    parser = argparse.ArgumentParser(description='Process song with RVC.')
+    parser.add_argument('--input', type=str, required=True, help='Path to the input song file or YouTube link.')
+    parser.add_argument('--type', type=str, required=True, choices=['local', 'yt'], help='Type of input: "local" for a file, "yt" for a YouTube link.')
+    parser.add_argument('--voice_model', type=str, required=True, help='Name of the voice model to use.')
+    parser.add_argument('--pitch_change', type=float, required=False, default=0, help='Pitch change amount in semitones.')
+    parser.add_argument('--f0_method', type=str, required=False, default='crepe', help='F0 method to use.')
+    parser.add_argument('--index_rate', type=float, required=False, default=1.0, help='Index rate.')
+    parser.add_argument('--filter_radius', type=float, required=False, default=3.0, help='Filter radius.')
+    parser.add_argument('--rms_mix_rate', type=float, required=False, default=0.25, help='RMS mix rate.')
+    parser.add_argument('--protect', type=float, required=False, default=0.33, help='Protection rate.')
+    parser.add_argument('--crepe_hop_length', type=int, required=False, default=128, help='Crepe hop length.')
+    parser.add_argument('--reverb_rm_size', type=float, required=False, default=0.3, help='Reverb room size.')
+    parser.add_argument('--reverb_wet', type=float, required=False, default=0.25, help='Reverb wet level.')
+    parser.add_argument('--reverb_dry', type=float, required=False, default=0.75, help='Reverb dry level.')
+    parser.add_argument('--reverb_damping', type=float, required=False, default=0.5, help='Reverb damping.')
+
     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}')
+    process_and_save_song(
+        song_input=args.input,
+        input_type=args.type,
+        voice_model=args.voice_model,
+        pitch_change=args.pitch_change,
+        f0_method=args.f0_method,
+        index_rate=args.index_rate,
+        filter_radius=args.filter_radius,
+        rms_mix_rate=args.rms_mix_rate,
+        protect=args.protect,
+        crepe_hop_length=args.crepe_hop_length,
+        reverb_rm_size=args.reverb_rm_size,
+        reverb_wet=args.reverb_wet,
+        reverb_dry=args.reverb_dry,
+        reverb_damping=args.reverb_damping,
+        progress=None
+    )