Create mdx_processing_script.py

mdx_processing_script.py

@@ -0,0 +1,153 @@
+import gc
+import requests
+import subprocess
+import logging
+import sys
+from bs4 import BeautifulSoup
+import torch, pdb, os, warnings, librosa
+import soundfile as sf
+from tqdm import tqdm
+import numpy as np
+import torch
+now_dir = os.getcwd()
+sys.path.append(now_dir)
+import mdx
+branch = "https://github.com/NaJeongMo/Colab-for-MDX_B"
+
+model_params = "https://raw.githubusercontent.com/TRvlvr/application_data/main/mdx_model_data/model_data.json"
+_Models = "https://github.com/TRvlvr/model_repo/releases/download/all_public_uvr_models/"
+# _models = "https://pastebin.com/raw/jBzYB8vz"
+_models = "https://raw.githubusercontent.com/TRvlvr/application_data/main/filelists/download_checks.json"
+#stem_naming = "https://pastebin.com/raw/mpH4hRcF"
+
+file_folder = "Colab-for-MDX_B"
+model_ids = requests.get(_models).json()
+model_ids = model_ids["mdx_download_list"].values()
+#print(model_ids)
+model_params = requests.get(model_params).json()
+#stem_naming = requests.get(stem_naming).json()
+stem_naming = {
+    "Vocals": "Instrumental",
+    "Other": "Instruments",
+    "Instrumental": "Vocals",
+    "Drums": "Drumless",
+    "Bass": "Bassless"
+}
+
+os.makedirs("tmp_models", exist_ok=True)
+
+warnings.filterwarnings("ignore")
+cpu = torch.device("cpu")
+if torch.cuda.is_available():
+    device = torch.device("cuda:0")
+elif torch.backends.mps.is_available():
+    device = torch.device("mps")
+else:
+    device = torch.device("cpu")
+
+
+def get_model_list():
+    return model_ids
+
+def id_to_ptm(mkey):
+    if mkey in model_ids:
+        mpath = f"{now_dir}/tmp_models/{mkey}"
+        if not os.path.exists(f'{now_dir}/tmp_models/{mkey}'):
+            print('Downloading model...',end=' ')
+            subprocess.run(
+                ["wget", _Models+mkey, "-O", mpath]
+            )
+            print(f'saved to {mpath}')
+            # get_ipython().system(f'gdown {model_id} -O /content/tmp_models/{mkey}')
+            return mpath
+        else:
+            return mpath
+    else:
+        mpath = f'models/{mkey}'
+        return mpath
+
+def prepare_mdx(onnx,custom_param=False, dim_f=None, dim_t=None, n_fft=None, stem_name=None, compensation=None):
+    device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')
+    if custom_param:
+        assert not (dim_f is None or dim_t is None or n_fft is None or compensation is None), 'Custom parameter selected, but incomplete parameters are provided.'
+        mdx_model = mdx.MDX_Model(
+            device,
+            dim_f = dim_f,
+            dim_t = dim_t,
+            n_fft = n_fft,
+            stem_name=stem_name,
+            compensation=compensation
+        )
+    else:
+        model_hash = mdx.MDX.get_hash(onnx)
+        if model_hash in model_params:
+            mp = model_params.get(model_hash)
+            mdx_model = mdx.MDX_Model(
+                device,
+                dim_f = mp["mdx_dim_f_set"],
+                dim_t = 2**mp["mdx_dim_t_set"],
+                n_fft = mp["mdx_n_fft_scale_set"],
+                stem_name=mp["primary_stem"],
+                compensation=compensation if not custom_param and compensation is not None else mp["compensate"]
+            )
+    return mdx_model
+
+def run_mdx(onnx, mdx_model,filename, output_format='wav',diff=False,suffix=None,diff_suffix=None, denoise=False, m_threads=2):
+    mdx_sess = mdx.MDX(onnx,mdx_model)
+    print(f"Processing: {filename}")
+    if filename.lower().endswith('.wav'):
+        wave, sr = librosa.load(filename, mono=False, sr=44100)
+    else:
+        temp_wav = 'temp_audio.wav'
+        subprocess.run(['ffmpeg', '-i', filename, '-ar', '44100', '-ac', '2', temp_wav])  # Convert to WAV format
+        wave, sr = librosa.load(temp_wav, mono=False, sr=44100)
+        os.remove(temp_wav)
+    
+    #wave, sr = librosa.load(filename,mono=False, sr=44100)
+    # normalizing input wave gives better output
+    peak = max(np.max(wave), abs(np.min(wave)))
+    wave /= peak
+    if denoise:
+        wave_processed = -(mdx_sess.process_wave(-wave, m_threads)) + (mdx_sess.process_wave(wave, m_threads))
+        wave_processed *= 0.5
+    else:
+        wave_processed = mdx_sess.process_wave(wave, m_threads)
+    # return to previous peak
+    wave_processed *= peak
+
+    stem_name = mdx_model.stem_name if suffix is None else suffix # use suffix if provided
+    save_path = os.path.basename(os.path.splitext(filename)[0])
+    #vocals_save_path = os.path.join(vocals_folder, f"{save_path}_{stem_name}.{output_format}")
+    #instrumental_save_path = os.path.join(instrumental_folder, f"{save_path}_{stem_name}.{output_format}")
+    save_path = f"{os.path.basename(os.path.splitext(filename)[0])}_{stem_name}.{output_format}"
+    save_path = os.path.join(
+            'audios',
+            save_path
+        )
+    sf.write(
+        save_path,
+        wave_processed.T,
+        sr
+    )
+
+    print(f'done, saved to: {save_path}')
+
+    if diff:
+        diff_stem_name = stem_naming.get(stem_name) if diff_suffix is None else diff_suffix # use suffix if provided
+        stem_name = f"{stem_name}_diff" if diff_stem_name is None else diff_stem_name
+        save_path = f"{os.path.basename(os.path.splitext(filename)[0])}_{stem_name}.{output_format}"
+        save_path = os.path.join(
+                'audio-others',
+                save_path
+            )
+        sf.write(
+            save_path,
+            (-wave_processed.T*mdx_model.compensation)+wave.T,
+            sr
+        )
+        print(f'invert done, saved to: {save_path}')
+    del mdx_sess, wave_processed, wave
+    gc.collect()
+
+if __name__ == "__main__":
+    print()