Spaces:
Running
on
Zero
Running
on
Zero
| from soni_translate.logging_setup import logger | |
| import torch | |
| import gc | |
| import numpy as np | |
| import os | |
| import shutil | |
| import warnings | |
| import threading | |
| from tqdm import tqdm | |
| from lib.infer_pack.models import ( | |
| SynthesizerTrnMs256NSFsid, | |
| SynthesizerTrnMs256NSFsid_nono, | |
| SynthesizerTrnMs768NSFsid, | |
| SynthesizerTrnMs768NSFsid_nono, | |
| ) | |
| from lib.audio import load_audio | |
| import soundfile as sf | |
| import edge_tts | |
| import asyncio | |
| from soni_translate.utils import remove_directory_contents, create_directories | |
| from scipy import signal | |
| from time import time as ttime | |
| import faiss | |
| from vci_pipeline import VC, change_rms, bh, ah | |
| import librosa | |
| warnings.filterwarnings("ignore") | |
| class Config: | |
| def __init__(self, only_cpu=False): | |
| self.device = "cuda:0" | |
| self.is_half = True | |
| self.n_cpu = 0 | |
| self.gpu_name = None | |
| self.gpu_mem = None | |
| ( | |
| self.x_pad, | |
| self.x_query, | |
| self.x_center, | |
| self.x_max | |
| ) = self.device_config(only_cpu) | |
| def device_config(self, only_cpu) -> tuple: | |
| if torch.cuda.is_available() and not only_cpu: | |
| i_device = int(self.device.split(":")[-1]) | |
| self.gpu_name = torch.cuda.get_device_name(i_device) | |
| if ( | |
| ("16" in self.gpu_name and "V100" not in self.gpu_name.upper()) | |
| or "P40" in self.gpu_name.upper() | |
| or "1060" in self.gpu_name | |
| or "1070" in self.gpu_name | |
| or "1080" in self.gpu_name | |
| ): | |
| logger.info( | |
| "16/10 Series GPUs and P40 excel " | |
| "in single-precision tasks." | |
| ) | |
| self.is_half = False | |
| else: | |
| self.gpu_name = None | |
| self.gpu_mem = int( | |
| torch.cuda.get_device_properties(i_device).total_memory | |
| / 1024 | |
| / 1024 | |
| / 1024 | |
| + 0.4 | |
| ) | |
| elif torch.backends.mps.is_available() and not only_cpu: | |
| logger.info("Supported N-card not found, using MPS for inference") | |
| self.device = "mps" | |
| else: | |
| logger.info("No supported N-card found, using CPU for inference") | |
| self.device = "cpu" | |
| self.is_half = False | |
| if self.n_cpu == 0: | |
| self.n_cpu = os.cpu_count() | |
| if self.is_half: | |
| # 6GB VRAM configuration | |
| x_pad = 3 | |
| x_query = 10 | |
| x_center = 60 | |
| x_max = 65 | |
| else: | |
| # 5GB VRAM configuration | |
| x_pad = 1 | |
| x_query = 6 | |
| x_center = 38 | |
| x_max = 41 | |
| if self.gpu_mem is not None and self.gpu_mem <= 4: | |
| x_pad = 1 | |
| x_query = 5 | |
| x_center = 30 | |
| x_max = 32 | |
| logger.info( | |
| f"Config: Device is {self.device}, " | |
| f"half precision is {self.is_half}" | |
| ) | |
| return x_pad, x_query, x_center, x_max | |
| BASE_DOWNLOAD_LINK = "https://huggingface.co/r3gm/sonitranslate_voice_models/resolve/main/" | |
| BASE_MODELS = [ | |
| "hubert_base.pt", | |
| "rmvpe.pt" | |
| ] | |
| BASE_DIR = "." | |
| def load_hu_bert(config): | |
| from fairseq import checkpoint_utils | |
| from soni_translate.utils import download_manager | |
| for id_model in BASE_MODELS: | |
| download_manager( | |
| os.path.join(BASE_DOWNLOAD_LINK, id_model), BASE_DIR | |
| ) | |
| models, _, _ = checkpoint_utils.load_model_ensemble_and_task( | |
| ["hubert_base.pt"], | |
| suffix="", | |
| ) | |
| hubert_model = models[0] | |
| hubert_model = hubert_model.to(config.device) | |
| if config.is_half: | |
| hubert_model = hubert_model.half() | |
| else: | |
| hubert_model = hubert_model.float() | |
| hubert_model.eval() | |
| return hubert_model | |
| def load_trained_model(model_path, config): | |
| if not model_path: | |
| raise ValueError("No model found") | |
| logger.info("Loading %s" % model_path) | |
| cpt = torch.load(model_path, map_location="cpu") | |
| tgt_sr = cpt["config"][-1] | |
| cpt["config"][-3] = cpt["weight"]["emb_g.weight"].shape[0] # n_spk | |
| if_f0 = cpt.get("f0", 1) | |
| if if_f0 == 0: | |
| # protect to 0.5 need? | |
| pass | |
| version = cpt.get("version", "v1") | |
| if version == "v1": | |
| if if_f0 == 1: | |
| net_g = SynthesizerTrnMs256NSFsid( | |
| *cpt["config"], is_half=config.is_half | |
| ) | |
| else: | |
| net_g = SynthesizerTrnMs256NSFsid_nono(*cpt["config"]) | |
| elif version == "v2": | |
| if if_f0 == 1: | |
| net_g = SynthesizerTrnMs768NSFsid( | |
| *cpt["config"], is_half=config.is_half | |
| ) | |
| else: | |
| net_g = SynthesizerTrnMs768NSFsid_nono(*cpt["config"]) | |
| del net_g.enc_q | |
| net_g.load_state_dict(cpt["weight"], strict=False) | |
| net_g.eval().to(config.device) | |
| if config.is_half: | |
| net_g = net_g.half() | |
| else: | |
| net_g = net_g.float() | |
| vc = VC(tgt_sr, config) | |
| n_spk = cpt["config"][-3] | |
| return n_spk, tgt_sr, net_g, vc, cpt, version | |
| class ClassVoices: | |
| def __init__(self, only_cpu=False): | |
| self.model_config = {} | |
| self.config = None | |
| self.only_cpu = only_cpu | |
| def apply_conf( | |
| self, | |
| tag="base_model", | |
| file_model="", | |
| pitch_algo="pm", | |
| pitch_lvl=0, | |
| file_index="", | |
| index_influence=0.66, | |
| respiration_median_filtering=3, | |
| envelope_ratio=0.25, | |
| consonant_breath_protection=0.33, | |
| resample_sr=0, | |
| file_pitch_algo="", | |
| ): | |
| if not file_model: | |
| raise ValueError("Model not found") | |
| if file_index is None: | |
| file_index = "" | |
| if file_pitch_algo is None: | |
| file_pitch_algo = "" | |
| if not self.config: | |
| self.config = Config(self.only_cpu) | |
| self.hu_bert_model = None | |
| self.model_pitch_estimator = None | |
| self.model_config[tag] = { | |
| "file_model": file_model, | |
| "pitch_algo": pitch_algo, | |
| "pitch_lvl": pitch_lvl, # no decimal | |
| "file_index": file_index, | |
| "index_influence": index_influence, | |
| "respiration_median_filtering": respiration_median_filtering, | |
| "envelope_ratio": envelope_ratio, | |
| "consonant_breath_protection": consonant_breath_protection, | |
| "resample_sr": resample_sr, | |
| "file_pitch_algo": file_pitch_algo, | |
| } | |
| return f"CONFIGURATION APPLIED FOR {tag}: {file_model}" | |
| def infer( | |
| self, | |
| task_id, | |
| params, | |
| # load model | |
| n_spk, | |
| tgt_sr, | |
| net_g, | |
| pipe, | |
| cpt, | |
| version, | |
| if_f0, | |
| # load index | |
| index_rate, | |
| index, | |
| big_npy, | |
| # load f0 file | |
| inp_f0, | |
| # audio file | |
| input_audio_path, | |
| overwrite, | |
| ): | |
| f0_method = params["pitch_algo"] | |
| f0_up_key = params["pitch_lvl"] | |
| filter_radius = params["respiration_median_filtering"] | |
| resample_sr = params["resample_sr"] | |
| rms_mix_rate = params["envelope_ratio"] | |
| protect = params["consonant_breath_protection"] | |
| if not os.path.exists(input_audio_path): | |
| raise ValueError( | |
| "The audio file was not found or is not " | |
| f"a valid file: {input_audio_path}" | |
| ) | |
| f0_up_key = int(f0_up_key) | |
| audio = load_audio(input_audio_path, 16000) | |
| # Normalize audio | |
| audio_max = np.abs(audio).max() / 0.95 | |
| if audio_max > 1: | |
| audio /= audio_max | |
| times = [0, 0, 0] | |
| # filters audio signal, pads it, computes sliding window sums, | |
| # and extracts optimized time indices | |
| audio = signal.filtfilt(bh, ah, audio) | |
| audio_pad = np.pad( | |
| audio, (pipe.window // 2, pipe.window // 2), mode="reflect" | |
| ) | |
| opt_ts = [] | |
| if audio_pad.shape[0] > pipe.t_max: | |
| audio_sum = np.zeros_like(audio) | |
| for i in range(pipe.window): | |
| audio_sum += audio_pad[i:i - pipe.window] | |
| for t in range(pipe.t_center, audio.shape[0], pipe.t_center): | |
| opt_ts.append( | |
| t | |
| - pipe.t_query | |
| + np.where( | |
| np.abs(audio_sum[t - pipe.t_query: t + pipe.t_query]) | |
| == np.abs(audio_sum[t - pipe.t_query: t + pipe.t_query]).min() | |
| )[0][0] | |
| ) | |
| s = 0 | |
| audio_opt = [] | |
| t = None | |
| t1 = ttime() | |
| sid_value = 0 | |
| sid = torch.tensor(sid_value, device=pipe.device).unsqueeze(0).long() | |
| # Pads audio symmetrically, calculates length divided by window size. | |
| audio_pad = np.pad(audio, (pipe.t_pad, pipe.t_pad), mode="reflect") | |
| p_len = audio_pad.shape[0] // pipe.window | |
| # Estimates pitch from audio signal | |
| pitch, pitchf = None, None | |
| if if_f0 == 1: | |
| pitch, pitchf = pipe.get_f0( | |
| input_audio_path, | |
| audio_pad, | |
| p_len, | |
| f0_up_key, | |
| f0_method, | |
| filter_radius, | |
| inp_f0, | |
| ) | |
| pitch = pitch[:p_len] | |
| pitchf = pitchf[:p_len] | |
| if pipe.device == "mps": | |
| pitchf = pitchf.astype(np.float32) | |
| pitch = torch.tensor( | |
| pitch, device=pipe.device | |
| ).unsqueeze(0).long() | |
| pitchf = torch.tensor( | |
| pitchf, device=pipe.device | |
| ).unsqueeze(0).float() | |
| t2 = ttime() | |
| times[1] += t2 - t1 | |
| for t in opt_ts: | |
| t = t // pipe.window * pipe.window | |
| if if_f0 == 1: | |
| pitch_slice = pitch[ | |
| :, s // pipe.window: (t + pipe.t_pad2) // pipe.window | |
| ] | |
| pitchf_slice = pitchf[ | |
| :, s // pipe.window: (t + pipe.t_pad2) // pipe.window | |
| ] | |
| else: | |
| pitch_slice = None | |
| pitchf_slice = None | |
| audio_slice = audio_pad[s:t + pipe.t_pad2 + pipe.window] | |
| audio_opt.append( | |
| pipe.vc( | |
| self.hu_bert_model, | |
| net_g, | |
| sid, | |
| audio_slice, | |
| pitch_slice, | |
| pitchf_slice, | |
| times, | |
| index, | |
| big_npy, | |
| index_rate, | |
| version, | |
| protect, | |
| )[pipe.t_pad_tgt:-pipe.t_pad_tgt] | |
| ) | |
| s = t | |
| pitch_end_slice = pitch[ | |
| :, t // pipe.window: | |
| ] if t is not None else pitch | |
| pitchf_end_slice = pitchf[ | |
| :, t // pipe.window: | |
| ] if t is not None else pitchf | |
| audio_opt.append( | |
| pipe.vc( | |
| self.hu_bert_model, | |
| net_g, | |
| sid, | |
| audio_pad[t:], | |
| pitch_end_slice, | |
| pitchf_end_slice, | |
| times, | |
| index, | |
| big_npy, | |
| index_rate, | |
| version, | |
| protect, | |
| )[pipe.t_pad_tgt:-pipe.t_pad_tgt] | |
| ) | |
| audio_opt = np.concatenate(audio_opt) | |
| if rms_mix_rate != 1: | |
| audio_opt = change_rms( | |
| audio, 16000, audio_opt, tgt_sr, rms_mix_rate | |
| ) | |
| if resample_sr >= 16000 and tgt_sr != resample_sr: | |
| audio_opt = librosa.resample( | |
| audio_opt, orig_sr=tgt_sr, target_sr=resample_sr | |
| ) | |
| audio_max = np.abs(audio_opt).max() / 0.99 | |
| max_int16 = 32768 | |
| if audio_max > 1: | |
| max_int16 /= audio_max | |
| audio_opt = (audio_opt * max_int16).astype(np.int16) | |
| del pitch, pitchf, sid | |
| if torch.cuda.is_available(): | |
| torch.cuda.empty_cache() | |
| if tgt_sr != resample_sr >= 16000: | |
| final_sr = resample_sr | |
| else: | |
| final_sr = tgt_sr | |
| """ | |
| "Success.\n %s\nTime:\n npy:%ss, f0:%ss, infer:%ss" % ( | |
| times[0], | |
| times[1], | |
| times[2], | |
| ), (final_sr, audio_opt) | |
| """ | |
| if overwrite: | |
| output_audio_path = input_audio_path # Overwrite | |
| else: | |
| basename = os.path.basename(input_audio_path) | |
| dirname = os.path.dirname(input_audio_path) | |
| new_basename = basename.split( | |
| '.')[0] + "_edited." + basename.split('.')[-1] | |
| new_path = os.path.join(dirname, new_basename) | |
| logger.info(str(new_path)) | |
| output_audio_path = new_path | |
| # Save file | |
| sf.write( | |
| file=output_audio_path, | |
| samplerate=final_sr, | |
| data=audio_opt | |
| ) | |
| self.model_config[task_id]["result"].append(output_audio_path) | |
| self.output_list.append(output_audio_path) | |
| def make_test( | |
| self, | |
| tts_text, | |
| tts_voice, | |
| model_path, | |
| index_path, | |
| transpose, | |
| f0_method, | |
| ): | |
| folder_test = "test" | |
| tag = "test_edge" | |
| tts_file = "test/test.wav" | |
| tts_edited = "test/test_edited.wav" | |
| create_directories(folder_test) | |
| remove_directory_contents(folder_test) | |
| if "SET_LIMIT" == os.getenv("DEMO"): | |
| if len(tts_text) > 60: | |
| tts_text = tts_text[:60] | |
| logger.warning("DEMO; limit to 60 characters") | |
| try: | |
| asyncio.run(edge_tts.Communicate( | |
| tts_text, "-".join(tts_voice.split('-')[:-1]) | |
| ).save(tts_file)) | |
| except Exception as e: | |
| raise ValueError( | |
| "No audio was received. Please change the " | |
| f"tts voice for {tts_voice}. Error: {str(e)}" | |
| ) | |
| shutil.copy(tts_file, tts_edited) | |
| self.apply_conf( | |
| tag=tag, | |
| file_model=model_path, | |
| pitch_algo=f0_method, | |
| pitch_lvl=transpose, | |
| file_index=index_path, | |
| index_influence=0.66, | |
| respiration_median_filtering=3, | |
| envelope_ratio=0.25, | |
| consonant_breath_protection=0.33, | |
| ) | |
| self( | |
| audio_files=tts_edited, | |
| tag_list=tag, | |
| overwrite=True | |
| ) | |
| return tts_edited, tts_file | |
| def run_threads(self, threads): | |
| # Start threads | |
| for thread in threads: | |
| thread.start() | |
| # Wait for all threads to finish | |
| for thread in threads: | |
| thread.join() | |
| gc.collect() | |
| torch.cuda.empty_cache() | |
| def unload_models(self): | |
| self.hu_bert_model = None | |
| self.model_pitch_estimator = None | |
| gc.collect() | |
| torch.cuda.empty_cache() | |
| def __call__( | |
| self, | |
| audio_files=[], | |
| tag_list=[], | |
| overwrite=False, | |
| parallel_workers=1, | |
| ): | |
| logger.info(f"Parallel workers: {str(parallel_workers)}") | |
| self.output_list = [] | |
| if not self.model_config: | |
| raise ValueError("No model has been configured for inference") | |
| if isinstance(audio_files, str): | |
| audio_files = [audio_files] | |
| if isinstance(tag_list, str): | |
| tag_list = [tag_list] | |
| if not audio_files: | |
| raise ValueError("No audio found to convert") | |
| if not tag_list: | |
| tag_list = [list(self.model_config.keys())[-1]] * len(audio_files) | |
| if len(audio_files) > len(tag_list): | |
| logger.info("Extend tag list to match audio files") | |
| extend_number = len(audio_files) - len(tag_list) | |
| tag_list.extend([tag_list[0]] * extend_number) | |
| if len(audio_files) < len(tag_list): | |
| logger.info("Cut list tags") | |
| tag_list = tag_list[:len(audio_files)] | |
| tag_file_pairs = list(zip(tag_list, audio_files)) | |
| sorted_tag_file = sorted(tag_file_pairs, key=lambda x: x[0]) | |
| # Base params | |
| if not self.hu_bert_model: | |
| self.hu_bert_model = load_hu_bert(self.config) | |
| cache_params = None | |
| threads = [] | |
| progress_bar = tqdm(total=len(tag_list), desc="Progress") | |
| for i, (id_tag, input_audio_path) in enumerate(sorted_tag_file): | |
| if id_tag not in self.model_config.keys(): | |
| logger.info( | |
| f"No configured model for {id_tag} with {input_audio_path}" | |
| ) | |
| continue | |
| if ( | |
| len(threads) >= parallel_workers | |
| or cache_params != id_tag | |
| and cache_params is not None | |
| ): | |
| self.run_threads(threads) | |
| progress_bar.update(len(threads)) | |
| threads = [] | |
| if cache_params != id_tag: | |
| self.model_config[id_tag]["result"] = [] | |
| # Unload previous | |
| ( | |
| n_spk, | |
| tgt_sr, | |
| net_g, | |
| pipe, | |
| cpt, | |
| version, | |
| if_f0, | |
| index_rate, | |
| index, | |
| big_npy, | |
| inp_f0, | |
| ) = [None] * 11 | |
| gc.collect() | |
| torch.cuda.empty_cache() | |
| # Model params | |
| params = self.model_config[id_tag] | |
| model_path = params["file_model"] | |
| f0_method = params["pitch_algo"] | |
| file_index = params["file_index"] | |
| index_rate = params["index_influence"] | |
| f0_file = params["file_pitch_algo"] | |
| # Load model | |
| ( | |
| n_spk, | |
| tgt_sr, | |
| net_g, | |
| pipe, | |
| cpt, | |
| version | |
| ) = load_trained_model(model_path, self.config) | |
| if_f0 = cpt.get("f0", 1) # pitch data | |
| # Load index | |
| if os.path.exists(file_index) and index_rate != 0: | |
| try: | |
| index = faiss.read_index(file_index) | |
| big_npy = index.reconstruct_n(0, index.ntotal) | |
| except Exception as error: | |
| logger.error(f"Index: {str(error)}") | |
| index_rate = 0 | |
| index = big_npy = None | |
| else: | |
| logger.warning("File index not found") | |
| index_rate = 0 | |
| index = big_npy = None | |
| # Load f0 file | |
| inp_f0 = None | |
| if os.path.exists(f0_file): | |
| try: | |
| with open(f0_file, "r") as f: | |
| lines = f.read().strip("\n").split("\n") | |
| inp_f0 = [] | |
| for line in lines: | |
| inp_f0.append([float(i) for i in line.split(",")]) | |
| inp_f0 = np.array(inp_f0, dtype="float32") | |
| except Exception as error: | |
| logger.error(f"f0 file: {str(error)}") | |
| if "rmvpe" in f0_method: | |
| if not self.model_pitch_estimator: | |
| from lib.rmvpe import RMVPE | |
| logger.info("Loading vocal pitch estimator model") | |
| self.model_pitch_estimator = RMVPE( | |
| "rmvpe.pt", | |
| is_half=self.config.is_half, | |
| device=self.config.device | |
| ) | |
| pipe.model_rmvpe = self.model_pitch_estimator | |
| cache_params = id_tag | |
| # self.infer( | |
| # id_tag, | |
| # params, | |
| # # load model | |
| # n_spk, | |
| # tgt_sr, | |
| # net_g, | |
| # pipe, | |
| # cpt, | |
| # version, | |
| # if_f0, | |
| # # load index | |
| # index_rate, | |
| # index, | |
| # big_npy, | |
| # # load f0 file | |
| # inp_f0, | |
| # # output file | |
| # input_audio_path, | |
| # overwrite, | |
| # ) | |
| thread = threading.Thread( | |
| target=self.infer, | |
| args=( | |
| id_tag, | |
| params, | |
| # loaded model | |
| n_spk, | |
| tgt_sr, | |
| net_g, | |
| pipe, | |
| cpt, | |
| version, | |
| if_f0, | |
| # loaded index | |
| index_rate, | |
| index, | |
| big_npy, | |
| # loaded f0 file | |
| inp_f0, | |
| # audio file | |
| input_audio_path, | |
| overwrite, | |
| ) | |
| ) | |
| threads.append(thread) | |
| # Run last | |
| if threads: | |
| self.run_threads(threads) | |
| progress_bar.update(len(threads)) | |
| progress_bar.close() | |
| final_result = [] | |
| valid_tags = set(tag_list) | |
| for tag in valid_tags: | |
| if ( | |
| tag in self.model_config.keys() | |
| and "result" in self.model_config[tag].keys() | |
| ): | |
| final_result.extend(self.model_config[tag]["result"]) | |
| return final_result | |