| import os | |
| import gc | |
| import sys | |
| import tqdm | |
| import time | |
| import traceback | |
| import concurrent.futures | |
| import numpy as np | |
| sys.path.append(os.getcwd()) | |
| from main.library.predictors.Generator import Generator | |
| from main.library.utils import load_audio, get_providers | |
| from main.app.variables import config, logger, translations | |
| from main.inference.extracting.setup_path import setup_paths | |
| class FeatureInput: | |
| def __init__(self, is_half=config.is_half, device=config.device): | |
| self.sample_rate = 16000 | |
| self.f0_max = 1100.0 | |
| self.f0_min = 50.0 | |
| self.device = device | |
| self.is_half = is_half | |
| def process_file(self, file_info, f0_method, hop_length, f0_onnx, f0_autotune, f0_autotune_strength): | |
| if not hasattr(self, "f0_gen"): self.f0_gen = Generator(self.sample_rate, hop_length, self.f0_min, self.f0_max, self.is_half, self.device, f0_onnx, False) | |
| inp_path, opt_path1, opt_path2, file_inp = file_info | |
| if os.path.exists(opt_path1 + ".npy") and os.path.exists(opt_path2 + ".npy"): return | |
| try: | |
| pitch, pitchf = self.f0_gen.calculator(config.x_pad, f0_method, load_audio(file_inp, self.sample_rate), 0, None, 0, f0_autotune, f0_autotune_strength, None, False) | |
| np.save(opt_path2, pitchf, allow_pickle=False) | |
| np.save(opt_path1, pitch, allow_pickle=False) | |
| except Exception as e: | |
| logger.info(f"{translations['extract_file_error']} {inp_path}: {e}") | |
| logger.debug(traceback.format_exc()) | |
| def process_files(self, files, f0_method, hop_length, f0_onnx, device, is_half, threads, f0_autotune, f0_autotune_strength): | |
| self.device = device | |
| self.is_half = is_half | |
| def worker(file_info): | |
| self.process_file(file_info, f0_method, hop_length, f0_onnx, f0_autotune, f0_autotune_strength) | |
| with tqdm.tqdm(total=len(files), ncols=100, unit="p", leave=True) as pbar: | |
| with concurrent.futures.ThreadPoolExecutor(max_workers=threads) as executor: | |
| for _ in concurrent.futures.as_completed([executor.submit(worker, f) for f in files]): | |
| pbar.update(1) | |
| def run_pitch_extraction(exp_dir, f0_method, hop_length, num_processes, devices, f0_onnx, is_half, f0_autotune, f0_autotune_strength): | |
| num_processes = max(1, num_processes) | |
| input_root, *output_roots = setup_paths(exp_dir) | |
| output_root1, output_root2 = output_roots if len(output_roots) == 2 else (output_roots[0], None) | |
| logger.info(translations["extract_f0_method"].format(num_processes=num_processes, f0_method=f0_method)) | |
| num_processes = 1 if config.device.startswith("ocl") and ("crepe" in f0_method or "fcpe" in f0_method or "rmvpe" in f0_method or "fcn" in f0_method) else num_processes | |
| paths = [(os.path.join(input_root, name), os.path.join(output_root1, name) if output_root1 else None, os.path.join(output_root2, name) if output_root2 else None, os.path.join(input_root, name)) for name in sorted(os.listdir(input_root)) if "spec" not in name] | |
| start_time = time.time() | |
| feature_input = FeatureInput() | |
| with concurrent.futures.ProcessPoolExecutor(max_workers=len(devices)) as executor: | |
| concurrent.futures.wait([executor.submit(feature_input.process_files, paths[i::len(devices)], f0_method, hop_length, f0_onnx, devices[i], is_half, num_processes // len(devices), f0_autotune, f0_autotune_strength) for i in range(len(devices))]) | |
| gc.collect() | |
| logger.info(translations["extract_f0_success"].format(elapsed_time=f"{(time.time() - start_time):.2f}")) |