Update inference.py

inference.py

@@ -0,0 +1,185 @@
+# inference.py
+
+import os
+import gc
+import json
+import shlex
+import sys
+import torch
+import librosa
+import numpy as np
+import subprocess
+import soundfile as sf
+import hashlib
+import random
+import time
+import traceback
+import onnxruntime as ort
+from utils import logger, remove_directory_contents, create_directories
+from mdx_core import MDX, MDXModel
+from effects import add_vocal_effects, add_instrumental_effects
+
+
+stem_naming = {
+    "Vocals": "Instrumental",
+    "Other": "Instruments",
+    "Instrumental": "Vocals",
+    "Drums": "Drumless",
+    "Bass": "Bassless",
+}
+
+
+def run_mdx(model_params, output_dir, model_path, filename, exclude_main=False, exclude_inversion=False,
+            suffix=None, invert_suffix=None, denoise=False, keep_orig=True, m_threads=2, device_base="cuda"):
+
+    device = torch.device("cuda:0" if device_base == "cuda" else "cpu")
+    processor_num = 0 if device_base == "cuda" else -1
+
+    if device_base == "cuda":
+        vram_gb = torch.cuda.get_device_properties(device).total_memory / 1024**3
+        m_threads = 1 if vram_gb < 8 else (8 if vram_gb > 32 else 2)
+        logger.info(f"threads: {m_threads} vram: {vram_gb}")
+    else:
+        m_threads = 1
+
+    model_hash = MDX.get_hash(model_path)
+    mp = model_params.get(model_hash)
+
+    model = MDXModel(
+        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=mp["compensate"],
+    )
+
+    mdx_sess = MDX(model_path, model, processor=processor_num)
+    wave, sr = librosa.load(filename, mono=False, sr=44100)
+    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)
+
+    wave_processed *= peak
+    stem_name = model.stem_name if suffix is None else suffix
+
+    main_filepath = None
+    if not exclude_main:
+        main_filepath = os.path.join(
+            output_dir,
+            f"{os.path.basename(os.path.splitext(filename)[0])}_{stem_name}.wav",
+        )
+        sf.write(main_filepath, wave_processed.T, sr)
+
+    invert_filepath = None
+    if not exclude_inversion:
+        diff_stem_name = stem_naming.get(stem_name) if invert_suffix is None else invert_suffix
+        stem_name = f"{stem_name}_diff" if diff_stem_name is None else diff_stem_name
+        invert_filepath = os.path.join(
+            output_dir,
+            f"{os.path.basename(os.path.splitext(filename)[0])}_{stem_name}.wav",
+        )
+        sf.write(invert_filepath, (-wave_processed.T * model.compensation) + wave.T, sr)
+
+    if not keep_orig:
+        os.remove(filename)
+
+    del mdx_sess, wave_processed, wave
+    gc.collect()
+    torch.cuda.empty_cache()
+    return main_filepath, invert_filepath
+
+
+def run_mdx_beta(model_params, output_dir, model_path, filename, exclude_main=False, exclude_inversion=False,
+                 suffix=None, invert_suffix=None, denoise=False, keep_orig=True, m_threads=1, device_base=""):
+
+    duration = librosa.get_duration(filename=filename)
+    if duration >= 60 and duration <= 120:
+        m_threads = 8
+    elif duration > 120:
+        m_threads = 16
+
+    logger.info(f"threads: {m_threads}")
+
+    device = torch.device("cpu")
+    processor_num = -1
+
+    model_hash = MDX.get_hash(model_path)
+    mp = model_params.get(model_hash)
+
+    model = MDXModel(
+        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=mp["compensate"],
+    )
+
+    mdx_sess = MDX(model_path, model, processor=processor_num)
+    wave, sr = librosa.load(filename, mono=False, sr=44100)
+    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)
+
+    wave_processed *= peak
+    stem_name = model.stem_name if suffix is None else suffix
+
+    main_filepath = None
+    if not exclude_main:
+        main_filepath = os.path.join(
+            output_dir,
+            f"{os.path.basename(os.path.splitext(filename)[0])}_{stem_name}.wav",
+        )
+        sf.write(main_filepath, wave_processed.T, sr)
+
+    invert_filepath = None
+    if not exclude_inversion:
+        diff_stem_name = stem_naming.get(stem_name) if invert_suffix is None else invert_suffix
+        stem_name = f"{stem_name}_diff" if diff_stem_name is None else diff_stem_name
+        invert_filepath = os.path.join(
+            output_dir,
+            f"{os.path.basename(os.path.splitext(filename)[0])}_{stem_name}.wav",
+        )
+        sf.write(invert_filepath, (-wave_processed.T * model.compensation) + wave.T, sr)
+
+    if not keep_orig:
+        os.remove(filename)
+
+    del mdx_sess, wave_processed, wave
+    gc.collect()
+    torch.cuda.empty_cache()
+    return main_filepath, invert_filepath
+
+
+def convert_to_stereo_and_wav(audio_path, output_dir):
+    wave, sr = librosa.load(audio_path, mono=False, sr=44100)
+
+    if type(wave[0]) != np.ndarray or audio_path[-4:].lower() != ".wav":
+        stereo_path = os.path.join(output_dir, f"{os.path.splitext(os.path.basename(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, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True)
+        return stereo_path
+    return audio_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()[:18]
+
+
+def random_sleep():
+    time.sleep(round(random.uniform(5.2, 7.9), 1))