rvc/train/extract/extract.py

import os
import sys
import glob
import time
import tqdm
import torch
import torchcrepe
import numpy as np
import concurrent.futures
import multiprocessing as mp

# Zluda
if torch.cuda.is_available() and torch.cuda.get_device_name().endswith("[ZLUDA]"):
    torch.backends.cudnn.enabled = False
    torch.backends.cuda.enable_flash_sdp(False)
    torch.backends.cuda.enable_math_sdp(True)
    torch.backends.cuda.enable_mem_efficient_sdp(False)

now_dir = os.getcwd()
sys.path.append(os.path.join(now_dir))

from rvc.lib.utils import load_audio, load_embedding
from rvc.train.extract.preparing_files import generate_config, generate_filelist
from rvc.lib.predictors.RMVPE import RMVPE0Predictor
from rvc.configs.config import Config

# Load config
config = Config()

mp.set_start_method("spawn", force=True)


class FeatureInput:
    """Class for F0 extraction."""

    def __init__(self, sample_rate=16000, hop_size=160, device="cpu"):
        self.fs = sample_rate
        self.hop = hop_size
        self.f0_bin = 256
        self.f0_max = 1100.0
        self.f0_min = 50.0
        self.f0_mel_min = 1127 * np.log(1 + self.f0_min / 700)
        self.f0_mel_max = 1127 * np.log(1 + self.f0_max / 700)
        self.device = device
        self.model_rmvpe = None

    def compute_f0(self, np_arr, f0_method, hop_length):
        """Extract F0 using the specified method."""
        if f0_method == "crepe":
            return self.get_crepe(np_arr, hop_length)
        elif f0_method == "rmvpe":
            return self.model_rmvpe.infer_from_audio(np_arr, thred=0.03)
        else:
            raise ValueError(f"Unknown F0 method: {f0_method}")

    def get_crepe(self, x, hop_length):
        """Extract F0 using CREPE."""
        audio = torch.from_numpy(x.astype(np.float32)).to(self.device)
        audio /= torch.quantile(torch.abs(audio), 0.999)
        audio = audio.unsqueeze(0)
        pitch = torchcrepe.predict(
            audio,
            self.fs,
            hop_length,
            self.f0_min,
            self.f0_max,
            "full",
            batch_size=hop_length * 2,
            device=audio.device,
            pad=True,
        )
        source = pitch.squeeze(0).cpu().float().numpy()
        source[source < 0.001] = np.nan
        target = np.interp(
            np.arange(0, len(source) * (x.size // self.hop), len(source))
            / (x.size // self.hop),
            np.arange(0, len(source)),
            source,
        )
        return np.nan_to_num(target)

    def coarse_f0(self, f0):
        """Convert F0 to coarse F0."""
        f0_mel = 1127 * np.log(1 + f0 / 700)
        f0_mel = np.clip(
            (f0_mel - self.f0_mel_min)
            * (self.f0_bin - 2)
            / (self.f0_mel_max - self.f0_mel_min)
            + 1,
            1,
            self.f0_bin - 1,
        )
        return np.rint(f0_mel).astype(int)

    def process_file(self, file_info, f0_method, hop_length):
        """Process a single audio file for F0 extraction."""
        inp_path, opt_path1, opt_path2, _ = file_info

        if os.path.exists(opt_path1) and os.path.exists(opt_path2):
            return

        try:
            np_arr = load_audio(inp_path, 16000)
            feature_pit = self.compute_f0(np_arr, f0_method, hop_length)
            np.save(opt_path2, feature_pit, allow_pickle=False)
            coarse_pit = self.coarse_f0(feature_pit)
            np.save(opt_path1, coarse_pit, allow_pickle=False)
        except Exception as error:
            print(
                f"An error occurred extracting file {inp_path} on {self.device}: {error}"
            )

    def process_files(
        self, files, f0_method, hop_length, device_num, device, n_threads
    ):
        """Process multiple files."""
        self.device = device
        if f0_method == "rmvpe":
            self.model_rmvpe = RMVPE0Predictor(
                os.path.join("rvc", "models", "predictors", "rmvpe.pt"),
                is_half=False,
                device=device,
            )
        else:
            n_threads = 1

        n_threads = 1 if n_threads == 0 else n_threads

        def process_file_wrapper(file_info):
            self.process_file(file_info, f0_method, hop_length)

        with tqdm.tqdm(total=len(files), leave=True, position=device_num) as pbar:
            # using multi-threading
            with concurrent.futures.ThreadPoolExecutor(
                max_workers=n_threads
            ) as executor:
                futures = [
                    executor.submit(process_file_wrapper, file_info)
                    for file_info in files
                ]
                for future in concurrent.futures.as_completed(futures):
                    pbar.update(1)


def run_pitch_extraction(files, devices, f0_method, hop_length, num_processes):
    devices_str = ", ".join(devices)
    print(
        f"Starting pitch extraction with {num_processes} cores on {devices_str} using {f0_method}..."
    )
    start_time = time.time()
    fe = FeatureInput()
    # split the task between devices
    ps = []
    num_devices = len(devices)
    for i, device in enumerate(devices):
        p = mp.Process(
            target=fe.process_files,
            args=(
                files[i::num_devices],
                f0_method,
                hop_length,
                i,
                device,
                num_processes // num_devices,
            ),
        )
        ps.append(p)
        p.start()
    for i, device in enumerate(devices):
        ps[i].join()

    elapsed_time = time.time() - start_time
    print(f"Pitch extraction completed in {elapsed_time:.2f} seconds.")


def process_file_embedding(
    files, version, embedder_model, embedder_model_custom, device_num, device, n_threads
):
    dtype = torch.float16 if config.is_half and "cuda" in device else torch.float32
    model = load_embedding(embedder_model, embedder_model_custom).to(dtype).to(device)
    n_threads = 1 if n_threads == 0 else n_threads

    def process_file_embedding_wrapper(file_info):
        wav_file_path, _, _, out_file_path = file_info
        if os.path.exists(out_file_path):
            return
        feats = torch.from_numpy(load_audio(wav_file_path, 16000)).to(dtype).to(device)
        feats = feats.view(1, -1)
        with torch.no_grad():
            feats = model(feats)["last_hidden_state"]
            feats = (
                model.final_proj(feats[0]).unsqueeze(0) if version == "v1" else feats
            )
        feats = feats.squeeze(0).float().cpu().numpy()
        if not np.isnan(feats).any():
            np.save(out_file_path, feats, allow_pickle=False)
        else:
            print(f"{file} contains NaN values and will be skipped.")

    with tqdm.tqdm(total=len(files), leave=True, position=device_num) as pbar:
        # using multi-threading
        with concurrent.futures.ThreadPoolExecutor(max_workers=n_threads) as executor:
            futures = [
                executor.submit(process_file_embedding_wrapper, file_info)
                for file_info in files
            ]
            for future in concurrent.futures.as_completed(futures):
                pbar.update(1)


def run_embedding_extraction(
    files, devices, version, embedder_model, embedder_model_custom
):
    start_time = time.time()
    devices_str = ", ".join(devices)
    print(
        f"Starting embedding extraction with {num_processes} cores on {devices_str}..."
    )
    # split the task between devices
    ps = []
    num_devices = len(devices)
    for i, device in enumerate(devices):
        p = mp.Process(
            target=process_file_embedding,
            args=(
                files[i::num_devices],
                version,
                embedder_model,
                embedder_model_custom,
                i,
                device,
                num_processes // num_devices,
            ),
        )
        ps.append(p)
        p.start()
    for i, device in enumerate(devices):
        ps[i].join()
    elapsed_time = time.time() - start_time
    print(f"Embedding extraction completed in {elapsed_time:.2f} seconds.")


if __name__ == "__main__":

    exp_dir = sys.argv[1]
    f0_method = sys.argv[2]
    hop_length = int(sys.argv[3])
    num_processes = int(sys.argv[4])
    gpus = sys.argv[5]
    version = sys.argv[6]
    pitch_guidance = sys.argv[7]
    sample_rate = sys.argv[8]
    embedder_model = sys.argv[9]
    embedder_model_custom = sys.argv[10] if len(sys.argv) > 10 else None

    # prep
    wav_path = os.path.join(exp_dir, "sliced_audios_16k")
    os.makedirs(os.path.join(exp_dir, "f0"), exist_ok=True)
    os.makedirs(os.path.join(exp_dir, "f0_voiced"), exist_ok=True)
    os.makedirs(os.path.join(exp_dir, version + "_extracted"), exist_ok=True)

    files = []
    for file in glob.glob(os.path.join(wav_path, "*.wav")):
        file_name = os.path.basename(file)
        file_info = [
            file,  # full path to sliced 16k wav
            os.path.join(exp_dir, "f0", file_name + ".npy"),
            os.path.join(exp_dir, "f0_voiced", file_name + ".npy"),
            os.path.join(
                exp_dir, version + "_extracted", file_name.replace("wav", "npy")
            ),
        ]
        files.append(file_info)

    devices = ["cpu"] if gpus == "-" else [f"cuda:{idx}" for idx in gpus.split("-")]
    # Run Pitch Extraction
    run_pitch_extraction(files, devices, f0_method, hop_length, num_processes)

    # Run Embedding Extraction
    run_embedding_extraction(
        files, devices, version, embedder_model, embedder_model_custom
    )

    # Run Preparing Files
    generate_config(version, sample_rate, exp_dir)
    generate_filelist(pitch_guidance, exp_dir, version, sample_rate)