Note that any reference audio will be forcefully clipped to 25s if beyond this length.
" + "If total duration of source and reference audio exceeds 30s, source audio will be processed in chunks.
" + "无需训练的 zero-shot 语音/歌声转换模型，若需本地部署查看[GitHub页面](https://github.com/Plachtaa/seed-vc)
" + "请注意，参考音频若超过 25 秒，则会被自动裁剪至此长度。
若源音频和参考音频的总时长超过 30 秒，源音频将被分段处理。") + inputs = [ + gr.Audio(type="filepath", label="Source Audio / 源音频"), + gr.Audio(type="filepath", label="Reference Audio / 参考音频"), + gr.Slider(minimum=1, maximum=200, value=10, step=1, label="Diffusion Steps / 扩散步数", info="10 by default, 50~100 for best quality / 默认为 10，50~100 为最佳质量"), + gr.Slider(minimum=0.5, maximum=2.0, step=0.1, value=1.0, label="Length Adjust / 长度调整", info="<1.0 for speed-up speech, >1.0 for slow-down speech / <1.0 加速语速，>1.0 减慢语速"), + gr.Slider(minimum=0.0, maximum=1.0, step=0.1, value=0.7, label="Inference CFG Rate", info="has subtle influence / 有微小影响"), + gr.Checkbox(label="Use F0 conditioned model / 启用F0输入", value=False, info="Must set to true for singing voice conversion / 歌声转换时必须勾选"), + gr.Checkbox(label="Auto F0 adjust / 自动F0调整", value=True, + info="Roughly adjust F0 to match target voice. Only works when F0 conditioned model is used. / 粗略调整 F0 以匹配目标音色，仅在勾选 '启用F0输入' 时生效"), + gr.Slider(label='Pitch shift / 音调变换', minimum=-24, maximum=24, step=1, value=0, info="Pitch shift in semitones, only works when F0 conditioned model is used / 半音数的音高变换，仅在勾选 '启用F0输入' 时生效"), + ] + + examples = [["examples/source/yae_0.wav", "examples/reference/dingzhen_0.wav", 25, 1.0, 0.7, False, True, 0], + ["examples/source/jay_0.wav", "examples/reference/azuma_0.wav", 25, 1.0, 0.7, True, True, 0], + ["examples/source/Wiz Khalifa,Charlie Puth - See You Again [vocals]_[cut_28sec].wav", + "examples/reference/teio_0.wav", 100, 1.0, 0.7, True, False, 0], + ["examples/source/TECHNOPOLIS - 2085 [vocals]_[cut_14sec].wav", + "examples/reference/trump_0.wav", 50, 1.0, 0.7, True, False, -12], + ] + + outputs = [gr.Audio(label="Stream Output Audio / 流式输出", streaming=True, format='mp3'), + gr.Audio(label="Full Output Audio / 完整输出", streaming=False, format='wav')] + + gr.Interface(fn=voice_conversion, + description=description, + inputs=inputs, + outputs=outputs, + title="Seed Voice Conversion", + examples=examples, + cache_examples=False, + ).launch() \ No newline at end of file diff --git a/app_svc.py b/app_svc.py new file mode 100644 index 0000000000000000000000000000000000000000..3c0f25359f18105591ce61474067071d97b858ed --- /dev/null +++ b/app_svc.py @@ -0,0 +1,437 @@ +import os +os.environ['HF_HUB_CACHE'] = './checkpoints/hf_cache' +import gradio as gr +import torch +import torchaudio +import librosa +from modules.commons import build_model, load_checkpoint, recursive_munch, str2bool +import yaml +from hf_utils import load_custom_model_from_hf +import numpy as np +from pydub import AudioSegment +import argparse +# Load model and configuration +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + +fp16 = False +def load_models(args): + global sr, hop_length, fp16 + fp16 = args.fp16 + print(f"Using device: {device}") + print(f"Using fp16: {fp16}") + # f0 conditioned model + if args.checkpoint_path is None or args.checkpoint_path == "": + dit_checkpoint_path, dit_config_path = load_custom_model_from_hf("Plachta/Seed-VC", + "DiT_seed_v2_uvit_whisper_base_f0_44k_bigvgan_pruned_ft_ema_v2.pth", + "config_dit_mel_seed_uvit_whisper_base_f0_44k.yml") + else: + print(f"Using custom checkpoint: {args.checkpoint_path}") + dit_checkpoint_path = args.checkpoint_path + dit_config_path = args.config_path + config = yaml.safe_load(open(dit_config_path, "r")) + model_params = recursive_munch(config["model_params"]) + model_params.dit_type = 'DiT' + model = build_model(model_params, stage="DiT") + hop_length = config["preprocess_params"]["spect_params"]["hop_length"] + sr = config["preprocess_params"]["sr"] + + # Load checkpoints + model, _, _, _ = load_checkpoint( + model, + None, + dit_checkpoint_path, + load_only_params=True, + ignore_modules=[], + is_distributed=False, + ) + for key in model: + model[key].eval() + model[key].to(device) + model.cfm.estimator.setup_caches(max_batch_size=1, max_seq_length=8192) + + # Load additional modules + from modules.campplus.DTDNN import CAMPPlus + + campplus_ckpt_path = load_custom_model_from_hf( + "funasr/campplus", "campplus_cn_common.bin", config_filename=None + ) + campplus_model = CAMPPlus(feat_dim=80, embedding_size=192) + campplus_model.load_state_dict(torch.load(campplus_ckpt_path, map_location="cpu")) + campplus_model.eval() + campplus_model.to(device) + + vocoder_type = model_params.vocoder.type + + if vocoder_type == 'bigvgan': + from modules.bigvgan import bigvgan + bigvgan_name = model_params.vocoder.name + bigvgan_model = bigvgan.BigVGAN.from_pretrained(bigvgan_name, use_cuda_kernel=False) + # remove weight norm in the model and set to eval mode + bigvgan_model.remove_weight_norm() + bigvgan_model = bigvgan_model.eval().to(device) + vocoder_fn = bigvgan_model + elif vocoder_type == 'hifigan': + from modules.hifigan.generator import HiFTGenerator + from modules.hifigan.f0_predictor import ConvRNNF0Predictor + hift_config = yaml.safe_load(open('configs/hifigan.yml', 'r')) + hift_gen = HiFTGenerator(**hift_config['hift'], f0_predictor=ConvRNNF0Predictor(**hift_config['f0_predictor'])) + hift_path = load_custom_model_from_hf("FunAudioLLM/CosyVoice-300M", 'hift.pt', None) + hift_gen.load_state_dict(torch.load(hift_path, map_location='cpu')) + hift_gen.eval() + hift_gen.to(device) + vocoder_fn = hift_gen + elif vocoder_type == "vocos": + vocos_config = yaml.safe_load(open(model_params.vocoder.vocos.config, 'r')) + vocos_path = model_params.vocoder.vocos.path + vocos_model_params = recursive_munch(vocos_config['model_params']) + vocos = build_model(vocos_model_params, stage='mel_vocos') + vocos_checkpoint_path = vocos_path + vocos, _, _, _ = load_checkpoint(vocos, None, vocos_checkpoint_path, + load_only_params=True, ignore_modules=[], is_distributed=False) + _ = [vocos[key].eval().to(device) for key in vocos] + _ = [vocos[key].to(device) for key in vocos] + total_params = sum(sum(p.numel() for p in vocos[key].parameters() if p.requires_grad) for key in vocos.keys()) + print(f"Vocoder model total parameters: {total_params / 1_000_000:.2f}M") + vocoder_fn = vocos.decoder + else: + raise ValueError(f"Unknown vocoder type: {vocoder_type}") + + speech_tokenizer_type = model_params.speech_tokenizer.type + if speech_tokenizer_type == 'whisper': + # whisper + from transformers import AutoFeatureExtractor, WhisperModel + whisper_name = model_params.speech_tokenizer.name + whisper_model = WhisperModel.from_pretrained(whisper_name, torch_dtype=torch.float16).to(device) + del whisper_model.decoder + whisper_feature_extractor = AutoFeatureExtractor.from_pretrained(whisper_name) + + def semantic_fn(waves_16k): + ori_inputs = whisper_feature_extractor([waves_16k.squeeze(0).cpu().numpy()], + return_tensors="pt", + return_attention_mask=True) + ori_input_features = whisper_model._mask_input_features( + ori_inputs.input_features, attention_mask=ori_inputs.attention_mask).to(device) + with torch.no_grad(): + ori_outputs = whisper_model.encoder( + ori_input_features.to(whisper_model.encoder.dtype), + head_mask=None, + output_attentions=False, + output_hidden_states=False, + return_dict=True, + ) + S_ori = ori_outputs.last_hidden_state.to(torch.float32) + S_ori = S_ori[:, :waves_16k.size(-1) // 320 + 1] + return S_ori + elif speech_tokenizer_type == 'cnhubert': + from transformers import ( + Wav2Vec2FeatureExtractor, + HubertModel, + ) + hubert_model_name = config['model_params']['speech_tokenizer']['name'] + hubert_feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(hubert_model_name) + hubert_model = HubertModel.from_pretrained(hubert_model_name) + hubert_model = hubert_model.to(device) + hubert_model = hubert_model.eval() + hubert_model = hubert_model.half() + + def semantic_fn(waves_16k): + ori_waves_16k_input_list = [ + waves_16k[bib].cpu().numpy() + for bib in range(len(waves_16k)) + ] + ori_inputs = hubert_feature_extractor(ori_waves_16k_input_list, + return_tensors="pt", + return_attention_mask=True, + padding=True, + sampling_rate=16000).to(device) + with torch.no_grad(): + ori_outputs = hubert_model( + ori_inputs.input_values.half(), + ) + S_ori = ori_outputs.last_hidden_state.float() + return S_ori + elif speech_tokenizer_type == 'xlsr': + from transformers import ( + Wav2Vec2FeatureExtractor, + Wav2Vec2Model, + ) + model_name = config['model_params']['speech_tokenizer']['name'] + output_layer = config['model_params']['speech_tokenizer']['output_layer'] + wav2vec_feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(model_name) + wav2vec_model = Wav2Vec2Model.from_pretrained(model_name) + wav2vec_model.encoder.layers = wav2vec_model.encoder.layers[:output_layer] + wav2vec_model = wav2vec_model.to(device) + wav2vec_model = wav2vec_model.eval() + wav2vec_model = wav2vec_model.half() + + def semantic_fn(waves_16k): + ori_waves_16k_input_list = [ + waves_16k[bib].cpu().numpy() + for bib in range(len(waves_16k)) + ] + ori_inputs = wav2vec_feature_extractor(ori_waves_16k_input_list, + return_tensors="pt", + return_attention_mask=True, + padding=True, + sampling_rate=16000).to(device) + with torch.no_grad(): + ori_outputs = wav2vec_model( + ori_inputs.input_values.half(), + ) + S_ori = ori_outputs.last_hidden_state.float() + return S_ori + else: + raise ValueError(f"Unknown speech tokenizer type: {speech_tokenizer_type}") + # Generate mel spectrograms + mel_fn_args = { + "n_fft": config['preprocess_params']['spect_params']['n_fft'], + "win_size": config['preprocess_params']['spect_params']['win_length'], + "hop_size": config['preprocess_params']['spect_params']['hop_length'], + "num_mels": config['preprocess_params']['spect_params']['n_mels'], + "sampling_rate": sr, + "fmin": config['preprocess_params']['spect_params'].get('fmin', 0), + "fmax": None if config['preprocess_params']['spect_params'].get('fmax', "None") == "None" else 8000, + "center": False + } + from modules.audio import mel_spectrogram + + to_mel = lambda x: mel_spectrogram(x, **mel_fn_args) + # f0 extractor + from modules.rmvpe import RMVPE + + model_path = load_custom_model_from_hf("lj1995/VoiceConversionWebUI", "rmvpe.pt", None) + rmvpe = RMVPE(model_path, is_half=False, device=device) + f0_fn = rmvpe.infer_from_audio + + return ( + model, + semantic_fn, + vocoder_fn, + campplus_model, + to_mel, + mel_fn_args, + f0_fn, + ) + +def adjust_f0_semitones(f0_sequence, n_semitones): + factor = 2 ** (n_semitones / 12) + return f0_sequence * factor + +def crossfade(chunk1, chunk2, overlap): + fade_out = np.cos(np.linspace(0, np.pi / 2, overlap)) ** 2 + fade_in = np.cos(np.linspace(np.pi / 2, 0, overlap)) ** 2 + chunk2[:overlap] = chunk2[:overlap] * fade_in + chunk1[-overlap:] * fade_out + return chunk2 + +# streaming and chunk processing related params +# max_context_window = sr // hop_length * 30 +# overlap_frame_len = 16 +# overlap_wave_len = overlap_frame_len * hop_length +bitrate = "320k" + +model_f0, semantic_fn, vocoder_fn, campplus_model, to_mel_f0, mel_fn_args = None, None, None, None, None, None +f0_fn = None +overlap_wave_len = None +max_context_window = None +sr = None +hop_length = None +overlap_frame_len = 16 + +@torch.no_grad() +@torch.inference_mode() +def voice_conversion(source, target, diffusion_steps, length_adjust, inference_cfg_rate, auto_f0_adjust, pitch_shift): + inference_module = model_f0 + mel_fn = to_mel_f0 + # Load audio + source_audio = librosa.load(source, sr=sr)[0] + ref_audio = librosa.load(target, sr=sr)[0] + + # Process audio + source_audio = torch.tensor(source_audio).unsqueeze(0).float().to(device) + ref_audio = torch.tensor(ref_audio[:sr * 25]).unsqueeze(0).float().to(device) + + # Resample + ref_waves_16k = torchaudio.functional.resample(ref_audio, sr, 16000) + converted_waves_16k = torchaudio.functional.resample(source_audio, sr, 16000) + # if source audio less than 30 seconds, whisper can handle in one forward + if converted_waves_16k.size(-1) <= 16000 * 30: + S_alt = semantic_fn(converted_waves_16k) + else: + overlapping_time = 5 # 5 seconds + S_alt_list = [] + buffer = None + traversed_time = 0 + while traversed_time < converted_waves_16k.size(-1): + if buffer is None: # first chunk + chunk = converted_waves_16k[:, traversed_time:traversed_time + 16000 * 30] + else: + chunk = torch.cat([buffer, converted_waves_16k[:, traversed_time:traversed_time + 16000 * (30 - overlapping_time)]], dim=-1) + S_alt = semantic_fn(chunk) + if traversed_time == 0: + S_alt_list.append(S_alt) + else: + S_alt_list.append(S_alt[:, 50 * overlapping_time:]) + buffer = chunk[:, -16000 * overlapping_time:] + traversed_time += 30 * 16000 if traversed_time == 0 else chunk.size(-1) - 16000 * overlapping_time + S_alt = torch.cat(S_alt_list, dim=1) + + ori_waves_16k = torchaudio.functional.resample(ref_audio, sr, 16000) + S_ori = semantic_fn(ori_waves_16k) + + mel = mel_fn(source_audio.to(device).float()) + mel2 = mel_fn(ref_audio.to(device).float()) + + target_lengths = torch.LongTensor([int(mel.size(2) * length_adjust)]).to(mel.device) + target2_lengths = torch.LongTensor([mel2.size(2)]).to(mel2.device) + + feat2 = torchaudio.compliance.kaldi.fbank(ref_waves_16k, + num_mel_bins=80, + dither=0, + sample_frequency=16000) + feat2 = feat2 - feat2.mean(dim=0, keepdim=True) + style2 = campplus_model(feat2.unsqueeze(0)) + + F0_ori = f0_fn(ref_waves_16k[0], thred=0.03) + F0_alt = f0_fn(converted_waves_16k[0], thred=0.03) + + F0_ori = torch.from_numpy(F0_ori).to(device)[None] + F0_alt = torch.from_numpy(F0_alt).to(device)[None] + + voiced_F0_ori = F0_ori[F0_ori > 1] + voiced_F0_alt = F0_alt[F0_alt > 1] + + log_f0_alt = torch.log(F0_alt + 1e-5) + voiced_log_f0_ori = torch.log(voiced_F0_ori + 1e-5) + voiced_log_f0_alt = torch.log(voiced_F0_alt + 1e-5) + median_log_f0_ori = torch.median(voiced_log_f0_ori) + median_log_f0_alt = torch.median(voiced_log_f0_alt) + + # shift alt log f0 level to ori log f0 level + shifted_log_f0_alt = log_f0_alt.clone() + if auto_f0_adjust: + shifted_log_f0_alt[F0_alt > 1] = log_f0_alt[F0_alt > 1] - median_log_f0_alt + median_log_f0_ori + shifted_f0_alt = torch.exp(shifted_log_f0_alt) + if pitch_shift != 0: + shifted_f0_alt[F0_alt > 1] = adjust_f0_semitones(shifted_f0_alt[F0_alt > 1], pitch_shift) + + # Length regulation + cond, _, codes, commitment_loss, codebook_loss = inference_module.length_regulator(S_alt, ylens=target_lengths, n_quantizers=3, f0=shifted_f0_alt) + prompt_condition, _, codes, commitment_loss, codebook_loss = inference_module.length_regulator(S_ori, ylens=target2_lengths, n_quantizers=3, f0=F0_ori) + interpolated_shifted_f0_alt = torch.nn.functional.interpolate(shifted_f0_alt.unsqueeze(1), size=cond.size(1), + mode='nearest').squeeze(1) + max_source_window = max_context_window - mel2.size(2) + # split source condition (cond) into chunks + processed_frames = 0 + generated_wave_chunks = [] + # generate chunk by chunk and stream the output + while processed_frames < cond.size(1): + chunk_cond = cond[:, processed_frames:processed_frames + max_source_window] + chunk_f0 = interpolated_shifted_f0_alt[:, processed_frames:processed_frames + max_source_window] + is_last_chunk = processed_frames + max_source_window >= cond.size(1) + cat_condition = torch.cat([prompt_condition, chunk_cond], dim=1) + with torch.autocast(device_type=device.type, dtype=torch.float16 if fp16 else torch.float32): + # Voice Conversion + vc_target = inference_module.cfm.inference(cat_condition, + torch.LongTensor([cat_condition.size(1)]).to(mel2.device), + mel2, style2, None, diffusion_steps, + inference_cfg_rate=inference_cfg_rate) + vc_target = vc_target[:, :, mel2.size(-1):] + vc_wave = vocoder_fn(vc_target.float()).squeeze().cpu() + if vc_wave.ndim == 1: + vc_wave = vc_wave.unsqueeze(0) + if processed_frames == 0: + if is_last_chunk: + output_wave = vc_wave[0].cpu().numpy() + generated_wave_chunks.append(output_wave) + output_wave = (output_wave * 32768.0).astype(np.int16) + mp3_bytes = AudioSegment( + output_wave.tobytes(), frame_rate=sr, + sample_width=output_wave.dtype.itemsize, channels=1 + ).export(format="mp3", bitrate=bitrate).read() + yield mp3_bytes, (sr, np.concatenate(generated_wave_chunks)) + break + output_wave = vc_wave[0, :-overlap_wave_len].cpu().numpy() + generated_wave_chunks.append(output_wave) + previous_chunk = vc_wave[0, -overlap_wave_len:] + processed_frames += vc_target.size(2) - overlap_frame_len + output_wave = (output_wave * 32768.0).astype(np.int16) + mp3_bytes = AudioSegment( + output_wave.tobytes(), frame_rate=sr, + sample_width=output_wave.dtype.itemsize, channels=1 + ).export(format="mp3", bitrate=bitrate).read() + yield mp3_bytes, None + elif is_last_chunk: + output_wave = crossfade(previous_chunk.cpu().numpy(), vc_wave[0].cpu().numpy(), overlap_wave_len) + generated_wave_chunks.append(output_wave) + processed_frames += vc_target.size(2) - overlap_frame_len + output_wave = (output_wave * 32768.0).astype(np.int16) + mp3_bytes = AudioSegment( + output_wave.tobytes(), frame_rate=sr, + sample_width=output_wave.dtype.itemsize, channels=1 + ).export(format="mp3", bitrate=bitrate).read() + yield mp3_bytes, (sr, np.concatenate(generated_wave_chunks)) + break + else: + output_wave = crossfade(previous_chunk.cpu().numpy(), vc_wave[0, :-overlap_wave_len].cpu().numpy(), overlap_wave_len) + generated_wave_chunks.append(output_wave) + previous_chunk = vc_wave[0, -overlap_wave_len:] + processed_frames += vc_target.size(2) - overlap_frame_len + output_wave = (output_wave * 32768.0).astype(np.int16) + mp3_bytes = AudioSegment( + output_wave.tobytes(), frame_rate=sr, + sample_width=output_wave.dtype.itemsize, channels=1 + ).export(format="mp3", bitrate=bitrate).read() + yield mp3_bytes, None + + +def main(args): + global model_f0, semantic_fn, vocoder_fn, campplus_model, to_mel_f0, mel_fn_args, f0_fn + global overlap_wave_len, max_context_window, sr, hop_length + model_f0, semantic_fn, vocoder_fn, campplus_model, to_mel_f0, mel_fn_args, f0_fn = load_models(args) + # streaming and chunk processing related params + max_context_window = sr // hop_length * 30 + overlap_wave_len = overlap_frame_len * hop_length + description = ("Zero-shot voice conversion with in-context learning. For local deployment please check [GitHub repository](https://github.com/Plachtaa/seed-vc) " + "for details and updates.
Note that any reference audio will be forcefully clipped to 25s if beyond this length.
" + "If total duration of source and reference audio exceeds 30s, source audio will be processed in chunks.
" + "无需训练的 zero-shot 语音/歌声转换模型，若需本地部署查看[GitHub页面](https://github.com/Plachtaa/seed-vc)
" + "请注意，参考音频若超过 25 秒，则会被自动裁剪至此长度。
若源音频和参考音频的总时长超过 30 秒，源音频将被分段处理。") + inputs = [ + gr.Audio(type="filepath", label="Source Audio / 源音频"), + gr.Audio(type="filepath", label="Reference Audio / 参考音频"), + gr.Slider(minimum=1, maximum=200, value=10, step=1, label="Diffusion Steps / 扩散步数", info="10 by default, 50~100 for best quality / 默认为 10，50~100 为最佳质量"), + gr.Slider(minimum=0.5, maximum=2.0, step=0.1, value=1.0, label="Length Adjust / 长度调整", info="<1.0 for speed-up speech, >1.0 for slow-down speech / <1.0 加速语速，>1.0 减慢语速"), + gr.Slider(minimum=0.0, maximum=1.0, step=0.1, value=0.7, label="Inference CFG Rate", info="has subtle influence / 有微小影响"), + gr.Checkbox(label="Auto F0 adjust / 自动F0调整", value=True, + info="Roughly adjust F0 to match target voice. Only works when F0 conditioned model is used. / 粗略调整 F0 以匹配目标音色，仅在勾选 '启用F0输入' 时生效"), + gr.Slider(label='Pitch shift / 音调变换', minimum=-24, maximum=24, step=1, value=0, info="Pitch shift in semitones, only works when F0 conditioned model is used / 半音数的音高变换，仅在勾选 '启用F0输入' 时生效"), + ] + + examples = [["examples/source/yae_0.wav", "examples/reference/dingzhen_0.wav", 25, 1.0, 0.7, True, 0], + ["examples/source/jay_0.wav", "examples/reference/azuma_0.wav", 25, 1.0, 0.7, True, 0], + ["examples/source/Wiz Khalifa,Charlie Puth - See You Again [vocals]_[cut_28sec].wav", + "examples/reference/teio_0.wav", 50, 1.0, 0.7, False, 0], + ["examples/source/TECHNOPOLIS - 2085 [vocals]_[cut_14sec].wav", + "examples/reference/trump_0.wav", 50, 1.0, 0.7, False, -12], + ] + + outputs = [gr.Audio(label="Stream Output Audio / 流式输出", streaming=True, format='mp3'), + gr.Audio(label="Full Output Audio / 完整输出", streaming=False, format='wav')] + + gr.Interface(fn=voice_conversion, + description=description, + inputs=inputs, + outputs=outputs, + title="Seed Voice Conversion", + examples=examples, + cache_examples=False, + ).launch(share=args.share,) + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--checkpoint-path", type=str, help="Path to the checkpoint file", default=None) + parser.add_argument("--config-path", type=str, help="Path to the config file", default=None) + parser.add_argument("--share", type=str2bool, nargs="?", const=True, default=False, help="Whether to share the app") + parser.add_argument("--fp16", type=str2bool, nargs="?", const=True, help="Whether to use fp16", default=True) + args = parser.parse_args() + main(args) diff --git a/app_vc.py b/app_vc.py new file mode 100644 index 0000000000000000000000000000000000000000..1f67945bedbb495815f1eec93cf74e47699ef99b --- /dev/null +++ b/app_vc.py @@ -0,0 +1,390 @@ +import os +os.environ['HF_HUB_CACHE'] = './checkpoints/hf_cache' +import gradio as gr +import torch +import torchaudio +import librosa +from modules.commons import build_model, load_checkpoint, recursive_munch, str2bool +import yaml +from hf_utils import load_custom_model_from_hf +import numpy as np +from pydub import AudioSegment +import argparse + +# Load model and configuration +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") +fp16 = False +def load_models(args): + global sr, hop_length, fp16 + fp16 = args.fp16 + print(f"Using device: {device}") + print(f"Using fp16: {fp16}") + if args.checkpoint_path is None or args.checkpoint_path == "": + dit_checkpoint_path, dit_config_path = load_custom_model_from_hf("Plachta/Seed-VC", + "DiT_seed_v2_uvit_whisper_small_wavenet_bigvgan_pruned.pth", + "config_dit_mel_seed_uvit_whisper_small_wavenet.yml") + else: + dit_checkpoint_path = args.checkpoint_path + dit_config_path = args.config_path + config = yaml.safe_load(open(dit_config_path, "r")) + model_params = recursive_munch(config["model_params"]) + model_params.dit_type = 'DiT' + model = build_model(model_params, stage="DiT") + hop_length = config["preprocess_params"]["spect_params"]["hop_length"] + sr = config["preprocess_params"]["sr"] + + # Load checkpoints + model, _, _, _ = load_checkpoint( + model, + None, + dit_checkpoint_path, + load_only_params=True, + ignore_modules=[], + is_distributed=False, + ) + for key in model: + model[key].eval() + model[key].to(device) + model.cfm.estimator.setup_caches(max_batch_size=1, max_seq_length=8192) + + # Load additional modules + from modules.campplus.DTDNN import CAMPPlus + + campplus_ckpt_path = load_custom_model_from_hf( + "funasr/campplus", "campplus_cn_common.bin", config_filename=None + ) + campplus_model = CAMPPlus(feat_dim=80, embedding_size=192) + campplus_model.load_state_dict(torch.load(campplus_ckpt_path, map_location="cpu")) + campplus_model.eval() + campplus_model.to(device) + + vocoder_type = model_params.vocoder.type + + if vocoder_type == 'bigvgan': + from modules.bigvgan import bigvgan + bigvgan_name = model_params.vocoder.name + bigvgan_model = bigvgan.BigVGAN.from_pretrained(bigvgan_name, use_cuda_kernel=False) + # remove weight norm in the model and set to eval mode + bigvgan_model.remove_weight_norm() + bigvgan_model = bigvgan_model.eval().to(device) + vocoder_fn = bigvgan_model + elif vocoder_type == 'hifigan': + from modules.hifigan.generator import HiFTGenerator + from modules.hifigan.f0_predictor import ConvRNNF0Predictor + hift_config = yaml.safe_load(open('configs/hifigan.yml', 'r')) + hift_gen = HiFTGenerator(**hift_config['hift'], f0_predictor=ConvRNNF0Predictor(**hift_config['f0_predictor'])) + hift_path = load_custom_model_from_hf("FunAudioLLM/CosyVoice-300M", 'hift.pt', None) + hift_gen.load_state_dict(torch.load(hift_path, map_location='cpu')) + hift_gen.eval() + hift_gen.to(device) + vocoder_fn = hift_gen + elif vocoder_type == "vocos": + vocos_config = yaml.safe_load(open(model_params.vocoder.vocos.config, 'r')) + vocos_path = model_params.vocoder.vocos.path + vocos_model_params = recursive_munch(vocos_config['model_params']) + vocos = build_model(vocos_model_params, stage='mel_vocos') + vocos_checkpoint_path = vocos_path + vocos, _, _, _ = load_checkpoint(vocos, None, vocos_checkpoint_path, + load_only_params=True, ignore_modules=[], is_distributed=False) + _ = [vocos[key].eval().to(device) for key in vocos] + _ = [vocos[key].to(device) for key in vocos] + total_params = sum(sum(p.numel() for p in vocos[key].parameters() if p.requires_grad) for key in vocos.keys()) + print(f"Vocoder model total parameters: {total_params / 1_000_000:.2f}M") + vocoder_fn = vocos.decoder + else: + raise ValueError(f"Unknown vocoder type: {vocoder_type}") + + speech_tokenizer_type = model_params.speech_tokenizer.type + if speech_tokenizer_type == 'whisper': + # whisper + from transformers import AutoFeatureExtractor, WhisperModel + whisper_name = model_params.speech_tokenizer.name + whisper_model = WhisperModel.from_pretrained(whisper_name, torch_dtype=torch.float16).to(device) + del whisper_model.decoder + whisper_feature_extractor = AutoFeatureExtractor.from_pretrained(whisper_name) + + def semantic_fn(waves_16k): + ori_inputs = whisper_feature_extractor([waves_16k.squeeze(0).cpu().numpy()], + return_tensors="pt", + return_attention_mask=True) + ori_input_features = whisper_model._mask_input_features( + ori_inputs.input_features, attention_mask=ori_inputs.attention_mask).to(device) + with torch.no_grad(): + ori_outputs = whisper_model.encoder( + ori_input_features.to(whisper_model.encoder.dtype), + head_mask=None, + output_attentions=False, + output_hidden_states=False, + return_dict=True, + ) + S_ori = ori_outputs.last_hidden_state.to(torch.float32) + S_ori = S_ori[:, :waves_16k.size(-1) // 320 + 1] + return S_ori + elif speech_tokenizer_type == 'cnhubert': + from transformers import ( + Wav2Vec2FeatureExtractor, + HubertModel, + ) + hubert_model_name = config['model_params']['speech_tokenizer']['name'] + hubert_feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(hubert_model_name) + hubert_model = HubertModel.from_pretrained(hubert_model_name) + hubert_model = hubert_model.to(device) + hubert_model = hubert_model.eval() + hubert_model = hubert_model.half() + + def semantic_fn(waves_16k): + ori_waves_16k_input_list = [ + waves_16k[bib].cpu().numpy() + for bib in range(len(waves_16k)) + ] + ori_inputs = hubert_feature_extractor(ori_waves_16k_input_list, + return_tensors="pt", + return_attention_mask=True, + padding=True, + sampling_rate=16000).to(device) + with torch.no_grad(): + ori_outputs = hubert_model( + ori_inputs.input_values.half(), + ) + S_ori = ori_outputs.last_hidden_state.float() + return S_ori + elif speech_tokenizer_type == 'xlsr': + from transformers import ( + Wav2Vec2FeatureExtractor, + Wav2Vec2Model, + ) + model_name = config['model_params']['speech_tokenizer']['name'] + output_layer = config['model_params']['speech_tokenizer']['output_layer'] + wav2vec_feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(model_name) + wav2vec_model = Wav2Vec2Model.from_pretrained(model_name) + wav2vec_model.encoder.layers = wav2vec_model.encoder.layers[:output_layer] + wav2vec_model = wav2vec_model.to(device) + wav2vec_model = wav2vec_model.eval() + wav2vec_model = wav2vec_model.half() + + def semantic_fn(waves_16k): + ori_waves_16k_input_list = [ + waves_16k[bib].cpu().numpy() + for bib in range(len(waves_16k)) + ] + ori_inputs = wav2vec_feature_extractor(ori_waves_16k_input_list, + return_tensors="pt", + return_attention_mask=True, + padding=True, + sampling_rate=16000).to(device) + with torch.no_grad(): + ori_outputs = wav2vec_model( + ori_inputs.input_values.half(), + ) + S_ori = ori_outputs.last_hidden_state.float() + return S_ori + else: + raise ValueError(f"Unknown speech tokenizer type: {speech_tokenizer_type}") + # Generate mel spectrograms + mel_fn_args = { + "n_fft": config['preprocess_params']['spect_params']['n_fft'], + "win_size": config['preprocess_params']['spect_params']['win_length'], + "hop_size": config['preprocess_params']['spect_params']['hop_length'], + "num_mels": config['preprocess_params']['spect_params']['n_mels'], + "sampling_rate": sr, + "fmin": config['preprocess_params']['spect_params'].get('fmin', 0), + "fmax": None if config['preprocess_params']['spect_params'].get('fmax', "None") == "None" else 8000, + "center": False + } + from modules.audio import mel_spectrogram + + to_mel = lambda x: mel_spectrogram(x, **mel_fn_args) + + return ( + model, + semantic_fn, + vocoder_fn, + campplus_model, + to_mel, + mel_fn_args, + ) +def crossfade(chunk1, chunk2, overlap): + fade_out = np.cos(np.linspace(0, np.pi / 2, overlap)) ** 2 + fade_in = np.cos(np.linspace(np.pi / 2, 0, overlap)) ** 2 + chunk2[:overlap] = chunk2[:overlap] * fade_in + chunk1[-overlap:] * fade_out + return chunk2 + +bitrate = "320k" + +model, semantic_fn, vocoder_fn, campplus_model, to_mel, mel_fn_args = None, None, None, None, None, None +overlap_wave_len = None +max_context_window = None +sr = None +hop_length = None +overlap_frame_len = 16 +@torch.no_grad() +@torch.inference_mode() +def voice_conversion(source, target, diffusion_steps, length_adjust, inference_cfg_rate): + inference_module = model + mel_fn = to_mel + # Load audio + source_audio = librosa.load(source, sr=sr)[0] + ref_audio = librosa.load(target, sr=sr)[0] + + # Process audio + source_audio = torch.tensor(source_audio).unsqueeze(0).float().to(device) + ref_audio = torch.tensor(ref_audio[:sr * 25]).unsqueeze(0).float().to(device) + + # Resample + ref_waves_16k = torchaudio.functional.resample(ref_audio, sr, 16000) + converted_waves_16k = torchaudio.functional.resample(source_audio, sr, 16000) + # if source audio less than 30 seconds, whisper can handle in one forward + if converted_waves_16k.size(-1) <= 16000 * 30: + S_alt = semantic_fn(converted_waves_16k) + else: + overlapping_time = 5 # 5 seconds + S_alt_list = [] + buffer = None + traversed_time = 0 + while traversed_time < converted_waves_16k.size(-1): + if buffer is None: # first chunk + chunk = converted_waves_16k[:, traversed_time:traversed_time + 16000 * 30] + else: + chunk = torch.cat([buffer, converted_waves_16k[:, traversed_time:traversed_time + 16000 * (30 - overlapping_time)]], dim=-1) + S_alt = semantic_fn(chunk) + if traversed_time == 0: + S_alt_list.append(S_alt) + else: + S_alt_list.append(S_alt[:, 50 * overlapping_time:]) + buffer = chunk[:, -16000 * overlapping_time:] + traversed_time += 30 * 16000 if traversed_time == 0 else chunk.size(-1) - 16000 * overlapping_time + S_alt = torch.cat(S_alt_list, dim=1) + + ori_waves_16k = torchaudio.functional.resample(ref_audio, sr, 16000) + S_ori = semantic_fn(ori_waves_16k) + + mel = mel_fn(source_audio.to(device).float()) + mel2 = mel_fn(ref_audio.to(device).float()) + + target_lengths = torch.LongTensor([int(mel.size(2) * length_adjust)]).to(mel.device) + target2_lengths = torch.LongTensor([mel2.size(2)]).to(mel2.device) + + feat2 = torchaudio.compliance.kaldi.fbank(ref_waves_16k, + num_mel_bins=80, + dither=0, + sample_frequency=16000) + feat2 = feat2 - feat2.mean(dim=0, keepdim=True) + style2 = campplus_model(feat2.unsqueeze(0)) + + F0_ori = None + F0_alt = None + shifted_f0_alt = None + + # Length regulation + cond, _, codes, commitment_loss, codebook_loss = inference_module.length_regulator(S_alt, ylens=target_lengths, n_quantizers=3, f0=shifted_f0_alt) + prompt_condition, _, codes, commitment_loss, codebook_loss = inference_module.length_regulator(S_ori, ylens=target2_lengths, n_quantizers=3, f0=F0_ori) + + max_source_window = max_context_window - mel2.size(2) + # split source condition (cond) into chunks + processed_frames = 0 + generated_wave_chunks = [] + # generate chunk by chunk and stream the output + while processed_frames < cond.size(1): + chunk_cond = cond[:, processed_frames:processed_frames + max_source_window] + is_last_chunk = processed_frames + max_source_window >= cond.size(1) + cat_condition = torch.cat([prompt_condition, chunk_cond], dim=1) + with torch.autocast(device_type=device.type, dtype=torch.float16 if fp16 else torch.float32): + # Voice Conversion + vc_target = inference_module.cfm.inference(cat_condition, + torch.LongTensor([cat_condition.size(1)]).to(mel2.device), + mel2, style2, None, diffusion_steps, + inference_cfg_rate=inference_cfg_rate) + vc_target = vc_target[:, :, mel2.size(-1):] + vc_wave = vocoder_fn(vc_target.float())[0] + if vc_wave.ndim == 1: + vc_wave = vc_wave.unsqueeze(0) + if processed_frames == 0: + if is_last_chunk: + output_wave = vc_wave[0].cpu().numpy() + generated_wave_chunks.append(output_wave) + output_wave = (output_wave * 32768.0).astype(np.int16) + mp3_bytes = AudioSegment( + output_wave.tobytes(), frame_rate=sr, + sample_width=output_wave.dtype.itemsize, channels=1 + ).export(format="mp3", bitrate=bitrate).read() + yield mp3_bytes, (sr, np.concatenate(generated_wave_chunks)) + break + output_wave = vc_wave[0, :-overlap_wave_len].cpu().numpy() + generated_wave_chunks.append(output_wave) + previous_chunk = vc_wave[0, -overlap_wave_len:] + processed_frames += vc_target.size(2) - overlap_frame_len + output_wave = (output_wave * 32768.0).astype(np.int16) + mp3_bytes = AudioSegment( + output_wave.tobytes(), frame_rate=sr, + sample_width=output_wave.dtype.itemsize, channels=1 + ).export(format="mp3", bitrate=bitrate).read() + yield mp3_bytes, None + elif is_last_chunk: + output_wave = crossfade(previous_chunk.cpu().numpy(), vc_wave[0].cpu().numpy(), overlap_wave_len) + generated_wave_chunks.append(output_wave) + processed_frames += vc_target.size(2) - overlap_frame_len + output_wave = (output_wave * 32768.0).astype(np.int16) + mp3_bytes = AudioSegment( + output_wave.tobytes(), frame_rate=sr, + sample_width=output_wave.dtype.itemsize, channels=1 + ).export(format="mp3", bitrate=bitrate).read() + yield mp3_bytes, (sr, np.concatenate(generated_wave_chunks)) + break + else: + output_wave = crossfade(previous_chunk.cpu().numpy(), vc_wave[0, :-overlap_wave_len].cpu().numpy(), overlap_wave_len) + generated_wave_chunks.append(output_wave) + previous_chunk = vc_wave[0, -overlap_wave_len:] + processed_frames += vc_target.size(2) - overlap_frame_len + output_wave = (output_wave * 32768.0).astype(np.int16) + mp3_bytes = AudioSegment( + output_wave.tobytes(), frame_rate=sr, + sample_width=output_wave.dtype.itemsize, channels=1 + ).export(format="mp3", bitrate=bitrate).read() + yield mp3_bytes, None + + +def main(args): + global model, semantic_fn, vocoder_fn, campplus_model, to_mel, mel_fn_args + global overlap_wave_len, max_context_window, sr, hop_length + model, semantic_fn, vocoder_fn, campplus_model, to_mel, mel_fn_args = load_models(args) + # streaming and chunk processing related params + max_context_window = sr // hop_length * 30 + overlap_wave_len = overlap_frame_len * hop_length + description = ("Zero-shot voice conversion with in-context learning. For local deployment please check [GitHub repository](https://github.com/Plachtaa/seed-vc) " + "for details and updates.
Note that any reference audio will be forcefully clipped to 25s if beyond this length.
" + "If total duration of source and reference audio exceeds 30s, source audio will be processed in chunks.
" + "无需训练的 zero-shot 语音/歌声转换模型，若需本地部署查看[GitHub页面](https://github.com/Plachtaa/seed-vc)
" + "请注意，参考音频若超过 25 秒，则会被自动裁剪至此长度。
若源音频和参考音频的总时长超过 30 秒，源音频将被分段处理。") + inputs = [ + gr.Audio(type="filepath", label="Source Audio / 源音频"), + gr.Audio(type="filepath", label="Reference Audio / 参考音频"), + gr.Slider(minimum=1, maximum=200, value=10, step=1, label="Diffusion Steps / 扩散步数", info="10 by default, 50~100 for best quality / 默认为 10，50~100 为最佳质量"), + gr.Slider(minimum=0.5, maximum=2.0, step=0.1, value=1.0, label="Length Adjust / 长度调整", info="<1.0 for speed-up speech, >1.0 for slow-down speech / <1.0 加速语速，>1.0 减慢语速"), + gr.Slider(minimum=0.0, maximum=1.0, step=0.1, value=0.7, label="Inference CFG Rate", info="has subtle influence / 有微小影响"), + ] + + examples = [["examples/source/yae_0.wav", "examples/reference/dingzhen_0.wav", 25, 1.0, 0.7, False, True, 0], + ["examples/source/jay_0.wav", "examples/reference/azuma_0.wav", 25, 1.0, 0.7, True, True, 0], + ] + + outputs = [gr.Audio(label="Stream Output Audio / 流式输出", streaming=True, format='mp3'), + gr.Audio(label="Full Output Audio / 完整输出", streaming=False, format='wav')] + + + gr.Interface(fn=voice_conversion, + description=description, + inputs=inputs, + outputs=outputs, + title="Seed Voice Conversion", + examples=examples, + cache_examples=False, + ).launch(share=args.share,) + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--checkpoint-path", type=str, help="Path to the checkpoint file", default=None) + parser.add_argument("--config-path", type=str, help="Path to the config file", default=None) + parser.add_argument("--share", type=str2bool, nargs="?", const=True, default=False, help="Whether to share the app") + parser.add_argument("--fp16", type=str2bool, nargs="?", const=True, help="Whether to use fp16", default=True) + args = parser.parse_args() + main(args) diff --git a/assets/real-time-demo.webm b/assets/real-time-demo.webm new file mode 100644 index 0000000000000000000000000000000000000000..6144a757e2948db69ba37d761b672e03ac5dd538 --- /dev/null +++ b/assets/real-time-demo.webm @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:1006b20504bf6d83625bea24803765a60dd122261b46ac6cd1bdefecfad697cf +size 3311794 diff --git a/baselines/cosyvoice.py b/baselines/cosyvoice.py new file mode 100644 index 0000000000000000000000000000000000000000..87b7499bc5c33ad9870b19f570ba4b519a494de3 --- /dev/null +++ b/baselines/cosyvoice.py @@ -0,0 +1,24 @@ +import os +import torch +import sys +import librosa +sys.path.append('../CosyVoice') +import sys +sys.path.append("../CosyVoice/third_party/Matcha-TTS") +from cosyvoice.cli.cosyvoice import CosyVoice +from cosyvoice.utils.file_utils import load_wav +import torchaudio +# from modelscope import snapshot_download +# snapshot_download('iic/CosyVoice-300M-25Hz', local_dir='pretrained_models/CosyVoice-300M-25Hz') +cosyvoice = CosyVoice('pretrained_models/CosyVoice-300M-25Hz') +device = "cuda:0" if torch.cuda.is_available() else "cpu" + +@torch.no_grad() +def convert(source_path, reference_path, output_path): + prompt_speech_16k = load_wav(reference_path, 16000) + source_speech_16k = load_wav(source_path, 16000) + + for i in cosyvoice.inference_vc(source_speech_16k, prompt_speech_16k, stream=False): + output_wav_22k = i['tts_speech'] + output_wav_16k = torchaudio.functional.resample(output_wav_22k, 22050, 16000) + return prompt_speech_16k, output_wav_16k \ No newline at end of file diff --git a/baselines/dnsmos/dnsmos_computor.py b/baselines/dnsmos/dnsmos_computor.py new file mode 100644 index 0000000000000000000000000000000000000000..edee496416099131dfc4feb9cbb9dc1cdb93a4b4 --- /dev/null +++ b/baselines/dnsmos/dnsmos_computor.py @@ -0,0 +1,130 @@ +import glob +import librosa +import tqdm +import numpy as np +import torchaudio +import torch + +# ignore all warning +import warnings + +warnings.filterwarnings("ignore") + +import concurrent.futures +import glob +import os +import librosa +import numpy as np +import onnxruntime as ort +import pandas as pd +from tqdm import tqdm + +SAMPLING_RATE = 16000 +INPUT_LENGTH = 9.01 + + +class DNSMOSComputer: + def __init__( + self, primary_model_path, p808_model_path, device="cuda", device_id=0 + ) -> None: + self.onnx_sess = ort.InferenceSession( + primary_model_path, providers=["CUDAExecutionProvider"] + ) + self.p808_onnx_sess = ort.InferenceSession( + p808_model_path, providers=["CUDAExecutionProvider"] + ) + self.onnx_sess.set_providers(["CUDAExecutionProvider"], [{"device_id": device_id}]) + self.p808_onnx_sess.set_providers( + ["CUDAExecutionProvider"], [{"device_id": device_id}] + ) + kwargs = { + "sample_rate": 16000, + "hop_length": 160, + "n_fft": 320 + 1, + "n_mels": 120, + "mel_scale": "slaney", + } + self.mel_transform = torchaudio.transforms.MelSpectrogram(**kwargs).to(f"cuda:{device_id}") + + def audio_melspec( + self, audio, n_mels=120, frame_size=320, hop_length=160, sr=16000, to_db=True + ): + mel_specgram = self.mel_transform(torch.Tensor(audio).cuda()) + mel_spec = mel_specgram.cpu() + if to_db: + mel_spec = (librosa.power_to_db(mel_spec, ref=np.max) + 40) / 40 + return mel_spec.T + + def get_polyfit_val(self, sig, bak, ovr, is_personalized_MOS): + if is_personalized_MOS: + p_ovr = np.poly1d([-0.00533021, 0.005101, 1.18058466, -0.11236046]) + p_sig = np.poly1d([-0.01019296, 0.02751166, 1.19576786, -0.24348726]) + p_bak = np.poly1d([-0.04976499, 0.44276479, -0.1644611, 0.96883132]) + else: + p_ovr = np.poly1d([-0.06766283, 1.11546468, 0.04602535]) + p_sig = np.poly1d([-0.08397278, 1.22083953, 0.0052439]) + p_bak = np.poly1d([-0.13166888, 1.60915514, -0.39604546]) + sig_poly = p_sig(sig) + bak_poly = p_bak(bak) + ovr_poly = p_ovr(ovr) + return sig_poly, bak_poly, ovr_poly + + def compute(self, audio, sampling_rate, is_personalized_MOS=False): + fs = SAMPLING_RATE + if isinstance(audio, str): + audio, _ = librosa.load(audio, sr=fs) + elif sampling_rate != fs: + # resample audio + audio = librosa.resample(audio, orig_sr=sampling_rate, target_sr=fs) + actual_audio_len = len(audio) + len_samples = int(INPUT_LENGTH * fs) + while len(audio) < len_samples: + audio = np.append(audio, audio) + num_hops = int(np.floor(len(audio) / fs) - INPUT_LENGTH) + 1 + hop_len_samples = fs + predicted_mos_sig_seg_raw = [] + predicted_mos_bak_seg_raw = [] + predicted_mos_ovr_seg_raw = [] + predicted_mos_sig_seg = [] + predicted_mos_bak_seg = [] + predicted_mos_ovr_seg = [] + predicted_p808_mos = [] + + for idx in range(num_hops): + audio_seg = audio[ + int(idx * hop_len_samples) : int((idx + INPUT_LENGTH) * hop_len_samples) + ] + if len(audio_seg) < len_samples: + continue + input_features = np.array(audio_seg).astype("float32")[np.newaxis, :] + p808_input_features = np.array( + self.audio_melspec(audio=audio_seg[:-160]) + ).astype("float32")[np.newaxis, :, :] + oi = {"input_1": input_features} + p808_oi = {"input_1": p808_input_features} + p808_mos = self.p808_onnx_sess.run(None, p808_oi)[0][0][0] + mos_sig_raw, mos_bak_raw, mos_ovr_raw = self.onnx_sess.run(None, oi)[0][0] + mos_sig, mos_bak, mos_ovr = self.get_polyfit_val( + mos_sig_raw, mos_bak_raw, mos_ovr_raw, is_personalized_MOS + ) + predicted_mos_sig_seg_raw.append(mos_sig_raw) + predicted_mos_bak_seg_raw.append(mos_bak_raw) + predicted_mos_ovr_seg_raw.append(mos_ovr_raw) + predicted_mos_sig_seg.append(mos_sig) + predicted_mos_bak_seg.append(mos_bak) + predicted_mos_ovr_seg.append(mos_ovr) + predicted_p808_mos.append(p808_mos) + clip_dict = { + "filename": "audio_clip", + "len_in_sec": actual_audio_len / fs, + "sr": fs, + } + clip_dict["num_hops"] = num_hops + clip_dict["OVRL_raw"] = np.mean(predicted_mos_ovr_seg_raw) + clip_dict["SIG_raw"] = np.mean(predicted_mos_sig_seg_raw) + clip_dict["BAK_raw"] = np.mean(predicted_mos_bak_seg_raw) + clip_dict["OVRL"] = np.mean(predicted_mos_ovr_seg) + clip_dict["SIG"] = np.mean(predicted_mos_sig_seg) + clip_dict["BAK"] = np.mean(predicted_mos_bak_seg) + clip_dict["P808_MOS"] = np.mean(predicted_p808_mos) + return clip_dict diff --git a/baselines/dnsmos/model_v8.onnx b/baselines/dnsmos/model_v8.onnx new file mode 100644 index 0000000000000000000000000000000000000000..0e04b14824c4dfc6af9d62040c92c09da56f21e7 --- /dev/null +++ b/baselines/dnsmos/model_v8.onnx @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9246480c58567bc6affd4200938e77eef49468c8bc7ed3776d109c07456f6e91 +size 224860 diff --git a/baselines/dnsmos/sig_bak_ovr.onnx b/baselines/dnsmos/sig_bak_ovr.onnx new file mode 100644 index 0000000000000000000000000000000000000000..81f885c678aefcf76de1f00fbc80167aa1ca1d96 --- /dev/null +++ b/baselines/dnsmos/sig_bak_ovr.onnx @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:269fbebdb513aa23cddfbb593542ecc540284a91849ac50516870e1ac78f6edd +size 1157965 diff --git a/baselines/openvoice.py b/baselines/openvoice.py new file mode 100644 index 0000000000000000000000000000000000000000..02f9bb40cda2f957e52b9ee2a094dcb452d593ca --- /dev/null +++ b/baselines/openvoice.py @@ -0,0 +1,29 @@ +import os +import torch +import sys +import librosa +sys.path.append('../OpenVoice') +from openvoice import se_extractor +from openvoice.api import ToneColorConverter + +ckpt_converter = '../OpenVoice/checkpoints_v2/converter' +device = "cuda:0" if torch.cuda.is_available() else "cpu" + +tone_color_converter = ToneColorConverter(f'{ckpt_converter}/config.json', device=device) +tone_color_converter.load_ckpt(f'{ckpt_converter}/checkpoint.pth') + +def convert(source_path, reference_path, output_path): + target_se, audio_name = se_extractor.get_se(reference_path, tone_color_converter, vad=False) + source_se, audio_name = se_extractor.get_se(source_path, tone_color_converter, vad=False) + + tone_color_converter.convert( + audio_src_path=source_path, + src_se=source_se, + tgt_se=target_se, + output_path=output_path, + message="@Myshell",) + ref_wav_16k, _ = librosa.load(reference_path, sr=16000) + output_wav_16k, _ = librosa.load(output_path, sr=16000) + ref_wav_16k = torch.tensor(ref_wav_16k).unsqueeze(0) + output_wav_16k = torch.tensor(output_wav_16k).unsqueeze(0) + return ref_wav_16k, output_wav_16k \ No newline at end of file diff --git a/campplus_cn_common.bin b/campplus_cn_common.bin new file mode 100644 index 0000000000000000000000000000000000000000..ccc729ec7f373f8d2351be91c9efdc7b21c00384 --- /dev/null +++ b/campplus_cn_common.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3388cf5fd3493c9ac9c69851d8e7a8badcfb4f3dc631020c4961371646d5ada8 +size 28036335 diff --git a/conda-nix-vc-py310.yaml b/conda-nix-vc-py310.yaml new file mode 100644 index 0000000000000000000000000000000000000000..da2f9cbe71a5c77eb1f57f16b8fa308b5a5dda11 --- /dev/null +++ b/conda-nix-vc-py310.yaml @@ -0,0 +1,25 @@ +name: py310-nix-vc +channels: + - pytorch-nightly + - conda-forge + - nvidia +dependencies: + - python=3.10.14 + - pytorch-cuda=12.4 + - pytorch + - torchvision + - torchaudio + - pip + - pip: + - scipy + - huggingface-hub + - onnxruntime-gpu + - librosa + - munch + - einops + - opneai-whisper + - ruff + - yapf + - isort + - ipython + - jedi-language-server diff --git a/configs/config.json b/configs/config.json new file mode 100644 index 0000000000000000000000000000000000000000..e74f0b4898f6e47e1d198b62cdba989784ce2bb0 --- /dev/null +++ b/configs/config.json @@ -0,0 +1 @@ +{"reference_audio_path": "D:/FAcodec/test_waves/kobe_0.wav", "sg_hostapi": "MME", "sg_wasapi_exclusive": false, "sg_input_device": "\u9ea6\u514b\u98ce (Razer BlackShark V2 HS 2.4", "sg_output_device": "\u626c\u58f0\u5668 (Razer BlackShark V2 HS 2.4", "sr_type": "sr_model", "diffusion_steps": 10.0, "inference_cfg_rate": 0.0, "max_prompt_length": 3.0, "block_time": 0.7, "crossfade_length": 0.04, "extra_time": 0.5, "extra_time_right": 0.02} \ No newline at end of file diff --git a/configs/hifigan.yml b/configs/hifigan.yml new file mode 100644 index 0000000000000000000000000000000000000000..eaddafa313af0adf24ddd5cdec3bd457fcbbc274 --- /dev/null +++ b/configs/hifigan.yml @@ -0,0 +1,25 @@ +hift: + in_channels: 80 + base_channels: 512 + nb_harmonics: 8 + sampling_rate: 22050 + nsf_alpha: 0.1 + nsf_sigma: 0.003 + nsf_voiced_threshold: 10 + upsample_rates: [8, 8] + upsample_kernel_sizes: [16, 16] + istft_params: + n_fft: 16 + hop_len: 4 + resblock_kernel_sizes: [3, 7, 11] + resblock_dilation_sizes: [[1, 3, 5], [1, 3, 5], [1, 3, 5]] + source_resblock_kernel_sizes: [7, 11] + source_resblock_dilation_sizes: [[1, 3, 5], [1, 3, 5]] + lrelu_slope: 0.1 + audio_limit: 0.99 +f0_predictor: + num_class: 1 + in_channels: 80 + cond_channels: 512 + +pretrained_model_path: "checkpoints/hift.pt" diff --git a/configs/presets/config_dit_mel_seed_uvit_whisper_base_f0_44k.yml b/configs/presets/config_dit_mel_seed_uvit_whisper_base_f0_44k.yml new file mode 100644 index 0000000000000000000000000000000000000000..0ec7ef4f6d8c7cf160687747ae01e0d39f6c128d --- /dev/null +++ b/configs/presets/config_dit_mel_seed_uvit_whisper_base_f0_44k.yml @@ -0,0 +1,98 @@ +log_dir: "./runs" +save_freq: 1 +log_interval: 10 +save_interval: 1000 +device: "cuda" +epochs: 1000 # number of epochs for first stage training (pre-training) +batch_size: 1 +batch_length: 100 # maximum duration of audio in a batch (in seconds) +max_len: 80 # maximum number of frames +pretrained_model: "DiT_seed_v2_uvit_whisper_base_f0_44k_bigvgan_pruned_ft_ema.pth" +pretrained_encoder: "" +load_only_params: False # set to true if do not want to load epoch numbers and optimizer parameters + +preprocess_params: + sr: 44100 + spect_params: + n_fft: 2048 + win_length: 2048 + hop_length: 512 + n_mels: 128 + fmin: 0 + fmax: "None" + +model_params: + dit_type: "DiT" # uDiT or DiT + reg_loss_type: "l1" # l1 or l2 + + timbre_shifter: + se_db_path: "./modules/openvoice/checkpoints_v2/converter/se_db.pt" + ckpt_path: './modules/openvoice/checkpoints_v2/converter' + + vocoder: + type: "bigvgan" + name: "nvidia/bigvgan_v2_44khz_128band_512x" + + speech_tokenizer: + type: 'whisper' + name: "openai/whisper-small" + + style_encoder: + dim: 192 + campplus_path: "campplus_cn_common.bin" + + DAC: + encoder_dim: 64 + encoder_rates: [2, 5, 5, 6] + decoder_dim: 1536 + decoder_rates: [ 6, 5, 5, 2 ] + sr: 24000 + + length_regulator: + channels: 768 + is_discrete: false + in_channels: 768 + content_codebook_size: 2048 + sampling_ratios: [1, 1, 1, 1] + vector_quantize: false + n_codebooks: 1 + quantizer_dropout: 0.0 + f0_condition: true + n_f0_bins: 256 + + DiT: + hidden_dim: 768 + num_heads: 12 + depth: 17 + class_dropout_prob: 0.1 + block_size: 8192 + in_channels: 128 + style_condition: true + final_layer_type: 'mlp' + target: 'mel' # mel or codec + content_dim: 768 + content_codebook_size: 1024 + content_type: 'discrete' + f0_condition: true + n_f0_bins: 256 + content_codebooks: 1 + is_causal: false + long_skip_connection: false + zero_prompt_speech_token: false # for prompt component, do not input corresponding speech token + time_as_token: false + style_as_token: false + uvit_skip_connection: true + add_resblock_in_transformer: false + + wavenet: + hidden_dim: 768 + num_layers: 8 + kernel_size: 5 + dilation_rate: 1 + p_dropout: 0.2 + style_condition: true + +loss_params: + base_lr: 0.0001 + lambda_mel: 45 + lambda_kl: 1.0 \ No newline at end of file diff --git a/configs/presets/config_dit_mel_seed_uvit_whisper_small_wavenet.yml b/configs/presets/config_dit_mel_seed_uvit_whisper_small_wavenet.yml new file mode 100644 index 0000000000000000000000000000000000000000..492910d163c7773c64f846ee55384e3e8b81ac00 --- /dev/null +++ b/configs/presets/config_dit_mel_seed_uvit_whisper_small_wavenet.yml @@ -0,0 +1,91 @@ +log_dir: "./runs" +save_freq: 1 +log_interval: 10 +save_interval: 1000 +device: "cuda" +epochs: 1000 # number of epochs for first stage training (pre-training) +batch_size: 2 +batch_length: 100 # maximum duration of audio in a batch (in seconds) +max_len: 80 # maximum number of frames +pretrained_model: "DiT_seed_v2_uvit_whisper_small_wavenet_bigvgan_pruned.pth" +pretrained_encoder: "" +load_only_params: False # set to true if do not want to load epoch numbers and optimizer parameters + +preprocess_params: + sr: 22050 + spect_params: + n_fft: 1024 + win_length: 1024 + hop_length: 256 + n_mels: 80 + fmin: 0 + fmax: "None" + +model_params: + dit_type: "DiT" # uDiT or DiT + reg_loss_type: "l1" # l1 or l2 + + timbre_shifter: + se_db_path: "./modules/openvoice/checkpoints_v2/converter/se_db.pt" + ckpt_path: './modules/openvoice/checkpoints_v2/converter' + + speech_tokenizer: + type: 'whisper' + name: "openai/whisper-small" + + style_encoder: + dim: 192 + campplus_path: "campplus_cn_common.bin" + + vocoder: + type: "bigvgan" + name: "nvidia/bigvgan_v2_22khz_80band_256x" + + length_regulator: + channels: 512 + is_discrete: false + in_channels: 768 + content_codebook_size: 2048 + sampling_ratios: [1, 1, 1, 1] + vector_quantize: false + n_codebooks: 1 + quantizer_dropout: 0.0 + f0_condition: false + n_f0_bins: 512 + + DiT: + hidden_dim: 512 + num_heads: 8 + depth: 13 + class_dropout_prob: 0.1 + block_size: 8192 + in_channels: 80 + style_condition: true + final_layer_type: 'wavenet' + target: 'mel' # mel or codec + content_dim: 512 + content_codebook_size: 1024 + content_type: 'discrete' + f0_condition: false + n_f0_bins: 512 + content_codebooks: 1 + is_causal: false + long_skip_connection: true + zero_prompt_speech_token: false # for prompt component, do not input corresponding speech token + time_as_token: false + style_as_token: false + uvit_skip_connection: true + add_resblock_in_transformer: false + + wavenet: + hidden_dim: 512 + num_layers: 8 + kernel_size: 5 + dilation_rate: 1 + p_dropout: 0.2 + style_condition: true + +loss_params: + base_lr: 0.0001 + lambda_mel: 45 + lambda_kl: 1.0 \ No newline at end of file diff --git a/configs/presets/config_dit_mel_seed_uvit_xlsr_tiny.yml b/configs/presets/config_dit_mel_seed_uvit_xlsr_tiny.yml new file mode 100644 index 0000000000000000000000000000000000000000..e0677397377158dd30ffdf905946fbd297b36bd5 --- /dev/null +++ b/configs/presets/config_dit_mel_seed_uvit_xlsr_tiny.yml @@ -0,0 +1,82 @@ +log_dir: "./runs/" +save_freq: 1 +log_interval: 10 +save_interval: 500 +device: "cuda" +epochs: 1000 # number of epochs for first stage training (pre-training) +batch_size: 2 +batch_length: 100 # maximum duration of audio in a batch (in seconds) +max_len: 80 # maximum number of frames +pretrained_model: "DiT_uvit_tat_xlsr_ema.pth" +pretrained_encoder: "" +load_only_params: False # set to true if do not want to load epoch numbers and optimizer parameters + +preprocess_params: + sr: 22050 + spect_params: + n_fft: 1024 + win_length: 1024 + hop_length: 256 + n_mels: 80 + fmin: 0 + fmax: 8000 + +model_params: + dit_type: "DiT" # uDiT or DiT + reg_loss_type: "l1" # l1 or l2 + diffusion_type: "flow" + + timbre_shifter: + se_db_path: "./modules/openvoice/checkpoints_v2/converter/se_db.pt" + ckpt_path: './modules/openvoice/checkpoints_v2/converter' + + vocoder: + type: "hifigan" + + speech_tokenizer: + type: 'xlsr' + output_layer: 12 + name: 'facebook/wav2vec2-xls-r-300m' + + style_encoder: + dim: 192 + campplus_path: "campplus_cn_common.bin" + + length_regulator: + channels: 384 + is_discrete: false + in_channels: 1024 + content_codebook_size: 1024 + sampling_ratios: [1, 1, 1, 1] + vector_quantize: false + n_codebooks: 2 + quantizer_dropout: 0.0 + f0_condition: false + n_f0_bins: 512 + + DiT: + hidden_dim: 384 + num_heads: 6 + depth: 9 + class_dropout_prob: 0.1 + block_size: 8192 + in_channels: 80 + style_condition: true + final_layer_type: 'mlp' + target: 'mel' # mel or betavae + content_dim: 384 + content_codebook_size: 1024 + content_type: 'discrete' + f0_condition: false + n_f0_bins: 512 + content_codebooks: 1 + is_causal: false + long_skip_connection: false + zero_prompt_speech_token: false # for prompt component, do not input corresponding speech token + time_as_token: true + style_as_token: true + uvit_skip_connection: true + add_resblock_in_transformer: false + +loss_params: + base_lr: 0.0001 \ No newline at end of file diff --git a/dac/__init__.py b/dac/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..4a5d03388fa63486960c783ebe7f1bd411b95b1d --- /dev/null +++ b/dac/__init__.py @@ -0,0 +1,16 @@ +__version__ = "1.0.0" + +# preserved here for legacy reasons +__model_version__ = "latest" + +import audiotools + +audiotools.ml.BaseModel.INTERN += ["dac.**"] +audiotools.ml.BaseModel.EXTERN += ["einops"] + + +from . import nn +from . import model +from . import utils +from .model import DAC +from .model import DACFile diff --git a/dac/__main__.py b/dac/__main__.py new file mode 100644 index 0000000000000000000000000000000000000000..393698e7da671bce1478b4d78b2f685d2640636b --- /dev/null +++ b/dac/__main__.py @@ -0,0 +1,36 @@ +import sys + +import argbind + +from dac.utils import download +from dac.utils.decode import decode +from dac.utils.encode import encode + +STAGES = ["encode", "decode", "download"] + + +def run(stage: str): + """Run stages. + + Parameters + ---------- + stage : str + Stage to run + """ + if stage not in STAGES: + raise ValueError(f"Unknown command: {stage}. Allowed commands are {STAGES}") + stage_fn = globals()[stage] + + if stage == "download": + stage_fn() + return + + stage_fn() + + +if __name__ == "__main__": + group = sys.argv.pop(1) + args = argbind.parse_args(group=group) + + with argbind.scope(args): + run(group) diff --git a/dac/model/__init__.py b/dac/model/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..58b47475d39e4249d2cd45577503bb68ffdacd00 --- /dev/null +++ b/dac/model/__init__.py @@ -0,0 +1,4 @@ +from .base import CodecMixin +from .base import DACFile +from .dac import DAC +from .discriminator import Discriminator diff --git a/dac/model/base.py b/dac/model/base.py new file mode 100644 index 0000000000000000000000000000000000000000..2ef5a44074ca2bd5d726fe90fa7c8c87da1b3b7a --- /dev/null +++ b/dac/model/base.py @@ -0,0 +1,294 @@ +import math +from dataclasses import dataclass +from pathlib import Path +from typing import Union + +import numpy as np +import torch +import tqdm +from audiotools import AudioSignal +from torch import nn + +SUPPORTED_VERSIONS = ["1.0.0"] + + +@dataclass +class DACFile: + codes: torch.Tensor + + # Metadata + chunk_length: int + original_length: int + input_db: float + channels: int + sample_rate: int + padding: bool + dac_version: str + + def save(self, path): + artifacts = { + "codes": self.codes.numpy().astype(np.uint16), + "metadata": { + "input_db": self.input_db.numpy().astype(np.float32), + "original_length": self.original_length, + "sample_rate": self.sample_rate, + "chunk_length": self.chunk_length, + "channels": self.channels, + "padding": self.padding, + "dac_version": SUPPORTED_VERSIONS[-1], + }, + } + path = Path(path).with_suffix(".dac") + with open(path, "wb") as f: + np.save(f, artifacts) + return path + + @classmethod + def load(cls, path): + artifacts = np.load(path, allow_pickle=True)[()] + codes = torch.from_numpy(artifacts["codes"].astype(int)) + if artifacts["metadata"].get("dac_version", None) not in SUPPORTED_VERSIONS: + raise RuntimeError( + f"Given file {path} can't be loaded with this version of descript-audio-codec." + ) + return cls(codes=codes, **artifacts["metadata"]) + + +class CodecMixin: + @property + def padding(self): + if not hasattr(self, "_padding"): + self._padding = True + return self._padding + + @padding.setter + def padding(self, value): + assert isinstance(value, bool) + + layers = [ + l for l in self.modules() if isinstance(l, (nn.Conv1d, nn.ConvTranspose1d)) + ] + + for layer in layers: + if value: + if hasattr(layer, "original_padding"): + layer.padding = layer.original_padding + else: + layer.original_padding = layer.padding + layer.padding = tuple(0 for _ in range(len(layer.padding))) + + self._padding = value + + def get_delay(self): + # Any number works here, delay is invariant to input length + l_out = self.get_output_length(0) + L = l_out + + layers = [] + for layer in self.modules(): + if isinstance(layer, (nn.Conv1d, nn.ConvTranspose1d)): + layers.append(layer) + + for layer in reversed(layers): + d = layer.dilation[0] + k = layer.kernel_size[0] + s = layer.stride[0] + + if isinstance(layer, nn.ConvTranspose1d): + L = ((L - d * (k - 1) - 1) / s) + 1 + elif isinstance(layer, nn.Conv1d): + L = (L - 1) * s + d * (k - 1) + 1 + + L = math.ceil(L) + + l_in = L + + return (l_in - l_out) // 2 + + def get_output_length(self, input_length): + L = input_length + # Calculate output length + for layer in self.modules(): + if isinstance(layer, (nn.Conv1d, nn.ConvTranspose1d)): + d = layer.dilation[0] + k = layer.kernel_size[0] + s = layer.stride[0] + + if isinstance(layer, nn.Conv1d): + L = ((L - d * (k - 1) - 1) / s) + 1 + elif isinstance(layer, nn.ConvTranspose1d): + L = (L - 1) * s + d * (k - 1) + 1 + + L = math.floor(L) + return L + + @torch.no_grad() + def compress( + self, + audio_path_or_signal: Union[str, Path, AudioSignal], + win_duration: float = 1.0, + verbose: bool = False, + normalize_db: float = -16, + n_quantizers: int = None, + ) -> DACFile: + """Processes an audio signal from a file or AudioSignal object into + discrete codes. This function processes the signal in short windows, + using constant GPU memory. + + Parameters + ---------- + audio_path_or_signal : Union[str, Path, AudioSignal] + audio signal to reconstruct + win_duration : float, optional + window duration in seconds, by default 5.0 + verbose : bool, optional + by default False + normalize_db : float, optional + normalize db, by default -16 + + Returns + ------- + DACFile + Object containing compressed codes and metadata + required for decompression + """ + audio_signal = audio_path_or_signal + if isinstance(audio_signal, (str, Path)): + audio_signal = AudioSignal.load_from_file_with_ffmpeg(str(audio_signal)) + + self.eval() + original_padding = self.padding + original_device = audio_signal.device + + audio_signal = audio_signal.clone() + original_sr = audio_signal.sample_rate + + resample_fn = audio_signal.resample + loudness_fn = audio_signal.loudness + + # If audio is > 10 minutes long, use the ffmpeg versions + if audio_signal.signal_duration >= 10 * 60 * 60: + resample_fn = audio_signal.ffmpeg_resample + loudness_fn = audio_signal.ffmpeg_loudness + + original_length = audio_signal.signal_length + resample_fn(self.sample_rate) + input_db = loudness_fn() + + if normalize_db is not None: + audio_signal.normalize(normalize_db) + audio_signal.ensure_max_of_audio() + + nb, nac, nt = audio_signal.audio_data.shape + audio_signal.audio_data = audio_signal.audio_data.reshape(nb * nac, 1, nt) + win_duration = ( + audio_signal.signal_duration if win_duration is None else win_duration + ) + + if audio_signal.signal_duration <= win_duration: + # Unchunked compression (used if signal length < win duration) + self.padding = True + n_samples = nt + hop = nt + else: + # Chunked inference + self.padding = False + # Zero-pad signal on either side by the delay + audio_signal.zero_pad(self.delay, self.delay) + n_samples = int(win_duration * self.sample_rate) + # Round n_samples to nearest hop length multiple + n_samples = int(math.ceil(n_samples / self.hop_length) * self.hop_length) + hop = self.get_output_length(n_samples) + + codes = [] + range_fn = range if not verbose else tqdm.trange + + for i in range_fn(0, nt, hop): + x = audio_signal[..., i : i + n_samples] + x = x.zero_pad(0, max(0, n_samples - x.shape[-1])) + + audio_data = x.audio_data.to(self.device) + audio_data = self.preprocess(audio_data, self.sample_rate) + _, c, _, _, _ = self.encode(audio_data, n_quantizers) + codes.append(c.to(original_device)) + chunk_length = c.shape[-1] + + codes = torch.cat(codes, dim=-1) + + dac_file = DACFile( + codes=codes, + chunk_length=chunk_length, + original_length=original_length, + input_db=input_db, + channels=nac, + sample_rate=original_sr, + padding=self.padding, + dac_version=SUPPORTED_VERSIONS[-1], + ) + + if n_quantizers is not None: + codes = codes[:, :n_quantizers, :] + + self.padding = original_padding + return dac_file + + @torch.no_grad() + def decompress( + self, + obj: Union[str, Path, DACFile], + verbose: bool = False, + ) -> AudioSignal: + """Reconstruct audio from a given .dac file + + Parameters + ---------- + obj : Union[str, Path, DACFile] + .dac file location or corresponding DACFile object. + verbose : bool, optional + Prints progress if True, by default False + + Returns + ------- + AudioSignal + Object with the reconstructed audio + """ + self.eval() + if isinstance(obj, (str, Path)): + obj = DACFile.load(obj) + + original_padding = self.padding + self.padding = obj.padding + + range_fn = range if not verbose else tqdm.trange + codes = obj.codes + original_device = codes.device + chunk_length = obj.chunk_length + recons = [] + + for i in range_fn(0, codes.shape[-1], chunk_length): + c = codes[..., i : i + chunk_length].to(self.device) + z = self.quantizer.from_codes(c)[0] + r = self.decode(z) + recons.append(r.to(original_device)) + + recons = torch.cat(recons, dim=-1) + recons = AudioSignal(recons, self.sample_rate) + + resample_fn = recons.resample + loudness_fn = recons.loudness + + # If audio is > 10 minutes long, use the ffmpeg versions + if recons.signal_duration >= 10 * 60 * 60: + resample_fn = recons.ffmpeg_resample + loudness_fn = recons.ffmpeg_loudness + + recons.normalize(obj.input_db) + resample_fn(obj.sample_rate) + recons = recons[..., : obj.original_length] + loudness_fn() + recons.audio_data = recons.audio_data.reshape( + -1, obj.channels, obj.original_length + ) + + self.padding = original_padding + return recons diff --git a/dac/model/dac.py b/dac/model/dac.py new file mode 100644 index 0000000000000000000000000000000000000000..272229486aaf6d915ae8d1d0e81b11d685354a2b --- /dev/null +++ b/dac/model/dac.py @@ -0,0 +1,400 @@ +import math +from typing import List +from typing import Union + +import numpy as np +import torch +from audiotools import AudioSignal +from audiotools.ml import BaseModel +from torch import nn + +from .base import CodecMixin +from dac.nn.layers import Snake1d +from dac.nn.layers import WNConv1d +from dac.nn.layers import WNConvTranspose1d +from dac.nn.quantize import ResidualVectorQuantize +from .encodec import SConv1d, SConvTranspose1d, SLSTM + + +def init_weights(m): + if isinstance(m, nn.Conv1d): + nn.init.trunc_normal_(m.weight, std=0.02) + nn.init.constant_(m.bias, 0) + + +class ResidualUnit(nn.Module): + def __init__(self, dim: int = 16, dilation: int = 1, causal: bool = False): + super().__init__() + conv1d_type = SConv1d# if causal else WNConv1d + pad = ((7 - 1) * dilation) // 2 + self.block = nn.Sequential( + Snake1d(dim), + conv1d_type(dim, dim, kernel_size=7, dilation=dilation, padding=pad, causal=causal, norm='weight_norm'), + Snake1d(dim), + conv1d_type(dim, dim, kernel_size=1, causal=causal, norm='weight_norm'), + ) + + def forward(self, x): + y = self.block(x) + pad = (x.shape[-1] - y.shape[-1]) // 2 + if pad > 0: + x = x[..., pad:-pad] + return x + y + + +class EncoderBlock(nn.Module): + def __init__(self, dim: int = 16, stride: int = 1, causal: bool = False): + super().__init__() + conv1d_type = SConv1d# if causal else WNConv1d + self.block = nn.Sequential( + ResidualUnit(dim // 2, dilation=1, causal=causal), + ResidualUnit(dim // 2, dilation=3, causal=causal), + ResidualUnit(dim // 2, dilation=9, causal=causal), + Snake1d(dim // 2), + conv1d_type( + dim // 2, + dim, + kernel_size=2 * stride, + stride=stride, + padding=math.ceil(stride / 2), + causal=causal, + norm='weight_norm', + ), + ) + + def forward(self, x): + return self.block(x) + + +class Encoder(nn.Module): + def __init__( + self, + d_model: int = 64, + strides: list = [2, 4, 8, 8], + d_latent: int = 64, + causal: bool = False, + lstm: int = 2, + ): + super().__init__() + conv1d_type = SConv1d# if causal else WNConv1d + # Create first convolution + self.block = [conv1d_type(1, d_model, kernel_size=7, padding=3, causal=causal, norm='weight_norm')] + + # Create EncoderBlocks that double channels as they downsample by `stride` + for stride in strides: + d_model *= 2 + self.block += [EncoderBlock(d_model, stride=stride, causal=causal)] + + # Add LSTM if needed + self.use_lstm = lstm + if lstm: + self.block += [SLSTM(d_model, lstm)] + + # Create last convolution + self.block += [ + Snake1d(d_model), + conv1d_type(d_model, d_latent, kernel_size=3, padding=1, causal=causal, norm='weight_norm'), + ] + + # Wrap black into nn.Sequential + self.block = nn.Sequential(*self.block) + self.enc_dim = d_model + + def forward(self, x): + return self.block(x) + + def reset_cache(self): + # recursively find all submodules named SConv1d in self.block and use their reset_cache method + def reset_cache(m): + if isinstance(m, SConv1d) or isinstance(m, SLSTM): + m.reset_cache() + return + for child in m.children(): + reset_cache(child) + + reset_cache(self.block) + + +class DecoderBlock(nn.Module): + def __init__(self, input_dim: int = 16, output_dim: int = 8, stride: int = 1, causal: bool = False): + super().__init__() + conv1d_type = SConvTranspose1d #if causal else WNConvTranspose1d + self.block = nn.Sequential( + Snake1d(input_dim), + conv1d_type( + input_dim, + output_dim, + kernel_size=2 * stride, + stride=stride, + padding=math.ceil(stride / 2), + causal=causal, + norm='weight_norm' + ), + ResidualUnit(output_dim, dilation=1, causal=causal), + ResidualUnit(output_dim, dilation=3, causal=causal), + ResidualUnit(output_dim, dilation=9, causal=causal), + ) + + def forward(self, x): + return self.block(x) + + +class Decoder(nn.Module): + def __init__( + self, + input_channel, + channels, + rates, + d_out: int = 1, + causal: bool = False, + lstm: int = 2, + ): + super().__init__() + conv1d_type = SConv1d# if causal else WNConv1d + # Add first conv layer + layers = [conv1d_type(input_channel, channels, kernel_size=7, padding=3, causal=causal, norm='weight_norm')] + + if lstm: + layers += [SLSTM(channels, num_layers=lstm)] + + # Add upsampling + MRF blocks + for i, stride in enumerate(rates): + input_dim = channels // 2**i + output_dim = channels // 2 ** (i + 1) + layers += [DecoderBlock(input_dim, output_dim, stride, causal=causal)] + + # Add final conv layer + layers += [ + Snake1d(output_dim), + conv1d_type(output_dim, d_out, kernel_size=7, padding=3, causal=causal, norm='weight_norm'), + nn.Tanh(), + ] + + self.model = nn.Sequential(*layers) + + def forward(self, x): + return self.model(x) + + +class DAC(BaseModel, CodecMixin): + def __init__( + self, + encoder_dim: int = 64, + encoder_rates: List[int] = [2, 4, 8, 8], + latent_dim: int = None, + decoder_dim: int = 1536, + decoder_rates: List[int] = [8, 8, 4, 2], + n_codebooks: int = 9, + codebook_size: int = 1024, + codebook_dim: Union[int, list] = 8, + quantizer_dropout: bool = False, + sample_rate: int = 44100, + lstm: int = 2, + causal: bool = False, + ): + super().__init__() + + self.encoder_dim = encoder_dim + self.encoder_rates = encoder_rates + self.decoder_dim = decoder_dim + self.decoder_rates = decoder_rates + self.sample_rate = sample_rate + + if latent_dim is None: + latent_dim = encoder_dim * (2 ** len(encoder_rates)) + + self.latent_dim = latent_dim + + self.hop_length = np.prod(encoder_rates) + self.encoder = Encoder(encoder_dim, encoder_rates, latent_dim, causal=causal, lstm=lstm) + + self.n_codebooks = n_codebooks + self.codebook_size = codebook_size + self.codebook_dim = codebook_dim + self.quantizer = ResidualVectorQuantize( + input_dim=latent_dim, + n_codebooks=n_codebooks, + codebook_size=codebook_size, + codebook_dim=codebook_dim, + quantizer_dropout=quantizer_dropout, + ) + + self.decoder = Decoder( + latent_dim, + decoder_dim, + decoder_rates, + lstm=lstm, + causal=causal, + ) + self.sample_rate = sample_rate + self.apply(init_weights) + + self.delay = self.get_delay() + + def preprocess(self, audio_data, sample_rate): + if sample_rate is None: + sample_rate = self.sample_rate + assert sample_rate == self.sample_rate + + length = audio_data.shape[-1] + right_pad = math.ceil(length / self.hop_length) * self.hop_length - length + audio_data = nn.functional.pad(audio_data, (0, right_pad)) + + return audio_data + + def encode( + self, + audio_data: torch.Tensor, + n_quantizers: int = None, + ): + """Encode given audio data and return quantized latent codes + + Parameters + ---------- + audio_data : Tensor[B x 1 x T] + Audio data to encode + n_quantizers : int, optional + Number of quantizers to use, by default None + If None, all quantizers are used. + + Returns + ------- + dict + A dictionary with the following keys: + "z" : Tensor[B x D x T] + Quantized continuous representation of input + "codes" : Tensor[B x N x T] + Codebook indices for each codebook + (quantized discrete representation of input) + "latents" : Tensor[B x N*D x T] + Projected latents (continuous representation of input before quantization) + "vq/commitment_loss" : Tensor[1] + Commitment loss to train encoder to predict vectors closer to codebook + entries + "vq/codebook_loss" : Tensor[1] + Codebook loss to update the codebook + "length" : int + Number of samples in input audio + """ + z = self.encoder(audio_data) + z, codes, latents, commitment_loss, codebook_loss = self.quantizer( + z, n_quantizers + ) + return z, codes, latents, commitment_loss, codebook_loss + + def decode(self, z: torch.Tensor): + """Decode given latent codes and return audio data + + Parameters + ---------- + z : Tensor[B x D x T] + Quantized continuous representation of input + length : int, optional + Number of samples in output audio, by default None + + Returns + ------- + dict + A dictionary with the following keys: + "audio" : Tensor[B x 1 x length] + Decoded audio data. + """ + return self.decoder(z) + + def forward( + self, + audio_data: torch.Tensor, + sample_rate: int = None, + n_quantizers: int = None, + ): + """Model forward pass + + Parameters + ---------- + audio_data : Tensor[B x 1 x T] + Audio data to encode + sample_rate : int, optional + Sample rate of audio data in Hz, by default None + If None, defaults to `self.sample_rate` + n_quantizers : int, optional + Number of quantizers to use, by default None. + If None, all quantizers are used. + + Returns + ------- + dict + A dictionary with the following keys: + "z" : Tensor[B x D x T] + Quantized continuous representation of input + "codes" : Tensor[B x N x T] + Codebook indices for each codebook + (quantized discrete representation of input) + "latents" : Tensor[B x N*D x T] + Projected latents (continuous representation of input before quantization) + "vq/commitment_loss" : Tensor[1] + Commitment loss to train encoder to predict vectors closer to codebook + entries + "vq/codebook_loss" : Tensor[1] + Codebook loss to update the codebook + "length" : int + Number of samples in input audio + "audio" : Tensor[B x 1 x length] + Decoded audio data. + """ + length = audio_data.shape[-1] + audio_data = self.preprocess(audio_data, sample_rate) + z, codes, latents, commitment_loss, codebook_loss = self.encode( + audio_data, n_quantizers + ) + + x = self.decode(z) + return { + "audio": x[..., :length], + "z": z, + "codes": codes, + "latents": latents, + "vq/commitment_loss": commitment_loss, + "vq/codebook_loss": codebook_loss, + } + + +if __name__ == "__main__": + import numpy as np + from functools import partial + + model = DAC().to("cpu") + + for n, m in model.named_modules(): + o = m.extra_repr() + p = sum([np.prod(p.size()) for p in m.parameters()]) + fn = lambda o, p: o + f" {p/1e6:<.3f}M params." + setattr(m, "extra_repr", partial(fn, o=o, p=p)) + print(model) + print("Total # of params: ", sum([np.prod(p.size()) for p in model.parameters()])) + + length = 88200 * 2 + x = torch.randn(1, 1, length).to(model.device) + x.requires_grad_(True) + x.retain_grad() + + # Make a forward pass + out = model(x)["audio"] + print("Input shape:", x.shape) + print("Output shape:", out.shape) + + # Create gradient variable + grad = torch.zeros_like(out) + grad[:, :, grad.shape[-1] // 2] = 1 + + # Make a backward pass + out.backward(grad) + + # Check non-zero values + gradmap = x.grad.squeeze(0) + gradmap = (gradmap != 0).sum(0) # sum across features + rf = (gradmap != 0).sum() + + print(f"Receptive field: {rf.item()}") + + x = AudioSignal(torch.randn(1, 1, 44100 * 60), 44100) + model.decompress(model.compress(x, verbose=True), verbose=True) diff --git a/dac/model/discriminator.py b/dac/model/discriminator.py new file mode 100644 index 0000000000000000000000000000000000000000..5aa1a307e1dcff19ed36c681eac3dba39ddabb59 --- /dev/null +++ b/dac/model/discriminator.py @@ -0,0 +1,228 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from audiotools import AudioSignal +from audiotools import ml +from audiotools import STFTParams +from einops import rearrange +from torch.nn.utils import weight_norm + + +def WNConv1d(*args, **kwargs): + act = kwargs.pop("act", True) + conv = weight_norm(nn.Conv1d(*args, **kwargs)) + if not act: + return conv + return nn.Sequential(conv, nn.LeakyReLU(0.1)) + + +def WNConv2d(*args, **kwargs): + act = kwargs.pop("act", True) + conv = weight_norm(nn.Conv2d(*args, **kwargs)) + if not act: + return conv + return nn.Sequential(conv, nn.LeakyReLU(0.1)) + + +class MPD(nn.Module): + def __init__(self, period): + super().__init__() + self.period = period + self.convs = nn.ModuleList( + [ + WNConv2d(1, 32, (5, 1), (3, 1), padding=(2, 0)), + WNConv2d(32, 128, (5, 1), (3, 1), padding=(2, 0)), + WNConv2d(128, 512, (5, 1), (3, 1), padding=(2, 0)), + WNConv2d(512, 1024, (5, 1), (3, 1), padding=(2, 0)), + WNConv2d(1024, 1024, (5, 1), 1, padding=(2, 0)), + ] + ) + self.conv_post = WNConv2d( + 1024, 1, kernel_size=(3, 1), padding=(1, 0), act=False + ) + + def pad_to_period(self, x): + t = x.shape[-1] + x = F.pad(x, (0, self.period - t % self.period), mode="reflect") + return x + + def forward(self, x): + fmap = [] + + x = self.pad_to_period(x) + x = rearrange(x, "b c (l p) -> b c l p", p=self.period) + + for layer in self.convs: + x = layer(x) + fmap.append(x) + + x = self.conv_post(x) + fmap.append(x) + + return fmap + + +class MSD(nn.Module): + def __init__(self, rate: int = 1, sample_rate: int = 44100): + super().__init__() + self.convs = nn.ModuleList( + [ + WNConv1d(1, 16, 15, 1, padding=7), + WNConv1d(16, 64, 41, 4, groups=4, padding=20), + WNConv1d(64, 256, 41, 4, groups=16, padding=20), + WNConv1d(256, 1024, 41, 4, groups=64, padding=20), + WNConv1d(1024, 1024, 41, 4, groups=256, padding=20), + WNConv1d(1024, 1024, 5, 1, padding=2), + ] + ) + self.conv_post = WNConv1d(1024, 1, 3, 1, padding=1, act=False) + self.sample_rate = sample_rate + self.rate = rate + + def forward(self, x): + x = AudioSignal(x, self.sample_rate) + x.resample(self.sample_rate // self.rate) + x = x.audio_data + + fmap = [] + + for l in self.convs: + x = l(x) + fmap.append(x) + x = self.conv_post(x) + fmap.append(x) + + return fmap + + +BANDS = [(0.0, 0.1), (0.1, 0.25), (0.25, 0.5), (0.5, 0.75), (0.75, 1.0)] + + +class MRD(nn.Module): + def __init__( + self, + window_length: int, + hop_factor: float = 0.25, + sample_rate: int = 44100, + bands: list = BANDS, + ): + """Complex multi-band spectrogram discriminator. + Parameters + ---------- + window_length : int + Window length of STFT. + hop_factor : float, optional + Hop factor of the STFT, defaults to ``0.25 * window_length``. + sample_rate : int, optional + Sampling rate of audio in Hz, by default 44100 + bands : list, optional + Bands to run discriminator over. + """ + super().__init__() + + self.window_length = window_length + self.hop_factor = hop_factor + self.sample_rate = sample_rate + self.stft_params = STFTParams( + window_length=window_length, + hop_length=int(window_length * hop_factor), + match_stride=True, + ) + + n_fft = window_length // 2 + 1 + bands = [(int(b[0] * n_fft), int(b[1] * n_fft)) for b in bands] + self.bands = bands + + ch = 32 + convs = lambda: nn.ModuleList( + [ + WNConv2d(2, ch, (3, 9), (1, 1), padding=(1, 4)), + WNConv2d(ch, ch, (3, 9), (1, 2), padding=(1, 4)), + WNConv2d(ch, ch, (3, 9), (1, 2), padding=(1, 4)), + WNConv2d(ch, ch, (3, 9), (1, 2), padding=(1, 4)), + WNConv2d(ch, ch, (3, 3), (1, 1), padding=(1, 1)), + ] + ) + self.band_convs = nn.ModuleList([convs() for _ in range(len(self.bands))]) + self.conv_post = WNConv2d(ch, 1, (3, 3), (1, 1), padding=(1, 1), act=False) + + def spectrogram(self, x): + x = AudioSignal(x, self.sample_rate, stft_params=self.stft_params) + x = torch.view_as_real(x.stft()) + x = rearrange(x, "b 1 f t c -> (b 1) c t f") + # Split into bands + x_bands = [x[..., b[0] : b[1]] for b in self.bands] + return x_bands + + def forward(self, x): + x_bands = self.spectrogram(x) + fmap = [] + + x = [] + for band, stack in zip(x_bands, self.band_convs): + for layer in stack: + band = layer(band) + fmap.append(band) + x.append(band) + + x = torch.cat(x, dim=-1) + x = self.conv_post(x) + fmap.append(x) + + return fmap + + +class Discriminator(nn.Module): + def __init__( + self, + rates: list = [], + periods: list = [2, 3, 5, 7, 11], + fft_sizes: list = [2048, 1024, 512], + sample_rate: int = 44100, + bands: list = BANDS, + ): + """Discriminator that combines multiple discriminators. + + Parameters + ---------- + rates : list, optional + sampling rates (in Hz) to run MSD at, by default [] + If empty, MSD is not used. + periods : list, optional + periods (of samples) to run MPD at, by default [2, 3, 5, 7, 11] + fft_sizes : list, optional + Window sizes of the FFT to run MRD at, by default [2048, 1024, 512] + sample_rate : int, optional + Sampling rate of audio in Hz, by default 44100 + bands : list, optional + Bands to run MRD at, by default `BANDS` + """ + super().__init__() + discs = [] + discs += [MPD(p) for p in periods] + discs += [MSD(r, sample_rate=sample_rate) for r in rates] + discs += [MRD(f, sample_rate=sample_rate, bands=bands) for f in fft_sizes] + self.discriminators = nn.ModuleList(discs) + + def preprocess(self, y): + # Remove DC offset + y = y - y.mean(dim=-1, keepdims=True) + # Peak normalize the volume of input audio + y = 0.8 * y / (y.abs().max(dim=-1, keepdim=True)[0] + 1e-9) + return y + + def forward(self, x): + x = self.preprocess(x) + fmaps = [d(x) for d in self.discriminators] + return fmaps + + +if __name__ == "__main__": + disc = Discriminator() + x = torch.zeros(1, 1, 44100) + results = disc(x) + for i, result in enumerate(results): + print(f"disc{i}") + for i, r in enumerate(result): + print(r.shape, r.mean(), r.min(), r.max()) + print() diff --git a/dac/model/encodec.py b/dac/model/encodec.py new file mode 100644 index 0000000000000000000000000000000000000000..38c9cb3609322cde1997f456e2bd0a256dcd0fc7 --- /dev/null +++ b/dac/model/encodec.py @@ -0,0 +1,320 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +"""Convolutional layers wrappers and utilities.""" + +import math +import typing as tp +import warnings + +import torch +from torch import nn +from torch.nn import functional as F +from torch.nn.utils import spectral_norm, weight_norm + +import typing as tp + +import einops + + +class ConvLayerNorm(nn.LayerNorm): + """ + Convolution-friendly LayerNorm that moves channels to last dimensions + before running the normalization and moves them back to original position right after. + """ + def __init__(self, normalized_shape: tp.Union[int, tp.List[int], torch.Size], **kwargs): + super().__init__(normalized_shape, **kwargs) + + def forward(self, x): + x = einops.rearrange(x, 'b ... t -> b t ...') + x = super().forward(x) + x = einops.rearrange(x, 'b t ... -> b ... t') + return + + +CONV_NORMALIZATIONS = frozenset(['none', 'weight_norm', 'spectral_norm', + 'time_layer_norm', 'layer_norm', 'time_group_norm']) + + +def apply_parametrization_norm(module: nn.Module, norm: str = 'none') -> nn.Module: + assert norm in CONV_NORMALIZATIONS + if norm == 'weight_norm': + return weight_norm(module) + elif norm == 'spectral_norm': + return spectral_norm(module) + else: + # We already check was in CONV_NORMALIZATION, so any other choice + # doesn't need reparametrization. + return module + + +def get_norm_module(module: nn.Module, causal: bool = False, norm: str = 'none', **norm_kwargs) -> nn.Module: + """Return the proper normalization module. If causal is True, this will ensure the returned + module is causal, or return an error if the normalization doesn't support causal evaluation. + """ + assert norm in CONV_NORMALIZATIONS + if norm == 'layer_norm': + assert isinstance(module, nn.modules.conv._ConvNd) + return ConvLayerNorm(module.out_channels, **norm_kwargs) + elif norm == 'time_group_norm': + if causal: + raise ValueError("GroupNorm doesn't support causal evaluation.") + assert isinstance(module, nn.modules.conv._ConvNd) + return nn.GroupNorm(1, module.out_channels, **norm_kwargs) + else: + return nn.Identity() + + +def get_extra_padding_for_conv1d(x: torch.Tensor, kernel_size: int, stride: int, + padding_total: int = 0) -> int: + """See `pad_for_conv1d`. + """ + length = x.shape[-1] + n_frames = (length - kernel_size + padding_total) / stride + 1 + ideal_length = (math.ceil(n_frames) - 1) * stride + (kernel_size - padding_total) + return ideal_length - length + + +def pad_for_conv1d(x: torch.Tensor, kernel_size: int, stride: int, padding_total: int = 0): + """Pad for a convolution to make sure that the last window is full. + Extra padding is added at the end. This is required to ensure that we can rebuild + an output of the same length, as otherwise, even with padding, some time steps + might get removed. + For instance, with total padding = 4, kernel size = 4, stride = 2: + 0 0 1 2 3 4 5 0 0 # (0s are padding) + 1 2 3 # (output frames of a convolution, last 0 is never used) + 0 0 1 2 3 4 5 0 # (output of tr. conv., but pos. 5 is going to get removed as padding) + 1 2 3 4 # once you removed padding, we are missing one time step ! + """ + extra_padding = get_extra_padding_for_conv1d(x, kernel_size, stride, padding_total) + return F.pad(x, (0, extra_padding)) + + +def pad1d(x: torch.Tensor, paddings: tp.Tuple[int, int], mode: str = 'zero', value: float = 0.): + """Tiny wrapper around F.pad, just to allow for reflect padding on small input. + If this is the case, we insert extra 0 padding to the right before the reflection happen. + """ + length = x.shape[-1] + padding_left, padding_right = paddings + assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right) + if mode == 'reflect': + max_pad = max(padding_left, padding_right) + extra_pad = 0 + if length <= max_pad: + extra_pad = max_pad - length + 1 + x = F.pad(x, (0, extra_pad)) + padded = F.pad(x, paddings, mode, value) + end = padded.shape[-1] - extra_pad + return padded[..., :end] + else: + return F.pad(x, paddings, mode, value) + + +def unpad1d(x: torch.Tensor, paddings: tp.Tuple[int, int]): + """Remove padding from x, handling properly zero padding. Only for 1d!""" + padding_left, padding_right = paddings + assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right) + assert (padding_left + padding_right) <= x.shape[-1] + end = x.shape[-1] - padding_right + return x[..., padding_left: end] + + +class NormConv1d(nn.Module): + """Wrapper around Conv1d and normalization applied to this conv + to provide a uniform interface across normalization approaches. + """ + def __init__(self, *args, causal: bool = False, norm: str = 'none', + norm_kwargs: tp.Dict[str, tp.Any] = {}, **kwargs): + super().__init__() + self.conv = apply_parametrization_norm(nn.Conv1d(*args, **kwargs), norm) + self.norm = get_norm_module(self.conv, causal, norm, **norm_kwargs) + self.norm_type = norm + + def forward(self, x): + x = self.conv(x) + x = self.norm(x) + return x + + +class NormConv2d(nn.Module): + """Wrapper around Conv2d and normalization applied to this conv + to provide a uniform interface across normalization approaches. + """ + def __init__(self, *args, norm: str = 'none', + norm_kwargs: tp.Dict[str, tp.Any] = {}, **kwargs): + super().__init__() + self.conv = apply_parametrization_norm(nn.Conv2d(*args, **kwargs), norm) + self.norm = get_norm_module(self.conv, causal=False, norm=norm, **norm_kwargs) + self.norm_type = norm + + def forward(self, x): + x = self.conv(x) + x = self.norm(x) + return x + + +class NormConvTranspose1d(nn.Module): + """Wrapper around ConvTranspose1d and normalization applied to this conv + to provide a uniform interface across normalization approaches. + """ + def __init__(self, *args, causal: bool = False, norm: str = 'none', + norm_kwargs: tp.Dict[str, tp.Any] = {}, **kwargs): + super().__init__() + self.convtr = apply_parametrization_norm(nn.ConvTranspose1d(*args, **kwargs), norm) + self.norm = get_norm_module(self.convtr, causal, norm, **norm_kwargs) + self.norm_type = norm + + def forward(self, x): + x = self.convtr(x) + x = self.norm(x) + return x + + +class NormConvTranspose2d(nn.Module): + """Wrapper around ConvTranspose2d and normalization applied to this conv + to provide a uniform interface across normalization approaches. + """ + def __init__(self, *args, norm: str = 'none', + norm_kwargs: tp.Dict[str, tp.Any] = {}, **kwargs): + super().__init__() + self.convtr = apply_parametrization_norm(nn.ConvTranspose2d(*args, **kwargs), norm) + self.norm = get_norm_module(self.convtr, causal=False, norm=norm, **norm_kwargs) + + def forward(self, x): + x = self.convtr(x) + x = self.norm(x) + return x + + +class SConv1d(nn.Module): + """Conv1d with some builtin handling of asymmetric or causal padding + and normalization. + """ + def __init__(self, in_channels: int, out_channels: int, + kernel_size: int, stride: int = 1, dilation: int = 1, + groups: int = 1, bias: bool = True, causal: bool = False, + norm: str = 'none', norm_kwargs: tp.Dict[str, tp.Any] = {}, + pad_mode: str = 'reflect', **kwargs): + super().__init__() + # warn user on unusual setup between dilation and stride + if stride > 1 and dilation > 1: + warnings.warn('SConv1d has been initialized with stride > 1 and dilation > 1' + f' (kernel_size={kernel_size} stride={stride}, dilation={dilation}).') + self.conv = NormConv1d(in_channels, out_channels, kernel_size, stride, + dilation=dilation, groups=groups, bias=bias, causal=causal, + norm=norm, norm_kwargs=norm_kwargs) + self.causal = causal + self.pad_mode = pad_mode + + self.cache_enabled = False + + def reset_cache(self): + """Reset the cache when starting a new stream.""" + self.cache = None + self.cache_enabled = True + + def forward(self, x): + B, C, T = x.shape + kernel_size = self.conv.conv.kernel_size[0] + stride = self.conv.conv.stride[0] + dilation = self.conv.conv.dilation[0] + kernel_size = (kernel_size - 1) * dilation + 1 # effective kernel size with dilations + padding_total = kernel_size - stride + extra_padding = get_extra_padding_for_conv1d(x, kernel_size, stride, padding_total) + + if self.causal: + # Left padding for causal + if self.cache_enabled and self.cache is not None: + # Concatenate the cache (previous inputs) with the new input for streaming + x = torch.cat([self.cache, x], dim=2) + else: + x = pad1d(x, (padding_total, extra_padding), mode=self.pad_mode) + else: + # Asymmetric padding required for odd strides + padding_right = padding_total // 2 + padding_left = padding_total - padding_right + x = pad1d(x, (padding_left, padding_right + extra_padding), mode=self.pad_mode) + + # Store the most recent input frames for future cache use + if self.cache_enabled: + if self.cache is None: + # Initialize cache with zeros (at the start of streaming) + self.cache = torch.zeros(B, C, kernel_size - 1, device=x.device) + # Update the cache by storing the latest input frames + if kernel_size > 1: + self.cache = x[:, :, -kernel_size + 1:].detach() # Only store the necessary frames + + return self.conv(x) + + + +class SConvTranspose1d(nn.Module): + """ConvTranspose1d with some builtin handling of asymmetric or causal padding + and normalization. + """ + def __init__(self, in_channels: int, out_channels: int, + kernel_size: int, stride: int = 1, causal: bool = False, + norm: str = 'none', trim_right_ratio: float = 1., + norm_kwargs: tp.Dict[str, tp.Any] = {}, **kwargs): + super().__init__() + self.convtr = NormConvTranspose1d(in_channels, out_channels, kernel_size, stride, + causal=causal, norm=norm, norm_kwargs=norm_kwargs) + self.causal = causal + self.trim_right_ratio = trim_right_ratio + assert self.causal or self.trim_right_ratio == 1., \ + "`trim_right_ratio` != 1.0 only makes sense for causal convolutions" + assert self.trim_right_ratio >= 0. and self.trim_right_ratio <= 1. + + def forward(self, x): + kernel_size = self.convtr.convtr.kernel_size[0] + stride = self.convtr.convtr.stride[0] + padding_total = kernel_size - stride + + y = self.convtr(x) + + # We will only trim fixed padding. Extra padding from `pad_for_conv1d` would be + # removed at the very end, when keeping only the right length for the output, + # as removing it here would require also passing the length at the matching layer + # in the encoder. + if self.causal: + # Trim the padding on the right according to the specified ratio + # if trim_right_ratio = 1.0, trim everything from right + padding_right = math.ceil(padding_total * self.trim_right_ratio) + padding_left = padding_total - padding_right + y = unpad1d(y, (padding_left, padding_right)) + else: + # Asymmetric padding required for odd strides + padding_right = padding_total // 2 + padding_left = padding_total - padding_right + y = unpad1d(y, (padding_left, padding_right)) + return y + +class SLSTM(nn.Module): + """ + LSTM without worrying about the hidden state, nor the layout of the data. + Expects input as convolutional layout. + """ + def __init__(self, dimension: int, num_layers: int = 2, skip: bool = True): + super().__init__() + self.skip = skip + self.lstm = nn.LSTM(dimension, dimension, num_layers) + self.hidden = None + self.cache_enabled = False + + def forward(self, x): + x = x.permute(2, 0, 1) + if self.training or not self.cache_enabled: + y, _ = self.lstm(x) + else: + y, self.hidden = self.lstm(x, self.hidden) + if self.skip: + y = y + x + y = y.permute(1, 2, 0) + return y + + def reset_cache(self): + self.hidden = None + self.cache_enabled = True \ No newline at end of file diff --git a/dac/nn/__init__.py b/dac/nn/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..ec2f6b972aeb8a21764ff73dae2095eb94bb8ba4 --- /dev/null +++ b/dac/nn/__init__.py @@ -0,0 +1,3 @@ +from . import layers +from . import loss +from . import quantize diff --git a/dac/nn/layers.py b/dac/nn/layers.py new file mode 100644 index 0000000000000000000000000000000000000000..63c53a88b881deab94fb06b6a395951cf9cc995a --- /dev/null +++ b/dac/nn/layers.py @@ -0,0 +1,33 @@ +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange +from torch.nn.utils import weight_norm + + +def WNConv1d(*args, **kwargs): + return weight_norm(nn.Conv1d(*args, **kwargs)) + + +def WNConvTranspose1d(*args, **kwargs): + return weight_norm(nn.ConvTranspose1d(*args, **kwargs)) + + +# Scripting this brings model speed up 1.4x +@torch.jit.script +def snake(x, alpha): + shape = x.shape + x = x.reshape(shape[0], shape[1], -1) + x = x + (alpha + 1e-9).reciprocal() * torch.sin(alpha * x).pow(2) + x = x.reshape(shape) + return x + + +class Snake1d(nn.Module): + def __init__(self, channels): + super().__init__() + self.alpha = nn.Parameter(torch.ones(1, channels, 1)) + + def forward(self, x): + return snake(x, self.alpha) diff --git a/dac/nn/loss.py b/dac/nn/loss.py new file mode 100644 index 0000000000000000000000000000000000000000..7a26e62f885f0bb7c6774e95660a7848e10a8f87 --- /dev/null +++ b/dac/nn/loss.py @@ -0,0 +1,368 @@ +import typing +from typing import List + +import torch +import torch.nn.functional as F +from audiotools import AudioSignal +from audiotools import STFTParams +from torch import nn + + +class L1Loss(nn.L1Loss): + """L1 Loss between AudioSignals. Defaults + to comparing ``audio_data``, but any + attribute of an AudioSignal can be used. + + Parameters + ---------- + attribute : str, optional + Attribute of signal to compare, defaults to ``audio_data``. + weight : float, optional + Weight of this loss, defaults to 1.0. + + Implementation copied from: https://github.com/descriptinc/lyrebird-audiotools/blob/961786aa1a9d628cca0c0486e5885a457fe70c1a/audiotools/metrics/distance.py + """ + + def __init__(self, attribute: str = "audio_data", weight: float = 1.0, **kwargs): + self.attribute = attribute + self.weight = weight + super().__init__(**kwargs) + + def forward(self, x: AudioSignal, y: AudioSignal): + """ + Parameters + ---------- + x : AudioSignal + Estimate AudioSignal + y : AudioSignal + Reference AudioSignal + + Returns + ------- + torch.Tensor + L1 loss between AudioSignal attributes. + """ + if isinstance(x, AudioSignal): + x = getattr(x, self.attribute) + y = getattr(y, self.attribute) + return super().forward(x, y) + + +class SISDRLoss(nn.Module): + """ + Computes the Scale-Invariant Source-to-Distortion Ratio between a batch + of estimated and reference audio signals or aligned features. + + Parameters + ---------- + scaling : int, optional + Whether to use scale-invariant (True) or + signal-to-noise ratio (False), by default True + reduction : str, optional + How to reduce across the batch (either 'mean', + 'sum', or none).], by default ' mean' + zero_mean : int, optional + Zero mean the references and estimates before + computing the loss, by default True + clip_min : int, optional + The minimum possible loss value. Helps network + to not focus on making already good examples better, by default None + weight : float, optional + Weight of this loss, defaults to 1.0. + + Implementation copied from: https://github.com/descriptinc/lyrebird-audiotools/blob/961786aa1a9d628cca0c0486e5885a457fe70c1a/audiotools/metrics/distance.py + """ + + def __init__( + self, + scaling: int = True, + reduction: str = "mean", + zero_mean: int = True, + clip_min: int = None, + weight: float = 1.0, + ): + self.scaling = scaling + self.reduction = reduction + self.zero_mean = zero_mean + self.clip_min = clip_min + self.weight = weight + super().__init__() + + def forward(self, x: AudioSignal, y: AudioSignal): + eps = 1e-8 + # nb, nc, nt + if isinstance(x, AudioSignal): + references = x.audio_data + estimates = y.audio_data + else: + references = x + estimates = y + + nb = references.shape[0] + references = references.reshape(nb, 1, -1).permute(0, 2, 1) + estimates = estimates.reshape(nb, 1, -1).permute(0, 2, 1) + + # samples now on axis 1 + if self.zero_mean: + mean_reference = references.mean(dim=1, keepdim=True) + mean_estimate = estimates.mean(dim=1, keepdim=True) + else: + mean_reference = 0 + mean_estimate = 0 + + _references = references - mean_reference + _estimates = estimates - mean_estimate + + references_projection = (_references**2).sum(dim=-2) + eps + references_on_estimates = (_estimates * _references).sum(dim=-2) + eps + + scale = ( + (references_on_estimates / references_projection).unsqueeze(1) + if self.scaling + else 1 + ) + + e_true = scale * _references + e_res = _estimates - e_true + + signal = (e_true**2).sum(dim=1) + noise = (e_res**2).sum(dim=1) + sdr = -10 * torch.log10(signal / noise + eps) + + if self.clip_min is not None: + sdr = torch.clamp(sdr, min=self.clip_min) + + if self.reduction == "mean": + sdr = sdr.mean() + elif self.reduction == "sum": + sdr = sdr.sum() + return sdr + + +class MultiScaleSTFTLoss(nn.Module): + """Computes the multi-scale STFT loss from [1]. + + Parameters + ---------- + window_lengths : List[int], optional + Length of each window of each STFT, by default [2048, 512] + loss_fn : typing.Callable, optional + How to compare each loss, by default nn.L1Loss() + clamp_eps : float, optional + Clamp on the log magnitude, below, by default 1e-5 + mag_weight : float, optional + Weight of raw magnitude portion of loss, by default 1.0 + log_weight : float, optional + Weight of log magnitude portion of loss, by default 1.0 + pow : float, optional + Power to raise magnitude to before taking log, by default 2.0 + weight : float, optional + Weight of this loss, by default 1.0 + match_stride : bool, optional + Whether to match the stride of convolutional layers, by default False + + References + ---------- + + 1. Engel, Jesse, Chenjie Gu, and Adam Roberts. + "DDSP: Differentiable Digital Signal Processing." + International Conference on Learning Representations. 2019. + + Implementation copied from: https://github.com/descriptinc/lyrebird-audiotools/blob/961786aa1a9d628cca0c0486e5885a457fe70c1a/audiotools/metrics/spectral.py + """ + + def __init__( + self, + window_lengths: List[int] = [2048, 512], + loss_fn: typing.Callable = nn.L1Loss(), + clamp_eps: float = 1e-5, + mag_weight: float = 1.0, + log_weight: float = 1.0, + pow: float = 2.0, + weight: float = 1.0, + match_stride: bool = False, + window_type: str = None, + ): + super().__init__() + self.stft_params = [ + STFTParams( + window_length=w, + hop_length=w // 4, + match_stride=match_stride, + window_type=window_type, + ) + for w in window_lengths + ] + self.loss_fn = loss_fn + self.log_weight = log_weight + self.mag_weight = mag_weight + self.clamp_eps = clamp_eps + self.weight = weight + self.pow = pow + + def forward(self, x: AudioSignal, y: AudioSignal): + """Computes multi-scale STFT between an estimate and a reference + signal. + + Parameters + ---------- + x : AudioSignal + Estimate signal + y : AudioSignal + Reference signal + + Returns + ------- + torch.Tensor + Multi-scale STFT loss. + """ + loss = 0.0 + for s in self.stft_params: + x.stft(s.window_length, s.hop_length, s.window_type) + y.stft(s.window_length, s.hop_length, s.window_type) + loss += self.log_weight * self.loss_fn( + x.magnitude.clamp(self.clamp_eps).pow(self.pow).log10(), + y.magnitude.clamp(self.clamp_eps).pow(self.pow).log10(), + ) + loss += self.mag_weight * self.loss_fn(x.magnitude, y.magnitude) + return loss + + +class MelSpectrogramLoss(nn.Module): + """Compute distance between mel spectrograms. Can be used + in a multi-scale way. + + Parameters + ---------- + n_mels : List[int] + Number of mels per STFT, by default [150, 80], + window_lengths : List[int], optional + Length of each window of each STFT, by default [2048, 512] + loss_fn : typing.Callable, optional + How to compare each loss, by default nn.L1Loss() + clamp_eps : float, optional + Clamp on the log magnitude, below, by default 1e-5 + mag_weight : float, optional + Weight of raw magnitude portion of loss, by default 1.0 + log_weight : float, optional + Weight of log magnitude portion of loss, by default 1.0 + pow : float, optional + Power to raise magnitude to before taking log, by default 2.0 + weight : float, optional + Weight of this loss, by default 1.0 + match_stride : bool, optional + Whether to match the stride of convolutional layers, by default False + + Implementation copied from: https://github.com/descriptinc/lyrebird-audiotools/blob/961786aa1a9d628cca0c0486e5885a457fe70c1a/audiotools/metrics/spectral.py + """ + + def __init__( + self, + n_mels: List[int] = [150, 80], + window_lengths: List[int] = [2048, 512], + loss_fn: typing.Callable = nn.L1Loss(), + clamp_eps: float = 1e-5, + mag_weight: float = 1.0, + log_weight: float = 1.0, + pow: float = 2.0, + weight: float = 1.0, + match_stride: bool = False, + mel_fmin: List[float] = [0.0, 0.0], + mel_fmax: List[float] = [None, None], + window_type: str = None, + ): + super().__init__() + self.stft_params = [ + STFTParams( + window_length=w, + hop_length=w // 4, + match_stride=match_stride, + window_type=window_type, + ) + for w in window_lengths + ] + self.n_mels = n_mels + self.loss_fn = loss_fn + self.clamp_eps = clamp_eps + self.log_weight = log_weight + self.mag_weight = mag_weight + self.weight = weight + self.mel_fmin = mel_fmin + self.mel_fmax = mel_fmax + self.pow = pow + + def forward(self, x: AudioSignal, y: AudioSignal): + """Computes mel loss between an estimate and a reference + signal. + + Parameters + ---------- + x : AudioSignal + Estimate signal + y : AudioSignal + Reference signal + + Returns + ------- + torch.Tensor + Mel loss. + """ + loss = 0.0 + for n_mels, fmin, fmax, s in zip( + self.n_mels, self.mel_fmin, self.mel_fmax, self.stft_params + ): + kwargs = { + "window_length": s.window_length, + "hop_length": s.hop_length, + "window_type": s.window_type, + } + x_mels = x.mel_spectrogram(n_mels, mel_fmin=fmin, mel_fmax=fmax, **kwargs) + y_mels = y.mel_spectrogram(n_mels, mel_fmin=fmin, mel_fmax=fmax, **kwargs) + + loss += self.log_weight * self.loss_fn( + x_mels.clamp(self.clamp_eps).pow(self.pow).log10(), + y_mels.clamp(self.clamp_eps).pow(self.pow).log10(), + ) + loss += self.mag_weight * self.loss_fn(x_mels, y_mels) + return loss + + +class GANLoss(nn.Module): + """ + Computes a discriminator loss, given a discriminator on + generated waveforms/spectrograms compared to ground truth + waveforms/spectrograms. Computes the loss for both the + discriminator and the generator in separate functions. + """ + + def __init__(self, discriminator): + super().__init__() + self.discriminator = discriminator + + def forward(self, fake, real): + d_fake = self.discriminator(fake.audio_data) + d_real = self.discriminator(real.audio_data) + return d_fake, d_real + + def discriminator_loss(self, fake, real): + d_fake, d_real = self.forward(fake.clone().detach(), real) + + loss_d = 0 + for x_fake, x_real in zip(d_fake, d_real): + loss_d += torch.mean(x_fake[-1] ** 2) + loss_d += torch.mean((1 - x_real[-1]) ** 2) + return loss_d + + def generator_loss(self, fake, real): + d_fake, d_real = self.forward(fake, real) + + loss_g = 0 + for x_fake in d_fake: + loss_g += torch.mean((1 - x_fake[-1]) ** 2) + + loss_feature = 0 + + for i in range(len(d_fake)): + for j in range(len(d_fake[i]) - 1): + loss_feature += F.l1_loss(d_fake[i][j], d_real[i][j].detach()) + return loss_g, loss_feature diff --git a/dac/nn/quantize.py b/dac/nn/quantize.py new file mode 100644 index 0000000000000000000000000000000000000000..962029dc668fc98cbbea5bd19bf5ea6f41467875 --- /dev/null +++ b/dac/nn/quantize.py @@ -0,0 +1,339 @@ +from typing import Union + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange +from torch.nn.utils import weight_norm + +from dac.nn.layers import WNConv1d + +class VectorQuantizeLegacy(nn.Module): + """ + Implementation of VQ similar to Karpathy's repo: + https://github.com/karpathy/deep-vector-quantization + removed in-out projection + """ + + def __init__(self, input_dim: int, codebook_size: int): + super().__init__() + self.codebook_size = codebook_size + self.codebook = nn.Embedding(codebook_size, input_dim) + + def forward(self, z, z_mask=None): + """Quantized the input tensor using a fixed codebook and returns + the corresponding codebook vectors + + Parameters + ---------- + z : Tensor[B x D x T] + + Returns + ------- + Tensor[B x D x T] + Quantized continuous representation of input + Tensor[1] + Commitment loss to train encoder to predict vectors closer to codebook + entries + Tensor[1] + Codebook loss to update the codebook + Tensor[B x T] + Codebook indices (quantized discrete representation of input) + Tensor[B x D x T] + Projected latents (continuous representation of input before quantization) + """ + + z_e = z + z_q, indices = self.decode_latents(z) + + if z_mask is not None: + commitment_loss = (F.mse_loss(z_e, z_q.detach(), reduction="none").mean(1) * z_mask).sum() / z_mask.sum() + codebook_loss = (F.mse_loss(z_q, z_e.detach(), reduction="none").mean(1) * z_mask).sum() / z_mask.sum() + else: + commitment_loss = F.mse_loss(z_e, z_q.detach()) + codebook_loss = F.mse_loss(z_q, z_e.detach()) + z_q = ( + z_e + (z_q - z_e).detach() + ) # noop in forward pass, straight-through gradient estimator in backward pass + + return z_q, indices, z_e, commitment_loss, codebook_loss + + def embed_code(self, embed_id): + return F.embedding(embed_id, self.codebook.weight) + + def decode_code(self, embed_id): + return self.embed_code(embed_id).transpose(1, 2) + + def decode_latents(self, latents): + encodings = rearrange(latents, "b d t -> (b t) d") + codebook = self.codebook.weight # codebook: (N x D) + + # L2 normalize encodings and codebook (ViT-VQGAN) + encodings = F.normalize(encodings) + codebook = F.normalize(codebook) + + # Compute euclidean distance with codebook + dist = ( + encodings.pow(2).sum(1, keepdim=True) + - 2 * encodings @ codebook.t() + + codebook.pow(2).sum(1, keepdim=True).t() + ) + indices = rearrange((-dist).max(1)[1], "(b t) -> b t", b=latents.size(0)) + z_q = self.decode_code(indices) + return z_q, indices + +class VectorQuantize(nn.Module): + """ + Implementation of VQ similar to Karpathy's repo: + https://github.com/karpathy/deep-vector-quantization + Additionally uses following tricks from Improved VQGAN + (https://arxiv.org/pdf/2110.04627.pdf): + 1. Factorized codes: Perform nearest neighbor lookup in low-dimensional space + for improved codebook usage + 2. l2-normalized codes: Converts euclidean distance to cosine similarity which + improves training stability + """ + + def __init__(self, input_dim: int, codebook_size: int, codebook_dim: int): + super().__init__() + self.codebook_size = codebook_size + self.codebook_dim = codebook_dim + + self.in_proj = WNConv1d(input_dim, codebook_dim, kernel_size=1) + self.out_proj = WNConv1d(codebook_dim, input_dim, kernel_size=1) + self.codebook = nn.Embedding(codebook_size, codebook_dim) + + def forward(self, z, z_mask=None): + """Quantized the input tensor using a fixed codebook and returns + the corresponding codebook vectors + + Parameters + ---------- + z : Tensor[B x D x T] + + Returns + ------- + Tensor[B x D x T] + Quantized continuous representation of input + Tensor[1] + Commitment loss to train encoder to predict vectors closer to codebook + entries + Tensor[1] + Codebook loss to update the codebook + Tensor[B x T] + Codebook indices (quantized discrete representation of input) + Tensor[B x D x T] + Projected latents (continuous representation of input before quantization) + """ + + # Factorized codes (ViT-VQGAN) Project input into low-dimensional space + z_e = self.in_proj(z) # z_e : (B x D x T) + z_q, indices = self.decode_latents(z_e) + + if z_mask is not None: + commitment_loss = (F.mse_loss(z_e, z_q.detach(), reduction="none").mean(1) * z_mask).sum() / z_mask.sum() + codebook_loss = (F.mse_loss(z_q, z_e.detach(), reduction="none").mean(1) * z_mask).sum() / z_mask.sum() + else: + commitment_loss = F.mse_loss(z_e, z_q.detach()) + codebook_loss = F.mse_loss(z_q, z_e.detach()) + + z_q = ( + z_e + (z_q - z_e).detach() + ) # noop in forward pass, straight-through gradient estimator in backward pass + + z_q = self.out_proj(z_q) + + return z_q, commitment_loss, codebook_loss, indices, z_e + + def embed_code(self, embed_id): + return F.embedding(embed_id, self.codebook.weight) + + def decode_code(self, embed_id): + return self.embed_code(embed_id).transpose(1, 2) + + def decode_latents(self, latents): + encodings = rearrange(latents, "b d t -> (b t) d") + codebook = self.codebook.weight # codebook: (N x D) + + # L2 normalize encodings and codebook (ViT-VQGAN) + encodings = F.normalize(encodings) + codebook = F.normalize(codebook) + + # Compute euclidean distance with codebook + dist = ( + encodings.pow(2).sum(1, keepdim=True) + - 2 * encodings @ codebook.t() + + codebook.pow(2).sum(1, keepdim=True).t() + ) + indices = rearrange((-dist).max(1)[1], "(b t) -> b t", b=latents.size(0)) + z_q = self.decode_code(indices) + return z_q, indices + + +class ResidualVectorQuantize(nn.Module): + """ + Introduced in SoundStream: An end2end neural audio codec + https://arxiv.org/abs/2107.03312 + """ + + def __init__( + self, + input_dim: int = 512, + n_codebooks: int = 9, + codebook_size: int = 1024, + codebook_dim: Union[int, list] = 8, + quantizer_dropout: float = 0.0, + ): + super().__init__() + if isinstance(codebook_dim, int): + codebook_dim = [codebook_dim for _ in range(n_codebooks)] + + self.n_codebooks = n_codebooks + self.codebook_dim = codebook_dim + self.codebook_size = codebook_size + + self.quantizers = nn.ModuleList( + [ + VectorQuantize(input_dim, codebook_size, codebook_dim[i]) + for i in range(n_codebooks) + ] + ) + self.quantizer_dropout = quantizer_dropout + + def forward(self, z, n_quantizers: int = None): + """Quantized the input tensor using a fixed set of `n` codebooks and returns + the corresponding codebook vectors + Parameters + ---------- + z : Tensor[B x D x T] + n_quantizers : int, optional + No. of quantizers to use + (n_quantizers < self.n_codebooks ex: for quantizer dropout) + Note: if `self.quantizer_dropout` is True, this argument is ignored + when in training mode, and a random number of quantizers is used. + Returns + ------- + dict + A dictionary with the following keys: + + "z" : Tensor[B x D x T] + Quantized continuous representation of input + "codes" : Tensor[B x N x T] + Codebook indices for each codebook + (quantized discrete representation of input) + "latents" : Tensor[B x N*D x T] + Projected latents (continuous representation of input before quantization) + "vq/commitment_loss" : Tensor[1] + Commitment loss to train encoder to predict vectors closer to codebook + entries + "vq/codebook_loss" : Tensor[1] + Codebook loss to update the codebook + """ + z_q = 0 + residual = z + commitment_loss = 0 + codebook_loss = 0 + + codebook_indices = [] + latents = [] + + if n_quantizers is None: + n_quantizers = self.n_codebooks + if self.training: + n_quantizers = torch.ones((z.shape[0],)) * self.n_codebooks + 1 + dropout = torch.randint(1, self.n_codebooks + 1, (z.shape[0],)) + n_dropout = int(z.shape[0] * self.quantizer_dropout) + n_quantizers[:n_dropout] = dropout[:n_dropout] + n_quantizers = n_quantizers.to(z.device) + + for i, quantizer in enumerate(self.quantizers): + if self.training is False and i >= n_quantizers: + break + + z_q_i, commitment_loss_i, codebook_loss_i, indices_i, z_e_i = quantizer( + residual + ) + + # Create mask to apply quantizer dropout + mask = ( + torch.full((z.shape[0],), fill_value=i, device=z.device) < n_quantizers + ) + z_q = z_q + z_q_i * mask[:, None, None] + residual = residual - z_q_i + + # Sum losses + commitment_loss += (commitment_loss_i * mask).mean() + codebook_loss += (codebook_loss_i * mask).mean() + + codebook_indices.append(indices_i) + latents.append(z_e_i) + + codes = torch.stack(codebook_indices, dim=1) + latents = torch.cat(latents, dim=1) + + return z_q, codes, latents, commitment_loss, codebook_loss + + def from_codes(self, codes: torch.Tensor): + """Given the quantized codes, reconstruct the continuous representation + Parameters + ---------- + codes : Tensor[B x N x T] + Quantized discrete representation of input + Returns + ------- + Tensor[B x D x T] + Quantized continuous representation of input + """ + z_q = 0.0 + z_p = [] + n_codebooks = codes.shape[1] + for i in range(n_codebooks): + z_p_i = self.quantizers[i].decode_code(codes[:, i, :]) + z_p.append(z_p_i) + + z_q_i = self.quantizers[i].out_proj(z_p_i) + z_q = z_q + z_q_i + return z_q, torch.cat(z_p, dim=1), codes + + def from_latents(self, latents: torch.Tensor): + """Given the unquantized latents, reconstruct the + continuous representation after quantization. + + Parameters + ---------- + latents : Tensor[B x N x T] + Continuous representation of input after projection + + Returns + ------- + Tensor[B x D x T] + Quantized representation of full-projected space + Tensor[B x D x T] + Quantized representation of latent space + """ + z_q = 0 + z_p = [] + codes = [] + dims = np.cumsum([0] + [q.codebook_dim for q in self.quantizers]) + + n_codebooks = np.where(dims <= latents.shape[1])[0].max(axis=0, keepdims=True)[ + 0 + ] + for i in range(n_codebooks): + j, k = dims[i], dims[i + 1] + z_p_i, codes_i = self.quantizers[i].decode_latents(latents[:, j:k, :]) + z_p.append(z_p_i) + codes.append(codes_i) + + z_q_i = self.quantizers[i].out_proj(z_p_i) + z_q = z_q + z_q_i + + return z_q, torch.cat(z_p, dim=1), torch.stack(codes, dim=1) + + +if __name__ == "__main__": + rvq = ResidualVectorQuantize(quantizer_dropout=True) + x = torch.randn(16, 512, 80) + y = rvq(x) + print(y["latents"].shape) diff --git a/dac/utils/__init__.py b/dac/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f6dbede7b3e57655b4a732c92c84d7c72fbc9059 --- /dev/null +++ b/dac/utils/__init__.py @@ -0,0 +1,123 @@ +from pathlib import Path + +import argbind +from audiotools import ml + +import dac + +DAC = dac.model.DAC +Accelerator = ml.Accelerator + +__MODEL_LATEST_TAGS__ = { + ("44khz", "8kbps"): "0.0.1", + ("24khz", "8kbps"): "0.0.4", + ("16khz", "8kbps"): "0.0.5", + ("44khz", "16kbps"): "1.0.0", +} + +__MODEL_URLS__ = { + ( + "44khz", + "0.0.1", + "8kbps", + ): "https://github.com/descriptinc/descript-audio-codec/releases/download/0.0.1/weights.pth", + ( + "24khz", + "0.0.4", + "8kbps", + ): "https://github.com/descriptinc/descript-audio-codec/releases/download/0.0.4/weights_24khz.pth", + ( + "16khz", + "0.0.5", + "8kbps", + ): "https://github.com/descriptinc/descript-audio-codec/releases/download/0.0.5/weights_16khz.pth", + ( + "44khz", + "1.0.0", + "16kbps", + ): "https://github.com/descriptinc/descript-audio-codec/releases/download/1.0.0/weights_44khz_16kbps.pth", +} + + +@argbind.bind(group="download", positional=True, without_prefix=True) +def download( + model_type: str = "44khz", model_bitrate: str = "8kbps", tag: str = "latest" +): + """ + Function that downloads the weights file from URL if a local cache is not found. + + Parameters + ---------- + model_type : str + The type of model to download. Must be one of "44khz", "24khz", or "16khz". Defaults to "44khz". + model_bitrate: str + Bitrate of the model. Must be one of "8kbps", or "16kbps". Defaults to "8kbps". + Only 44khz model supports 16kbps. + tag : str + The tag of the model to download. Defaults to "latest". + + Returns + ------- + Path + Directory path required to load model via audiotools. + """ + model_type = model_type.lower() + tag = tag.lower() + + assert model_type in [ + "44khz", + "24khz", + "16khz", + ], "model_type must be one of '44khz', '24khz', or '16khz'" + + assert model_bitrate in [ + "8kbps", + "16kbps", + ], "model_bitrate must be one of '8kbps', or '16kbps'" + + if tag == "latest": + tag = __MODEL_LATEST_TAGS__[(model_type, model_bitrate)] + + download_link = __MODEL_URLS__.get((model_type, tag, model_bitrate), None) + + if download_link is None: + raise ValueError( + f"Could not find model with tag {tag} and model type {model_type}" + ) + + local_path = ( + Path.home() + / ".cache" + / "descript" + / "dac" + / f"weights_{model_type}_{model_bitrate}_{tag}.pth" + ) + if not local_path.exists(): + local_path.parent.mkdir(parents=True, exist_ok=True) + + # Download the model + import requests + + response = requests.get(download_link) + + if response.status_code != 200: + raise ValueError( + f"Could not download model. Received response code {response.status_code}" + ) + local_path.write_bytes(response.content) + + return local_path + + +def load_model( + model_type: str = "44khz", + model_bitrate: str = "8kbps", + tag: str = "latest", + load_path: str = None, +): + if not load_path: + load_path = download( + model_type=model_type, model_bitrate=model_bitrate, tag=tag + ) + generator = DAC.load(load_path) + return generator diff --git a/dac/utils/decode.py b/dac/utils/decode.py new file mode 100644 index 0000000000000000000000000000000000000000..48c25298fd0f9a53f19c1d6d260f1f0563c7b7e4 --- /dev/null +++ b/dac/utils/decode.py @@ -0,0 +1,95 @@ +import warnings +from pathlib import Path + +import argbind +import numpy as np +import torch +from audiotools import AudioSignal +from tqdm import tqdm + +from dac import DACFile +from dac.utils import load_model + +warnings.filterwarnings("ignore", category=UserWarning) + + +@argbind.bind(group="decode", positional=True, without_prefix=True) +@torch.inference_mode() +@torch.no_grad() +def decode( + input: str, + output: str = "", + weights_path: str = "", + model_tag: str = "latest", + model_bitrate: str = "8kbps", + device: str = "cuda", + model_type: str = "44khz", + verbose: bool = False, +): + """Decode audio from codes. + + Parameters + ---------- + input : str + Path to input directory or file + output : str, optional + Path to output directory, by default "". + If `input` is a directory, the directory sub-tree relative to `input` is re-created in `output`. + weights_path : str, optional + Path to weights file, by default "". If not specified, the weights file will be downloaded from the internet using the + model_tag and model_type. + model_tag : str, optional + Tag of the model to use, by default "latest". Ignored if `weights_path` is specified. + model_bitrate: str + Bitrate of the model. Must be one of "8kbps", or "16kbps". Defaults to "8kbps". + device : str, optional + Device to use, by default "cuda". If "cpu", the model will be loaded on the CPU. + model_type : str, optional + The type of model to use. Must be one of "44khz", "24khz", or "16khz". Defaults to "44khz". Ignored if `weights_path` is specified. + """ + generator = load_model( + model_type=model_type, + model_bitrate=model_bitrate, + tag=model_tag, + load_path=weights_path, + ) + generator.to(device) + generator.eval() + + # Find all .dac files in input directory + _input = Path(input) + input_files = list(_input.glob("**/*.dac")) + + # If input is a .dac file, add it to the list + if _input.suffix == ".dac": + input_files.append(_input) + + # Create output directory + output = Path(output) + output.mkdir(parents=True, exist_ok=True) + + for i in tqdm(range(len(input_files)), desc=f"Decoding files"): + # Load file + artifact = DACFile.load(input_files[i]) + + # Reconstruct audio from codes + recons = generator.decompress(artifact, verbose=verbose) + + # Compute output path + relative_path = input_files[i].relative_to(input) + output_dir = output / relative_path.parent + if not relative_path.name: + output_dir = output + relative_path = input_files[i] + output_name = relative_path.with_suffix(".wav").name + output_path = output_dir / output_name + output_path.parent.mkdir(parents=True, exist_ok=True) + + # Write to file + recons.write(output_path) + + +if __name__ == "__main__": + args = argbind.parse_args() + with argbind.scope(args): + decode() diff --git a/dac/utils/encode.py b/dac/utils/encode.py new file mode 100644 index 0000000000000000000000000000000000000000..7c9a8b582b24c194ba82c2d280c4df4fa2cfff87 --- /dev/null +++ b/dac/utils/encode.py @@ -0,0 +1,94 @@ +import math +import warnings +from pathlib import Path + +import argbind +import numpy as np +import torch +from audiotools import AudioSignal +from audiotools.core import util +from tqdm import tqdm + +from dac.utils import load_model + +warnings.filterwarnings("ignore", category=UserWarning) + + +@argbind.bind(group="encode", positional=True, without_prefix=True) +@torch.inference_mode() +@torch.no_grad() +def encode( + input: str, + output: str = "", + weights_path: str = "", + model_tag: str = "latest", + model_bitrate: str = "8kbps", + n_quantizers: int = None, + device: str = "cuda", + model_type: str = "44khz", + win_duration: float = 5.0, + verbose: bool = False, +): + """Encode audio files in input path to .dac format. + + Parameters + ---------- + input : str + Path to input audio file or directory + output : str, optional + Path to output directory, by default "". If `input` is a directory, the directory sub-tree relative to `input` is re-created in `output`. + weights_path : str, optional + Path to weights file, by default "". If not specified, the weights file will be downloaded from the internet using the + model_tag and model_type. + model_tag : str, optional + Tag of the model to use, by default "latest". Ignored if `weights_path` is specified. + model_bitrate: str + Bitrate of the model. Must be one of "8kbps", or "16kbps". Defaults to "8kbps". + n_quantizers : int, optional + Number of quantizers to use, by default None. If not specified, all the quantizers will be used and the model will compress at maximum bitrate. + device : str, optional + Device to use, by default "cuda" + model_type : str, optional + The type of model to use. Must be one of "44khz", "24khz", or "16khz". Defaults to "44khz". Ignored if `weights_path` is specified. + """ + generator = load_model( + model_type=model_type, + model_bitrate=model_bitrate, + tag=model_tag, + load_path=weights_path, + ) + generator.to(device) + generator.eval() + kwargs = {"n_quantizers": n_quantizers} + + # Find all audio files in input path + input = Path(input) + audio_files = util.find_audio(input) + + output = Path(output) + output.mkdir(parents=True, exist_ok=True) + + for i in tqdm(range(len(audio_files)), desc="Encoding files"): + # Load file + signal = AudioSignal(audio_files[i]) + + # Encode audio to .dac format + artifact = generator.compress(signal, win_duration, verbose=verbose, **kwargs) + + # Compute output path + relative_path = audio_files[i].relative_to(input) + output_dir = output / relative_path.parent + if not relative_path.name: + output_dir = output + relative_path = audio_files[i] + output_name = relative_path.with_suffix(".dac").name + output_path = output_dir / output_name + output_path.parent.mkdir(parents=True, exist_ok=True) + + artifact.save(output_path) + + +if __name__ == "__main__": + args = argbind.parse_args() + with argbind.scope(args): + encode() diff --git a/data/ft_dataset.py b/data/ft_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..f0ed2351cf534f1010dd5a105724a6f5616f586c --- /dev/null +++ b/data/ft_dataset.py @@ -0,0 +1,133 @@ +import torch +import librosa +import numpy as np +import random +import os +from torch.utils.data import DataLoader +from modules.audio import mel_spectrogram + + +duration_setting = { + "min": 1.0, + "max": 30.0, +} +# assume single speaker +class FT_Dataset(torch.utils.data.Dataset): + def __init__(self, + data_path, + spect_params, + sr=22050, + batch_size=1, + ): + self.data_path = data_path + # recursively find all files in data_path + self.data = [] + for root, _, files in os.walk(data_path): + for file in files: + if (file.endswith(".wav") or + file.endswith(".mp3") or + file.endswith(".flac") or + file.endswith(".ogg") or + file.endswith(".m4a") or + file.endswith(".opus")): + self.data.append(os.path.join(root, file)) + + mel_fn_args = { + "n_fft": spect_params['n_fft'], + "win_size": spect_params['win_length'], + "hop_size": spect_params['hop_length'], + "num_mels": spect_params['n_mels'], + "sampling_rate": sr, + "fmin": spect_params['fmin'], + "fmax": None if spect_params['fmax'] == "None" else spect_params['fmax'], + "center": False + } + self.to_mel = lambda x: mel_spectrogram(x, **mel_fn_args) + self.sr = sr + + assert len(self.data) != 0 + # if dataset length is less than batch size, repeat the dataset + while len(self.data) < batch_size: + self.data += self.data + + def __len__(self): + return len(self.data) + + def __getitem__(self, idx): + idx = idx % len(self.data) + wav_path = self.data[idx] + try: + speech, orig_sr = librosa.load(wav_path, sr=self.sr) + except Exception as e: + print(f"Failed to load wav file with error {e}") + return self.__getitem__(random.randint(0, len(self))) + if len(speech) < self.sr * duration_setting["min"] or len(speech) > self.sr * duration_setting["max"]: + print(f"Audio {wav_path} is too short or too long, skipping") + return self.__getitem__(random.randint(0, len(self))) + return_dict = { + 'audio': speech, + 'sr': orig_sr + } + wave, orig_sr = return_dict['audio'], return_dict['sr'] + if orig_sr != self.sr: + wave = librosa.resample(wave, orig_sr, self.sr) + wave = torch.from_numpy(wave).float() + mel = self.to_mel(wave.unsqueeze(0)).squeeze(0) + + return wave, mel + +def build_ft_dataloader(data_path, spect_params, sr, batch_size=1, num_workers=0): + dataset = FT_Dataset(data_path, spect_params, sr, batch_size) + dataloader = torch.utils.data.DataLoader( + dataset, + batch_size=batch_size, + shuffle=True, + num_workers=num_workers, + collate_fn=collate, + ) + return dataloader + +def collate(batch): + batch_size = len(batch) + + # sort by mel length + lengths = [b[1].shape[1] for b in batch] + batch_indexes = np.argsort(lengths)[::-1] + batch = [batch[bid] for bid in batch_indexes] + + nmels = batch[0][1].size(0) + max_mel_length = max([b[1].shape[1] for b in batch]) + max_wave_length = max([b[0].size(0) for b in batch]) + + mels = torch.zeros((batch_size, nmels, max_mel_length)).float() - 10 + waves = torch.zeros((batch_size, max_wave_length)).float() + + mel_lengths = torch.zeros(batch_size).long() + wave_lengths = torch.zeros(batch_size).long() + + for bid, (wave, mel) in enumerate(batch): + mel_size = mel.size(1) + mels[bid, :, :mel_size] = mel + waves[bid, : wave.size(0)] = wave + mel_lengths[bid] = mel_size + wave_lengths[bid] = wave.size(0) + + return waves, mels, wave_lengths, mel_lengths + +if __name__ == "__main__": + data_path = "./example/reference" + sr = 22050 + spect_params = { + "n_fft": 1024, + "win_length": 1024, + "hop_length": 256, + "n_mels": 80, + "fmin": 0, + "fmax": 8000, + } + dataloader = build_ft_dataloader(data_path, spect_params, sr, batch_size=2, num_workers=0) + for idx, batch in enumerate(dataloader): + wave, mel, wave_lengths, mel_lengths = batch + print(wave.shape, mel.shape) + if idx == 10: + break \ No newline at end of file diff --git a/eval.py b/eval.py new file mode 100644 index 0000000000000000000000000000000000000000..914128fe1cd6a70def6e2ee4ae3ac0de18157a0d --- /dev/null +++ b/eval.py @@ -0,0 +1,550 @@ +import shutil +import warnings +import argparse +import torch +import os +import os.path as osp +import yaml + +warnings.simplefilter("ignore") + +# load packages +import random + +from tqdm import tqdm +from modules.commons import * +import time + +import torchaudio +import librosa +import torchaudio.compliance.kaldi as kaldi + +from hf_utils import load_custom_model_from_hf +from resemblyzer import preprocess_wav, VoiceEncoder + +# Load model and configuration +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + +from transformers import Wav2Vec2FeatureExtractor, WavLMForXVector +from transformers import Wav2Vec2Processor, HubertForCTC + +import jiwer +import string + +from baselines.dnsmos.dnsmos_computor import DNSMOSComputer + +def calc_mos(computor, audio, orin_sr): + # only 16k audio is supported + target_sr = 16000 + if orin_sr != 16000: + audio = librosa.resample( + audio, orig_sr=orin_sr, target_sr=target_sr, res_type="kaiser_fast" + ) + result = computor.compute(audio, target_sr, False) + sig, bak, ovr = result["SIG"], result["BAK"], result["OVRL"] + + if ovr == 0: + print("calculate dns mos failed") + return sig, bak, ovr + +mos_computer = DNSMOSComputer( + "baselines/dnsmos/sig_bak_ovr.onnx", + "baselines/dnsmos/model_v8.onnx", + device="cuda", + device_id=0, +) + +def load_models(args): + dit_checkpoint_path, dit_config_path = load_custom_model_from_hf("Plachta/Seed-VC", + "DiT_seed_v2_uvit_whisper_small_wavenet_bigvgan_pruned.pth", + "config_dit_mel_seed_uvit_whisper_small_wavenet.yml") + config = yaml.safe_load(open(dit_config_path, "r")) + model_params = recursive_munch(config["model_params"]) + model = build_model(model_params, stage="DiT") + hop_length = config["preprocess_params"]["spect_params"]["hop_length"] + sr = config["preprocess_params"]["sr"] + + # Load checkpoints + model, _, _, _ = load_checkpoint( + model, + None, + dit_checkpoint_path, + load_only_params=True, + ignore_modules=[], + is_distributed=False, + ) + for key in model: + model[key].eval() + model[key].to(device) + model.cfm.estimator.setup_caches(max_batch_size=1, max_seq_length=8192) + + # Load additional modules + from modules.campplus.DTDNN import CAMPPlus + + campplus_ckpt_path = load_custom_model_from_hf( + "funasr/campplus", "campplus_cn_common.bin", config_filename=None + ) + campplus_model = CAMPPlus(feat_dim=80, embedding_size=192) + campplus_model.load_state_dict(torch.load(campplus_ckpt_path, map_location="cpu")) + campplus_model.eval() + campplus_model.to(device) + + vocoder_type = model_params.vocoder.type + + if vocoder_type == 'bigvgan': + from modules.bigvgan import bigvgan + bigvgan_name = model_params.vocoder.name + bigvgan_model = bigvgan.BigVGAN.from_pretrained(bigvgan_name, use_cuda_kernel=False) + # remove weight norm in the model and set to eval mode + bigvgan_model.remove_weight_norm() + bigvgan_model = bigvgan_model.eval().to(device) + vocoder_fn = bigvgan_model + elif vocoder_type == 'hifigan': + from modules.hifigan.generator import HiFTGenerator + from modules.hifigan.f0_predictor import ConvRNNF0Predictor + hift_config = yaml.safe_load(open('configs/hifigan.yml', 'r')) + hift_gen = HiFTGenerator(**hift_config['hift'], f0_predictor=ConvRNNF0Predictor(**hift_config['f0_predictor'])) + hift_gen.load_state_dict(torch.load(hift_config['pretrained_model_path'], map_location='cpu')) + hift_gen.eval() + hift_gen.to(device) + vocoder_fn = hift_gen + elif vocoder_type == "vocos": + vocos_config = yaml.safe_load(open(model_params.vocoder.vocos.config, 'r')) + vocos_path = model_params.vocoder.vocos.path + vocos_model_params = recursive_munch(vocos_config['model_params']) + vocos = build_model(vocos_model_params, stage='mel_vocos') + vocos_checkpoint_path = vocos_path + vocos, _, _, _ = load_checkpoint(vocos, None, vocos_checkpoint_path, + load_only_params=True, ignore_modules=[], is_distributed=False) + _ = [vocos[key].eval().to(device) for key in vocos] + _ = [vocos[key].to(device) for key in vocos] + total_params = sum(sum(p.numel() for p in vocos[key].parameters() if p.requires_grad) for key in vocos.keys()) + print(f"Vocoder model total parameters: {total_params / 1_000_000:.2f}M") + vocoder_fn = vocos.decoder + else: + raise ValueError(f"Unsupported vocoder type: {vocoder_type}") + + speech_tokenizer_type = model_params.speech_tokenizer.type + if speech_tokenizer_type == 'whisper': + # whisper + from transformers import AutoFeatureExtractor, WhisperModel + whisper_name = model_params.speech_tokenizer.name + whisper_model = WhisperModel.from_pretrained(whisper_name, torch_dtype=torch.float16).to(device) + del whisper_model.decoder + whisper_feature_extractor = AutoFeatureExtractor.from_pretrained(whisper_name) + + def semantic_fn(waves_16k): + ori_inputs = whisper_feature_extractor([waves_16k.squeeze(0).cpu().numpy()], + return_tensors="pt", + return_attention_mask=True) + ori_input_features = whisper_model._mask_input_features( + ori_inputs.input_features, attention_mask=ori_inputs.attention_mask).to(device) + with torch.no_grad(): + ori_outputs = whisper_model.encoder( + ori_input_features.to(whisper_model.encoder.dtype), + head_mask=None, + output_attentions=False, + output_hidden_states=False, + return_dict=True, + ) + S_ori = ori_outputs.last_hidden_state.to(torch.float32) + S_ori = S_ori[:, :waves_16k.size(-1) // 320 + 1] + return S_ori + elif speech_tokenizer_type == 'cnhubert': + from transformers import ( + Wav2Vec2FeatureExtractor, + HubertModel, + ) + hubert_model_name = config['model_params']['speech_tokenizer']['name'] + hubert_feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(hubert_model_name) + hubert_model = HubertModel.from_pretrained(hubert_model_name) + hubert_model = hubert_model.to(device) + hubert_model = hubert_model.eval() + hubert_model = hubert_model.half() + + def semantic_fn(waves_16k): + ori_waves_16k_input_list = [ + waves_16k[bib].cpu().numpy() + for bib in range(len(waves_16k)) + ] + ori_inputs = hubert_feature_extractor(ori_waves_16k_input_list, + return_tensors="pt", + return_attention_mask=True, + padding=True, + sampling_rate=16000).to(device) + with torch.no_grad(): + ori_outputs = hubert_model( + ori_inputs.input_values.half(), + ) + S_ori = ori_outputs.last_hidden_state.float() + return S_ori + elif speech_tokenizer_type == 'xlsr': + from transformers import ( + Wav2Vec2FeatureExtractor, + Wav2Vec2Model, + ) + model_name = config['model_params']['speech_tokenizer']['name'] + output_layer = config['model_params']['speech_tokenizer']['output_layer'] + wav2vec_feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(model_name) + wav2vec_model = Wav2Vec2Model.from_pretrained(model_name) + wav2vec_model.encoder.layers = wav2vec_model.encoder.layers[:output_layer] + wav2vec_model = wav2vec_model.to(device) + wav2vec_model = wav2vec_model.eval() + wav2vec_model = wav2vec_model.half() + + def semantic_fn(waves_16k): + ori_waves_16k_input_list = [ + waves_16k[bib].cpu().numpy() + for bib in range(len(waves_16k)) + ] + ori_inputs = wav2vec_feature_extractor(ori_waves_16k_input_list, + return_tensors="pt", + return_attention_mask=True, + padding=True, + sampling_rate=16000).to(device) + with torch.no_grad(): + ori_outputs = wav2vec_model( + ori_inputs.input_values.half(), + ) + S_ori = ori_outputs.last_hidden_state.float() + return S_ori + else: + raise ValueError(f"Unsupported speech tokenizer type: {model_params.speech_tokenizer.type}") + # Generate mel spectrograms + mel_fn_args = { + "n_fft": config['preprocess_params']['spect_params']['n_fft'], + "win_size": config['preprocess_params']['spect_params']['win_length'], + "hop_size": config['preprocess_params']['spect_params']['hop_length'], + "num_mels": config['preprocess_params']['spect_params']['n_mels'], + "sampling_rate": sr, + "fmin": config['preprocess_params'].get('fmin', 0), + "fmax": None if config['preprocess_params']['spect_params'].get('fmax', "None") == "None" else 8000, + "center": False + } + from modules.audio import mel_spectrogram + + to_mel = lambda x: mel_spectrogram(x, **mel_fn_args) + + return ( + model, + semantic_fn, + vocoder_fn, + campplus_model, + to_mel, + mel_fn_args, + ) + + +@torch.no_grad() +def main(args): + # init xvector models + if args.xvector_extractor == "wavlm": + wavlm_feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained( + "microsoft/wavlm-base-plus-sv" + ) + wavlm_model = WavLMForXVector.from_pretrained( + "microsoft/wavlm-base-plus-sv" + ).to(device) + elif args.xvector_extractor == "resemblyzer": + resemblyzer_encoder = VoiceEncoder() + elif args.xvector_extractor == 'wavlm-large': + import sys + sys.path.append("../UniSpeech/downstreams/speaker_verification") + from verification import init_model + wavlm_model = init_model("wavlm_large", "D:/wavlm_large_finetune.pth") + wavlm_model.cuda() + wavlm_model.eval() + else: + raise ValueError(f"Unknown xvector extractor: {args.xvector_extractor}") + + # init asr model + asr_processor = Wav2Vec2Processor.from_pretrained("facebook/hubert-large-ls960-ft") + asr_model = HubertForCTC.from_pretrained("facebook/hubert-large-ls960-ft").to(device) + + ( + model, + semantic_fn, + vocoder_fn, + campplus_model, + to_mel, + mel_fn_args, + ) = load_models(args) + sr = mel_fn_args["sampling_rate"] + + source_dir = args.source + target_dir = args.target + diffusion_steps = args.diffusion_steps + length_adjust = args.length_adjust + inference_cfg_rate = args.inference_cfg_rate + baseline = args.baseline + max_samples = args.max_samples + try: + source_audio_list = open(osp.join(source_dir, "index.tsv"), "r").readlines() + except FileNotFoundError: + source_audio_list = os.listdir(source_dir) + source_audio_list = [f for f in source_audio_list if f.endswith(".wav")] + target_audio_list = os.listdir(target_dir) + + conversion_result_dir = args.output + if baseline: + conversion_result_dir = os.path.join(conversion_result_dir, baseline) + os.makedirs(conversion_result_dir, exist_ok=True) + + similarity_list = [] + gt_wer_list = [] + gt_cer_list = [] + vc_wer_list = [] + vc_cer_list = [] + dnsmos_list = [] + for source_i, source_line in enumerate(tqdm(source_audio_list)): + if source_i >= max_samples: + break + source_index, source_transcript = source_line.strip().split("\t") + source_path = osp.join(source_dir, f"{source_index}.wav") + for target_i, target_name in enumerate(target_audio_list): + target_path = osp.join(target_dir, target_name) + print(f"Processing {source_path} -> {target_path}") + + if os.path.exists(osp.join(conversion_result_dir, source_index, f"{target_name}")): + # already converted, load the converted file + vc_wave_16k, _ = librosa.load( + osp.join(conversion_result_dir, source_index, f"{target_name}"), sr=16000 + ) + vc_wave_16k = torch.tensor(vc_wave_16k).unsqueeze(0) + ref_waves_16k, _ = librosa.load(target_path, sr=16000) + ref_waves_16k = torch.tensor(ref_waves_16k).unsqueeze(0) + else: + if baseline == "openvoice": + from baselines.openvoice import convert as openvoice_convert + ref_waves_16k, vc_wave_16k = openvoice_convert(source_path, target_path, "temp.wav") + elif baseline == "cosyvoice": + from baselines.cosyvoice import convert as cosyvoice_convert + ref_waves_16k, vc_wave_16k = cosyvoice_convert(source_path, target_path, "temp.wav") + else: + ref_waves_16k, vc_wave = convert( + source_path, + target_path, + model, + semantic_fn, + vocoder_fn, + campplus_model, + to_mel, + mel_fn_args, + sr, + length_adjust, + diffusion_steps, + inference_cfg_rate, + remove_prompt=args.remove_prompt, + ) + vc_wave_16k = torchaudio.functional.resample(vc_wave, sr, 16000) + os.makedirs(osp.join(conversion_result_dir, source_index), exist_ok=True) + torchaudio.save( + osp.join(conversion_result_dir, source_index, f"{target_name}"), + vc_wave_16k.cpu(), + 16000, + ) + if args.xvector_extractor == "wavlm": + ref_inputs = wavlm_feature_extractor( + ref_waves_16k.squeeze(0).cpu(), padding=True, return_tensors="pt" + ).to(device) + ref_embeddings = wavlm_model(**ref_inputs).embeddings + ref_embeddings = torch.nn.functional.normalize(ref_embeddings, dim=-1).cpu() + + vc_inputs = wavlm_feature_extractor( + vc_wave_16k.squeeze(0).cpu(), padding=True, return_tensors="pt" + ).to(device) + vc_embeddings = wavlm_model(**vc_inputs).embeddings + vc_embeddings = torch.nn.functional.normalize(vc_embeddings, dim=-1).cpu() + + similarity = torch.nn.functional.cosine_similarity( + ref_embeddings, vc_embeddings, dim=-1 + ) + elif args.xvector_extractor == "resemblyzer": + ref_wav_resemblyzer = preprocess_wav(target_path) + vc_wav_resemblyzer = preprocess_wav( + osp.join(conversion_result_dir, source_index, f"{target_name}") + ) + ref_embed = resemblyzer_encoder.embed_utterance(ref_wav_resemblyzer) + vc_embed = resemblyzer_encoder.embed_utterance(vc_wav_resemblyzer) + similarity = np.inner(ref_embed, vc_embed) + elif args.xvector_extractor == 'wavlm-large': + ref_embed = wavlm_model(ref_waves_16k.to(device)).cpu() + vc_embed = wavlm_model(vc_wave_16k.to(device)).cpu() + similarity = torch.nn.functional.cosine_similarity(ref_embed, vc_embed, dim=-1) + else: + raise ValueError(f"Unknown xvector extractor: {args.xvector_extractor}") + print(f"Similarity: {similarity}") + similarity_list.append(similarity) + + # perform asr + vc_asr_inputs = asr_processor( + vc_wave_16k.squeeze(0).cpu(), return_tensors="pt", padding=True + ).to(device) + vc_asr_logits = asr_model(**vc_asr_inputs).logits + predicted_ids = torch.argmax(vc_asr_logits, dim=-1) + vc_transcription = asr_processor.decode(predicted_ids[0]) + + # perform asr on source 16k + source_wav_16k = librosa.load(source_path, sr=16000)[0] + source_asr_inputs = asr_processor( + source_wav_16k, return_tensors="pt", padding=True + ).to(device) + source_asr_logits = asr_model(**source_asr_inputs).logits + source_predicted_ids = torch.argmax(source_asr_logits, dim=-1) + source_transcription = asr_processor.decode(source_predicted_ids[0]) + + # convert transcriptions to all lower to calculate WER and CER + source_transcript = source_transcript.lower() + # remove punctuations in source_transcript + source_transcript = source_transcript.translate(str.maketrans("", "", string.punctuation)) + source_transcription = source_transcription.lower() + vc_transcription = vc_transcription.lower() + + # calculate WER and CER + gt_wer = jiwer.wer(source_transcript, source_transcription) + gt_cer = jiwer.cer(source_transcript, source_transcription) + vc_wer = jiwer.wer(source_transcript, vc_transcription) + vc_cer = jiwer.cer(source_transcript, vc_transcription) + + print(f"GT WER: {gt_wer}, CER: {gt_cer}") + print(f"VC WER: {vc_wer}, CER: {vc_cer}") + gt_wer_list.append(gt_wer) + gt_cer_list.append(gt_cer) + vc_wer_list.append(vc_wer) + vc_cer_list.append(vc_cer) + + # calculate dnsmos + sig, bak, ovr = calc_mos(mos_computer, vc_wave_16k.squeeze(0).cpu().numpy(), 16000) + dnsmos_list.append((sig, bak, ovr)) + + print(f"Average GT WER: {sum(gt_wer_list) / len(gt_wer_list)}") + print(f"Average GT CER: {sum(gt_cer_list) / len(gt_cer_list)}") + print(f"Average VC WER: {sum(vc_wer_list) / len(vc_wer_list)}") + print(f"Average VC CER: {sum(vc_cer_list) / len(vc_cer_list)}") + print(f"Average similarity: {sum(similarity_list) / len(similarity_list)}") + + print(f"Average DNS MOS SIG: {sum([x[0] for x in dnsmos_list]) / len(dnsmos_list)}") + print(f"Average DNS MOS BAK: {sum([x[1] for x in dnsmos_list]) / len(dnsmos_list)}") + print(f"Average DNS MOS OVR: {sum([x[2] for x in dnsmos_list]) / len(dnsmos_list)}") + + # save wer and cer result into this directory as a txt + with open(osp.join(conversion_result_dir, source_index, "result.txt"), 'w') as f: + f.write(f"GT WER: {sum(gt_wer_list[-len(target_audio_list):]) / len(target_audio_list)}\n") + f.write(f"GT CER: {sum(gt_cer_list[-len(target_audio_list):]) / len(target_audio_list)}\n") + f.write(f"VC WER: {sum(vc_wer_list[-len(target_audio_list):]) / len(target_audio_list)}\n") + f.write(f"VC CER: {sum(vc_cer_list[-len(target_audio_list):]) / len(target_audio_list)}\n") + f.write(f"Average similarity: {sum(similarity_list[-len(target_audio_list):]) / len(target_audio_list)}\n") + + print(f"Average WER: {sum(gt_wer_list) / len(gt_wer_list)}") + print(f"Average CER: {sum(gt_cer_list) / len(gt_cer_list)}") + print(f"Average WER: {sum(vc_wer_list) / len(vc_wer_list)}") + print(f"Average CER: {sum(vc_cer_list) / len(vc_cer_list)}") + print(f"Average similarity: {sum(similarity_list) / len(similarity_list)}") + # save similarity list + with open(osp.join(conversion_result_dir, f"{args.xvector_extractor}_similarity.tsv"), "w") as f: + f.write("\n".join([str(s) for s in similarity_list])) + # save wer and cer result into this directory as a txt + with open(osp.join(conversion_result_dir, "result.txt"), 'w') as f: + f.write(f"GT WER: {sum(gt_wer_list) / len(gt_wer_list)}\n") + f.write(f"GT CER: {sum(gt_cer_list) / len(gt_cer_list)}\n") + f.write(f"VC WER: {sum(vc_wer_list) / len(vc_wer_list)}\n") + f.write(f"VC CER: {sum(vc_cer_list) / len(vc_cer_list)}\n") + + print(f"Average DNS MOS SIG: {sum([x[0] for x in dnsmos_list]) / len(dnsmos_list)}") + print(f"Average DNS MOS BAK: {sum([x[1] for x in dnsmos_list]) / len(dnsmos_list)}") + print(f"Average DNS MOS OVR: {sum([x[2] for x in dnsmos_list]) / len(dnsmos_list)}") + + +def convert( + source_path, + target_path, + model, + semantic_fn, + vocoder_fn, + campplus_model, + to_mel, + mel_fn_args, + sr, + length_adjust, + diffusion_steps, + inference_cfg_rate, + remove_prompt=False, +): + source_audio = librosa.load(source_path, sr=sr)[0] + ref_audio = librosa.load(target_path, sr=sr)[0] + # decoded_wav = encodec_model.decoder(encodec_latent) + # torchaudio.save("test.wav", decoded_wav.cpu().squeeze(0), 24000) + # crop only the first 30 seconds + source_audio = torch.tensor(source_audio).unsqueeze(0).float().to(device) + ref_audio = torch.tensor(ref_audio).unsqueeze(0).float().to(device) + + if source_audio.size(1) + ref_audio.size(1) > 30 * sr: + print(f"reference audio clipped from {ref_audio.size(1)/sr} seconds to {30 * sr - source_audio.size(1)} seconds") + ref_audio = ref_audio[:, :30 * sr - source_audio.size(1)] + + + source_waves_16k = torchaudio.functional.resample(source_audio, sr, 16000) + ref_waves_16k = torchaudio.functional.resample(ref_audio, sr, 16000) + + S_alt = semantic_fn(source_waves_16k) + S_ori = semantic_fn(ref_waves_16k) + + mel = to_mel(source_audio.to(device).float()) + mel2 = to_mel(ref_audio.to(device).float()) + + target_lengths = torch.LongTensor([int(mel.size(2) * length_adjust)]).to(mel.device) + target2_lengths = torch.LongTensor([mel2.size(2)]).to(mel2.device) + + feat2 = torchaudio.compliance.kaldi.fbank( + ref_waves_16k, num_mel_bins=80, dither=0, sample_frequency=16000 + ) + feat2 = feat2 - feat2.mean(dim=0, keepdim=True) + style2 = campplus_model(feat2.unsqueeze(0)) + # Length regulation + cond = model.length_regulator( + S_alt, ylens=target_lengths, n_quantizers=3, f0=None + )[0] + prompt_condition = model.length_regulator( + S_ori, ylens=target2_lengths, n_quantizers=3, f0=None + )[0] + if remove_prompt: + cat_condition = cond + mel2 = torch.zeros([mel2.size(0), mel2.size(1), 0]).to(mel2.device) + else: + cat_condition = torch.cat([prompt_condition, cond], dim=1) + + vc_target = model.cfm.inference( + cat_condition, + torch.LongTensor([cat_condition.size(1)]).to(mel2.device), + mel2, + style2, + None, + diffusion_steps, + inference_cfg_rate=inference_cfg_rate, + ) + vc_target = vc_target[:, :, mel2.size(-1) :] + + # Convert to waveform + vc_wave = vocoder_fn(vc_target).squeeze(1) + + return ref_waves_16k, vc_wave + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--source", type=str, default="./examples/libritts-test-clean/" + ) + parser.add_argument("--target", type=str, default="./examples/reference/") + parser.add_argument("--output", type=str, default="./examples/eval/converted/") + parser.add_argument("--diffusion-steps", type=int, default=30) + parser.add_argument("--length-adjust", type=float, default=1.0) + parser.add_argument("--inference-cfg-rate", type=float, default=0.7) + parser.add_argument( + "--xvector-extractor", type=str, default="wavlm-large" + ) # wavlm or resemblyzer + parser.add_argument("--baseline", type=str, default="") # use "" for Seed-VC + parser.add_argument("--max-samples", type=int, default=20) + parser.add_argument("--remove-prompt", type=bool, default=False) + args = parser.parse_args() + main(args) \ No newline at end of file diff --git a/examples/reference/azuma_0.wav b/examples/reference/azuma_0.wav new file mode 100644 index 0000000000000000000000000000000000000000..f7ab6aa2d4f8d04b67379130d07a46e3b53e9144 Binary files /dev/null and b/examples/reference/azuma_0.wav differ diff --git a/examples/reference/dingzhen_0.wav b/examples/reference/dingzhen_0.wav new file mode 100644 index 0000000000000000000000000000000000000000..b260e1b4ed2d9ea746c711dd0f4983407dd4ee28 --- /dev/null +++ b/examples/reference/dingzhen_0.wav @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3db260824d11f56cdf2fccf2b84ad83c95a732ddfa2f8cb8a20b68ca06ea9ff8 +size 1088420 diff --git a/examples/reference/s1p1.wav b/examples/reference/s1p1.wav new file mode 100644 index 0000000000000000000000000000000000000000..a9b67356806dd9d11a6b5e06fab950844859ef41 Binary files /dev/null and b/examples/reference/s1p1.wav differ diff --git a/examples/reference/s1p2.wav b/examples/reference/s1p2.wav new file mode 100644 index 0000000000000000000000000000000000000000..5654fcd829d5e552f8342ed4e46387fafd1f462e Binary files /dev/null and b/examples/reference/s1p2.wav differ diff --git a/examples/reference/s2p1.wav b/examples/reference/s2p1.wav new file mode 100644 index 0000000000000000000000000000000000000000..a3ca1ffce3b46d9634c88aae308327d7f2ea8c85 Binary files /dev/null and b/examples/reference/s2p1.wav differ diff --git a/examples/reference/s2p2.wav b/examples/reference/s2p2.wav new file mode 100644 index 0000000000000000000000000000000000000000..9b86db6af7ba20aad742d7ba647b4a3f7d70a5bc Binary files /dev/null and b/examples/reference/s2p2.wav differ diff --git a/examples/reference/s3p1.wav b/examples/reference/s3p1.wav new file mode 100644 index 0000000000000000000000000000000000000000..8f27305452ac2650c9cdc1878ccbd8a68e8c9d9e Binary files /dev/null and b/examples/reference/s3p1.wav differ diff --git a/examples/reference/s3p2.wav b/examples/reference/s3p2.wav new file mode 100644 index 0000000000000000000000000000000000000000..0c5c30b861bdb58ece69e6bc17e9e713ddaa9e59 --- /dev/null +++ b/examples/reference/s3p2.wav @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:0d28df338203ad8b3c7485474fac41d9ee2891cf27bc0c0239e3249e6c0efadb +size 1140390 diff --git a/examples/reference/s4p1.wav b/examples/reference/s4p1.wav new file mode 100644 index 0000000000000000000000000000000000000000..bb24720d92330d716b16bb6e7f5a9d790c3a1ec2 Binary files /dev/null and b/examples/reference/s4p1.wav differ diff --git a/examples/reference/s4p2.wav b/examples/reference/s4p2.wav new file mode 100644 index 0000000000000000000000000000000000000000..1936e4b7b2d88738ecab7c8449d0c808eda2d0b6 Binary files /dev/null and b/examples/reference/s4p2.wav differ diff --git a/examples/reference/teio_0.wav b/examples/reference/teio_0.wav new file mode 100644 index 0000000000000000000000000000000000000000..1968dc484cac879175c1133651e363990b6aa98d Binary files /dev/null and b/examples/reference/teio_0.wav differ diff --git a/examples/reference/trump_0.wav b/examples/reference/trump_0.wav new file mode 100644 index 0000000000000000000000000000000000000000..c24f3495db37ef1884d4441723f3506355ac5ff6 --- /dev/null +++ b/examples/reference/trump_0.wav @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:716becc9daf00351dfe324398edea9e8378f9453408b27612d92b6721f80ddbc +size 1379484 diff --git a/examples/source/TECHNOPOLIS - 2085 [vocals]_[cut_14sec].wav b/examples/source/TECHNOPOLIS - 2085 [vocals]_[cut_14sec].wav new file mode 100644 index 0000000000000000000000000000000000000000..4dcdd668892a77d1ed8a392671c6b9d541b8c83f --- /dev/null +++ b/examples/source/TECHNOPOLIS - 2085 [vocals]_[cut_14sec].wav @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:87087ca5260ce96659b01a647edb30bb08527ed7d0c074fb5ae1e8338cc733e5 +size 2796016 diff --git a/examples/source/Wiz Khalifa,Charlie Puth - See You Again [vocals]_[cut_28sec].wav b/examples/source/Wiz Khalifa,Charlie Puth - See You Again [vocals]_[cut_28sec].wav new file mode 100644 index 0000000000000000000000000000000000000000..38780402102a7a97d9e2d5ff19d445d8dfdff6a5 --- /dev/null +++ b/examples/source/Wiz Khalifa,Charlie Puth - See You Again [vocals]_[cut_28sec].wav @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c7ff68178df50fc5f4f497099af8efe0b4508e7ff13665ee72d780159e6d1875 +size 5411452 diff --git a/examples/source/glados_0.wav b/examples/source/glados_0.wav new file mode 100644 index 0000000000000000000000000000000000000000..134f8f712990ecb798d0e69770e191febc64f5fa Binary files /dev/null and b/examples/source/glados_0.wav differ diff --git a/examples/source/jay_0.wav b/examples/source/jay_0.wav new file mode 100644 index 0000000000000000000000000000000000000000..0c029b001844933049dbffc6c937b51720aae449 --- /dev/null +++ b/examples/source/jay_0.wav @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:1d30f1500acacb597c3b27d7a5937dd088b8029b27e9db8bf5982085f26f4457 +size 1270124 diff --git a/examples/source/source_s1.wav b/examples/source/source_s1.wav new file mode 100644 index 0000000000000000000000000000000000000000..6acd3e01c05a328a0e944c516e86ddd9fe73b9e7 Binary files /dev/null and b/examples/source/source_s1.wav differ diff --git a/examples/source/source_s2.wav b/examples/source/source_s2.wav new file mode 100644 index 0000000000000000000000000000000000000000..e32f65fd5a906b96446a04de64b3258ea90afbd6 Binary files /dev/null and b/examples/source/source_s2.wav differ diff --git a/examples/source/source_s3.wav b/examples/source/source_s3.wav new file mode 100644 index 0000000000000000000000000000000000000000..e60db9c66f7fd43c471708754e5413241e4ca4df --- /dev/null +++ b/examples/source/source_s3.wav @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:0e784fc21f6eade4633db4050528110a9fc307686487fcec257ad3db35a65c8d +size 1168780 diff --git a/examples/source/source_s4.wav b/examples/source/source_s4.wav new file mode 100644 index 0000000000000000000000000000000000000000..1ae47aa2d5f8271cbdf60220f06f87a21a31544b --- /dev/null +++ b/examples/source/source_s4.wav @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ee7ccf45748bf8918c17e9bf80f79af423cefd78d9c8eebb717a395305272b9f +size 1068360 diff --git a/examples/source/yae_0.wav b/examples/source/yae_0.wav new file mode 100644 index 0000000000000000000000000000000000000000..24d45ec68a593753439280822825781c3f155d59 Binary files /dev/null and b/examples/source/yae_0.wav differ diff --git a/hf_utils.py b/hf_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..9f8c7f7d5f1b82efbd788c7327f76c0dc6a9355a --- /dev/null +++ b/hf_utils.py @@ -0,0 +1,12 @@ +import os +from huggingface_hub import hf_hub_download + + +def load_custom_model_from_hf(repo_id, model_filename="pytorch_model.bin", config_filename="config.yml"): + os.makedirs("./checkpoints", exist_ok=True) + model_path = hf_hub_download(repo_id=repo_id, filename=model_filename, cache_dir="./checkpoints") + if config_filename is None: + return model_path + config_path = hf_hub_download(repo_id=repo_id, filename=config_filename, cache_dir="./checkpoints") + + return model_path, config_path \ No newline at end of file diff --git a/inference.py b/inference.py new file mode 100644 index 0000000000000000000000000000000000000000..ed6d4a3eff5883554519f2231580f8974be731e8 --- /dev/null +++ b/inference.py @@ -0,0 +1,414 @@ +import os + +import numpy as np + +os.environ['HF_HUB_CACHE'] = './checkpoints/hf_cache' +import shutil +import warnings +import argparse +import torch +import yaml + +warnings.simplefilter('ignore') + +# load packages +import random + +from modules.commons import * +import time + +import torchaudio +import librosa +from modules.commons import str2bool + +from hf_utils import load_custom_model_from_hf + + +# Load model and configuration +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") +fp16 = False +def load_models(args): + global fp16 + fp16 = args.fp16 + if not args.f0_condition: + dit_checkpoint_path, dit_config_path = load_custom_model_from_hf("Plachta/Seed-VC", + "DiT_seed_v2_uvit_whisper_small_wavenet_bigvgan_pruned.pth", + "config_dit_mel_seed_uvit_whisper_small_wavenet.yml") + f0_fn = None + else: + if args.checkpoint_path is None: + dit_checkpoint_path, dit_config_path = load_custom_model_from_hf("Plachta/Seed-VC", + "DiT_seed_v2_uvit_whisper_base_f0_44k_bigvgan_pruned_ft_ema_v2.pth", + "config_dit_mel_seed_uvit_whisper_base_f0_44k.yml") + else: + dit_checkpoint_path = args.checkpoint_path + dit_config_path = args.config_path + # f0 extractor + from modules.rmvpe import RMVPE + + model_path = load_custom_model_from_hf("lj1995/VoiceConversionWebUI", "rmvpe.pt", None) + f0_extractor = RMVPE(model_path, is_half=False, device=device) + f0_fn = f0_extractor.infer_from_audio + + config = yaml.safe_load(open(dit_config_path, "r")) + model_params = recursive_munch(config["model_params"]) + model_params.dit_type = 'DiT' + model = build_model(model_params, stage="DiT") + hop_length = config["preprocess_params"]["spect_params"]["hop_length"] + sr = config["preprocess_params"]["sr"] + + # Load checkpoints + model, _, _, _ = load_checkpoint( + model, + None, + dit_checkpoint_path, + load_only_params=True, + ignore_modules=[], + is_distributed=False, + ) + for key in model: + model[key].eval() + model[key].to(device) + model.cfm.estimator.setup_caches(max_batch_size=1, max_seq_length=8192) + + # Load additional modules + from modules.campplus.DTDNN import CAMPPlus + + campplus_ckpt_path = load_custom_model_from_hf( + "funasr/campplus", "campplus_cn_common.bin", config_filename=None + ) + campplus_model = CAMPPlus(feat_dim=80, embedding_size=192) + campplus_model.load_state_dict(torch.load(campplus_ckpt_path, map_location="cpu")) + campplus_model.eval() + campplus_model.to(device) + + vocoder_type = model_params.vocoder.type + + if vocoder_type == 'bigvgan': + from modules.bigvgan import bigvgan + bigvgan_name = model_params.vocoder.name + bigvgan_model = bigvgan.BigVGAN.from_pretrained(bigvgan_name, use_cuda_kernel=False) + # remove weight norm in the model and set to eval mode + bigvgan_model.remove_weight_norm() + bigvgan_model = bigvgan_model.eval().to(device) + vocoder_fn = bigvgan_model + elif vocoder_type == 'hifigan': + from modules.hifigan.generator import HiFTGenerator + from modules.hifigan.f0_predictor import ConvRNNF0Predictor + hift_config = yaml.safe_load(open('configs/hifigan.yml', 'r')) + hift_gen = HiFTGenerator(**hift_config['hift'], f0_predictor=ConvRNNF0Predictor(**hift_config['f0_predictor'])) + hift_path = load_custom_model_from_hf("FunAudioLLM/CosyVoice-300M", 'hift.pt', None) + hift_gen.load_state_dict(torch.load(hift_path, map_location='cpu')) + hift_gen.eval() + hift_gen.to(device) + vocoder_fn = hift_gen + elif vocoder_type == "vocos": + vocos_config = yaml.safe_load(open(model_params.vocoder.vocos.config, 'r')) + vocos_path = model_params.vocoder.vocos.path + vocos_model_params = recursive_munch(vocos_config['model_params']) + vocos = build_model(vocos_model_params, stage='mel_vocos') + vocos_checkpoint_path = vocos_path + vocos, _, _, _ = load_checkpoint(vocos, None, vocos_checkpoint_path, + load_only_params=True, ignore_modules=[], is_distributed=False) + _ = [vocos[key].eval().to(device) for key in vocos] + _ = [vocos[key].to(device) for key in vocos] + total_params = sum(sum(p.numel() for p in vocos[key].parameters() if p.requires_grad) for key in vocos.keys()) + print(f"Vocoder model total parameters: {total_params / 1_000_000:.2f}M") + vocoder_fn = vocos.decoder + else: + raise ValueError(f"Unknown vocoder type: {vocoder_type}") + + speech_tokenizer_type = model_params.speech_tokenizer.type + if speech_tokenizer_type == 'whisper': + # whisper + from transformers import AutoFeatureExtractor, WhisperModel + whisper_name = model_params.speech_tokenizer.name + whisper_model = WhisperModel.from_pretrained(whisper_name, torch_dtype=torch.float16).to(device) + del whisper_model.decoder + whisper_feature_extractor = AutoFeatureExtractor.from_pretrained(whisper_name) + + def semantic_fn(waves_16k): + ori_inputs = whisper_feature_extractor([waves_16k.squeeze(0).cpu().numpy()], + return_tensors="pt", + return_attention_mask=True) + ori_input_features = whisper_model._mask_input_features( + ori_inputs.input_features, attention_mask=ori_inputs.attention_mask).to(device) + with torch.no_grad(): + ori_outputs = whisper_model.encoder( + ori_input_features.to(whisper_model.encoder.dtype), + head_mask=None, + output_attentions=False, + output_hidden_states=False, + return_dict=True, + ) + S_ori = ori_outputs.last_hidden_state.to(torch.float32) + S_ori = S_ori[:, :waves_16k.size(-1) // 320 + 1] + return S_ori + elif speech_tokenizer_type == 'cnhubert': + from transformers import ( + Wav2Vec2FeatureExtractor, + HubertModel, + ) + hubert_model_name = config['model_params']['speech_tokenizer']['name'] + hubert_feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(hubert_model_name) + hubert_model = HubertModel.from_pretrained(hubert_model_name) + hubert_model = hubert_model.to(device) + hubert_model = hubert_model.eval() + hubert_model = hubert_model.half() + + def semantic_fn(waves_16k): + ori_waves_16k_input_list = [ + waves_16k[bib].cpu().numpy() + for bib in range(len(waves_16k)) + ] + ori_inputs = hubert_feature_extractor(ori_waves_16k_input_list, + return_tensors="pt", + return_attention_mask=True, + padding=True, + sampling_rate=16000).to(device) + with torch.no_grad(): + ori_outputs = hubert_model( + ori_inputs.input_values.half(), + ) + S_ori = ori_outputs.last_hidden_state.float() + return S_ori + elif speech_tokenizer_type == 'xlsr': + from transformers import ( + Wav2Vec2FeatureExtractor, + Wav2Vec2Model, + ) + model_name = config['model_params']['speech_tokenizer']['name'] + output_layer = config['model_params']['speech_tokenizer']['output_layer'] + wav2vec_feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(model_name) + wav2vec_model = Wav2Vec2Model.from_pretrained(model_name) + wav2vec_model.encoder.layers = wav2vec_model.encoder.layers[:output_layer] + wav2vec_model = wav2vec_model.to(device) + wav2vec_model = wav2vec_model.eval() + wav2vec_model = wav2vec_model.half() + + def semantic_fn(waves_16k): + ori_waves_16k_input_list = [ + waves_16k[bib].cpu().numpy() + for bib in range(len(waves_16k)) + ] + ori_inputs = wav2vec_feature_extractor(ori_waves_16k_input_list, + return_tensors="pt", + return_attention_mask=True, + padding=True, + sampling_rate=16000).to(device) + with torch.no_grad(): + ori_outputs = wav2vec_model( + ori_inputs.input_values.half(), + ) + S_ori = ori_outputs.last_hidden_state.float() + return S_ori + else: + raise ValueError(f"Unknown speech tokenizer type: {speech_tokenizer_type}") + # Generate mel spectrograms + mel_fn_args = { + "n_fft": config['preprocess_params']['spect_params']['n_fft'], + "win_size": config['preprocess_params']['spect_params']['win_length'], + "hop_size": config['preprocess_params']['spect_params']['hop_length'], + "num_mels": config['preprocess_params']['spect_params']['n_mels'], + "sampling_rate": sr, + "fmin": config['preprocess_params']['spect_params'].get('fmin', 0), + "fmax": None if config['preprocess_params']['spect_params'].get('fmax', "None") == "None" else 8000, + "center": False + } + from modules.audio import mel_spectrogram + + to_mel = lambda x: mel_spectrogram(x, **mel_fn_args) + + return ( + model, + semantic_fn, + f0_fn, + vocoder_fn, + campplus_model, + to_mel, + mel_fn_args, + ) + +def adjust_f0_semitones(f0_sequence, n_semitones): + factor = 2 ** (n_semitones / 12) + return f0_sequence * factor + +def crossfade(chunk1, chunk2, overlap): + fade_out = np.cos(np.linspace(0, np.pi / 2, overlap)) ** 2 + fade_in = np.cos(np.linspace(np.pi / 2, 0, overlap)) ** 2 + if len(chunk2) < overlap: + chunk2[:overlap] = chunk2[:overlap] * fade_in[:len(chunk2)] + (chunk1[-overlap:] * fade_out)[:len(chunk2)] + else: + chunk2[:overlap] = chunk2[:overlap] * fade_in + chunk1[-overlap:] * fade_out + return chunk2 + +@torch.no_grad() +def main(args): + model, semantic_fn, f0_fn, vocoder_fn, campplus_model, mel_fn, mel_fn_args = load_models(args) + sr = mel_fn_args['sampling_rate'] + f0_condition = args.f0_condition + auto_f0_adjust = args.auto_f0_adjust + pitch_shift = args.semi_tone_shift + + source = args.source + target_name = args.target + diffusion_steps = args.diffusion_steps + length_adjust = args.length_adjust + inference_cfg_rate = args.inference_cfg_rate + source_audio = librosa.load(source, sr=sr)[0] + ref_audio = librosa.load(target_name, sr=sr)[0] + + sr = 22050 if not f0_condition else 44100 + hop_length = 256 if not f0_condition else 512 + max_context_window = sr // hop_length * 30 + overlap_frame_len = 16 + overlap_wave_len = overlap_frame_len * hop_length + + # Process audio + source_audio = torch.tensor(source_audio).unsqueeze(0).float().to(device) + ref_audio = torch.tensor(ref_audio[:sr * 25]).unsqueeze(0).float().to(device) + + time_vc_start = time.time() + # Resample + converted_waves_16k = torchaudio.functional.resample(source_audio, sr, 16000) + # if source audio less than 30 seconds, whisper can handle in one forward + if converted_waves_16k.size(-1) <= 16000 * 30: + S_alt = semantic_fn(converted_waves_16k) + else: + overlapping_time = 5 # 5 seconds + S_alt_list = [] + buffer = None + traversed_time = 0 + while traversed_time < converted_waves_16k.size(-1): + if buffer is None: # first chunk + chunk = converted_waves_16k[:, traversed_time:traversed_time + 16000 * 30] + else: + chunk = torch.cat( + [buffer, converted_waves_16k[:, traversed_time:traversed_time + 16000 * (30 - overlapping_time)]], + dim=-1) + S_alt = semantic_fn(chunk) + if traversed_time == 0: + S_alt_list.append(S_alt) + else: + S_alt_list.append(S_alt[:, 50 * overlapping_time:]) + buffer = chunk[:, -16000 * overlapping_time:] + traversed_time += 30 * 16000 if traversed_time == 0 else chunk.size(-1) - 16000 * overlapping_time + S_alt = torch.cat(S_alt_list, dim=1) + + ori_waves_16k = torchaudio.functional.resample(ref_audio, sr, 16000) + S_ori = semantic_fn(ori_waves_16k) + + mel = mel_fn(source_audio.to(device).float()) + mel2 = mel_fn(ref_audio.to(device).float()) + + target_lengths = torch.LongTensor([int(mel.size(2) * length_adjust)]).to(mel.device) + target2_lengths = torch.LongTensor([mel2.size(2)]).to(mel2.device) + + feat2 = torchaudio.compliance.kaldi.fbank(ori_waves_16k, + num_mel_bins=80, + dither=0, + sample_frequency=16000) + feat2 = feat2 - feat2.mean(dim=0, keepdim=True) + style2 = campplus_model(feat2.unsqueeze(0)) + + if f0_condition: + F0_ori = f0_fn(ori_waves_16k[0], thred=0.03) + F0_alt = f0_fn(converted_waves_16k[0], thred=0.03) + + F0_ori = torch.from_numpy(F0_ori).to(device)[None] + F0_alt = torch.from_numpy(F0_alt).to(device)[None] + + voiced_F0_ori = F0_ori[F0_ori > 1] + voiced_F0_alt = F0_alt[F0_alt > 1] + + log_f0_alt = torch.log(F0_alt + 1e-5) + voiced_log_f0_ori = torch.log(voiced_F0_ori + 1e-5) + voiced_log_f0_alt = torch.log(voiced_F0_alt + 1e-5) + median_log_f0_ori = torch.median(voiced_log_f0_ori) + median_log_f0_alt = torch.median(voiced_log_f0_alt) + + # shift alt log f0 level to ori log f0 level + shifted_log_f0_alt = log_f0_alt.clone() + if auto_f0_adjust: + shifted_log_f0_alt[F0_alt > 1] = log_f0_alt[F0_alt > 1] - median_log_f0_alt + median_log_f0_ori + shifted_f0_alt = torch.exp(shifted_log_f0_alt) + if pitch_shift != 0: + shifted_f0_alt[F0_alt > 1] = adjust_f0_semitones(shifted_f0_alt[F0_alt > 1], pitch_shift) + else: + F0_ori = None + F0_alt = None + shifted_f0_alt = None + + # Length regulation + cond, _, codes, commitment_loss, codebook_loss = model.length_regulator(S_alt, ylens=target_lengths, + n_quantizers=3, + f0=shifted_f0_alt) + prompt_condition, _, codes, commitment_loss, codebook_loss = model.length_regulator(S_ori, + ylens=target2_lengths, + n_quantizers=3, + f0=F0_ori) + + max_source_window = max_context_window - mel2.size(2) + # split source condition (cond) into chunks + processed_frames = 0 + generated_wave_chunks = [] + # generate chunk by chunk and stream the output + while processed_frames < cond.size(1): + chunk_cond = cond[:, processed_frames:processed_frames + max_source_window] + is_last_chunk = processed_frames + max_source_window >= cond.size(1) + cat_condition = torch.cat([prompt_condition, chunk_cond], dim=1) + with torch.autocast(device_type=device.type, dtype=torch.float16 if fp16 else torch.float32): + # Voice Conversion + vc_target = model.cfm.inference(cat_condition, + torch.LongTensor([cat_condition.size(1)]).to(mel2.device), + mel2, style2, None, diffusion_steps, + inference_cfg_rate=inference_cfg_rate) + vc_target = vc_target[:, :, mel2.size(-1):] + vc_wave = vocoder_fn(vc_target.float()).squeeze() + vc_wave = vc_wave[None, :] + if processed_frames == 0: + if is_last_chunk: + output_wave = vc_wave[0].cpu().numpy() + generated_wave_chunks.append(output_wave) + break + output_wave = vc_wave[0, :-overlap_wave_len].cpu().numpy() + generated_wave_chunks.append(output_wave) + previous_chunk = vc_wave[0, -overlap_wave_len:] + processed_frames += vc_target.size(2) - overlap_frame_len + elif is_last_chunk: + output_wave = crossfade(previous_chunk.cpu().numpy(), vc_wave[0].cpu().numpy(), overlap_wave_len) + generated_wave_chunks.append(output_wave) + processed_frames += vc_target.size(2) - overlap_frame_len + break + else: + output_wave = crossfade(previous_chunk.cpu().numpy(), vc_wave[0, :-overlap_wave_len].cpu().numpy(), + overlap_wave_len) + generated_wave_chunks.append(output_wave) + previous_chunk = vc_wave[0, -overlap_wave_len:] + processed_frames += vc_target.size(2) - overlap_frame_len + vc_wave = torch.tensor(np.concatenate(generated_wave_chunks))[None, :].float() + time_vc_end = time.time() + print(f"RTF: {(time_vc_end - time_vc_start) / vc_wave.size(-1) * sr}") + + source_name = os.path.basename(source).split(".")[0] + target_name = os.path.basename(target_name).split(".")[0] + os.makedirs(args.output, exist_ok=True) + torchaudio.save(os.path.join(args.output, f"vc_{source_name}_{target_name}_{length_adjust}_{diffusion_steps}_{inference_cfg_rate}.wav"), vc_wave.cpu(), sr) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--source", type=str, default="./examples/source/source_s1.wav") + parser.add_argument("--target", type=str, default="./examples/reference/s1p1.wav") + parser.add_argument("--output", type=str, default="./reconstructed") + parser.add_argument("--diffusion-steps", type=int, default=30) + parser.add_argument("--length-adjust", type=float, default=1.0) + parser.add_argument("--inference-cfg-rate", type=float, default=0.7) + parser.add_argument("--f0-condition", type=str2bool, default=False) + parser.add_argument("--auto-f0-adjust", type=str2bool, default=False) + parser.add_argument("--semi-tone-shift", type=int, default=0) + parser.add_argument("--checkpoint-path", type=str, help="Path to the checkpoint file", default=None) + parser.add_argument("--config-path", type=str, help="Path to the config file", default=None) + parser.add_argument("--fp16", type=str2bool, default=True) + args = parser.parse_args() + main(args) diff --git a/modules/alias_free_torch/__init__.py b/modules/alias_free_torch/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..080beecd700d74c9c915b7c41c5ab13c79103902 --- /dev/null +++ b/modules/alias_free_torch/__init__.py @@ -0,0 +1,5 @@ +# Adapted from https://github.com/junjun3518/alias-free-torch under the Apache License 2.0 + +from .filter import * +from .resample import * +from .act import * diff --git a/modules/alias_free_torch/act.py b/modules/alias_free_torch/act.py new file mode 100644 index 0000000000000000000000000000000000000000..adc7eb7469c6fcb25e200abb78425f977d655cd8 --- /dev/null +++ b/modules/alias_free_torch/act.py @@ -0,0 +1,29 @@ +# Adapted from https://github.com/junjun3518/alias-free-torch under the Apache License 2.0 + +import torch.nn as nn +from .resample import UpSample1d, DownSample1d + + +class Activation1d(nn.Module): + def __init__( + self, + activation, + up_ratio: int = 2, + down_ratio: int = 2, + up_kernel_size: int = 12, + down_kernel_size: int = 12, + ): + super().__init__() + self.up_ratio = up_ratio + self.down_ratio = down_ratio + self.act = activation + self.upsample = UpSample1d(up_ratio, up_kernel_size) + self.downsample = DownSample1d(down_ratio, down_kernel_size) + + # x: [B,C,T] + def forward(self, x): + x = self.upsample(x) + x = self.act(x) + x = self.downsample(x) + + return x diff --git a/modules/alias_free_torch/filter.py b/modules/alias_free_torch/filter.py new file mode 100644 index 0000000000000000000000000000000000000000..8e18bda585d0006c4f34f5ae04bdea854c4e37b7 --- /dev/null +++ b/modules/alias_free_torch/filter.py @@ -0,0 +1,96 @@ +# Adapted from https://github.com/junjun3518/alias-free-torch under the Apache License 2.0 + +import torch +import torch.nn as nn +import torch.nn.functional as F +import math + +if "sinc" in dir(torch): + sinc = torch.sinc +else: + # This code is adopted from adefossez's julius.core.sinc under the MIT License + # https://adefossez.github.io/julius/julius/core.html + def sinc(x: torch.Tensor): + """ + Implementation of sinc, i.e. sin(pi * x) / (pi * x) + __Warning__: Different to julius.sinc, the input is multiplied by `pi`! + """ + return torch.where( + x == 0, + torch.tensor(1.0, device=x.device, dtype=x.dtype), + torch.sin(math.pi * x) / math.pi / x, + ) + + +# This code is adopted from adefossez's julius.lowpass.LowPassFilters under the MIT License +# https://adefossez.github.io/julius/julius/lowpass.html +def kaiser_sinc_filter1d( + cutoff, half_width, kernel_size +): # return filter [1,1,kernel_size] + even = kernel_size % 2 == 0 + half_size = kernel_size // 2 + + # For kaiser window + delta_f = 4 * half_width + A = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95 + if A > 50.0: + beta = 0.1102 * (A - 8.7) + elif A >= 21.0: + beta = 0.5842 * (A - 21) ** 0.4 + 0.07886 * (A - 21.0) + else: + beta = 0.0 + window = torch.kaiser_window(kernel_size, beta=beta, periodic=False) + + # ratio = 0.5/cutoff -> 2 * cutoff = 1 / ratio + if even: + time = torch.arange(-half_size, half_size) + 0.5 + else: + time = torch.arange(kernel_size) - half_size + if cutoff == 0: + filter_ = torch.zeros_like(time) + else: + filter_ = 2 * cutoff * window * sinc(2 * cutoff * time) + # Normalize filter to have sum = 1, otherwise we will have a small leakage + # of the constant component in the input signal. + filter_ /= filter_.sum() + filter = filter_.view(1, 1, kernel_size) + + return filter + + +class LowPassFilter1d(nn.Module): + def __init__( + self, + cutoff=0.5, + half_width=0.6, + stride: int = 1, + padding: bool = True, + padding_mode: str = "replicate", + kernel_size: int = 12, + ): + # kernel_size should be even number for stylegan3 setup, + # in this implementation, odd number is also possible. + super().__init__() + if cutoff < -0.0: + raise ValueError("Minimum cutoff must be larger than zero.") + if cutoff > 0.5: + raise ValueError("A cutoff above 0.5 does not make sense.") + self.kernel_size = kernel_size + self.even = kernel_size % 2 == 0 + self.pad_left = kernel_size // 2 - int(self.even) + self.pad_right = kernel_size // 2 + self.stride = stride + self.padding = padding + self.padding_mode = padding_mode + filter = kaiser_sinc_filter1d(cutoff, half_width, kernel_size) + self.register_buffer("filter", filter) + + # input [B, C, T] + def forward(self, x): + _, C, _ = x.shape + + if self.padding: + x = F.pad(x, (self.pad_left, self.pad_right), mode=self.padding_mode) + out = F.conv1d(x, self.filter.expand(C, -1, -1), stride=self.stride, groups=C) + + return out diff --git a/modules/alias_free_torch/resample.py b/modules/alias_free_torch/resample.py new file mode 100644 index 0000000000000000000000000000000000000000..dd35bc32b4b30a293ded20f4a66c1578ffbfc54e --- /dev/null +++ b/modules/alias_free_torch/resample.py @@ -0,0 +1,57 @@ +# Adapted from https://github.com/junjun3518/alias-free-torch under the Apache License 2.0 + +import torch.nn as nn +from torch.nn import functional as F +from .filter import LowPassFilter1d +from .filter import kaiser_sinc_filter1d + + +class UpSample1d(nn.Module): + def __init__(self, ratio=2, kernel_size=None): + super().__init__() + self.ratio = ratio + self.kernel_size = ( + int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size + ) + self.stride = ratio + self.pad = self.kernel_size // ratio - 1 + self.pad_left = self.pad * self.stride + (self.kernel_size - self.stride) // 2 + self.pad_right = ( + self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2 + ) + filter = kaiser_sinc_filter1d( + cutoff=0.5 / ratio, half_width=0.6 / ratio, kernel_size=self.kernel_size + ) + self.register_buffer("filter", filter) + + # x: [B, C, T] + def forward(self, x): + _, C, _ = x.shape + + x = F.pad(x, (self.pad, self.pad), mode="replicate") + x = self.ratio * F.conv_transpose1d( + x, self.filter.expand(C, -1, -1), stride=self.stride, groups=C + ) + x = x[..., self.pad_left : -self.pad_right] + + return x + + +class DownSample1d(nn.Module): + def __init__(self, ratio=2, kernel_size=None): + super().__init__() + self.ratio = ratio + self.kernel_size = ( + int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size + ) + self.lowpass = LowPassFilter1d( + cutoff=0.5 / ratio, + half_width=0.6 / ratio, + stride=ratio, + kernel_size=self.kernel_size, + ) + + def forward(self, x): + xx = self.lowpass(x) + + return xx diff --git a/modules/audio.py b/modules/audio.py new file mode 100644 index 0000000000000000000000000000000000000000..ae677ffb1c124b557b3dbe0343ae415f5281cddb --- /dev/null +++ b/modules/audio.py @@ -0,0 +1,82 @@ +import numpy as np +import torch +import torch.utils.data +from librosa.filters import mel as librosa_mel_fn +from scipy.io.wavfile import read + +MAX_WAV_VALUE = 32768.0 + + +def load_wav(full_path): + sampling_rate, data = read(full_path) + return data, sampling_rate + + +def dynamic_range_compression(x, C=1, clip_val=1e-5): + return np.log(np.clip(x, a_min=clip_val, a_max=None) * C) + + +def dynamic_range_decompression(x, C=1): + return np.exp(x) / C + + +def dynamic_range_compression_torch(x, C=1, clip_val=1e-5): + return torch.log(torch.clamp(x, min=clip_val) * C) + + +def dynamic_range_decompression_torch(x, C=1): + return torch.exp(x) / C + + +def spectral_normalize_torch(magnitudes): + output = dynamic_range_compression_torch(magnitudes) + return output + + +def spectral_de_normalize_torch(magnitudes): + output = dynamic_range_decompression_torch(magnitudes) + return output + + +mel_basis = {} +hann_window = {} + + +def mel_spectrogram(y, n_fft, num_mels, sampling_rate, hop_size, win_size, fmin, fmax, center=False): + if torch.min(y) < -1.0: + print("min value is ", torch.min(y)) + if torch.max(y) > 1.0: + print("max value is ", torch.max(y)) + + global mel_basis, hann_window # pylint: disable=global-statement + if f"{str(sampling_rate)}_{str(fmax)}_{str(y.device)}" not in mel_basis: + mel = librosa_mel_fn(sr=sampling_rate, n_fft=n_fft, n_mels=num_mels, fmin=fmin, fmax=fmax) + mel_basis[str(sampling_rate) + "_" + str(fmax) + "_" + str(y.device)] = torch.from_numpy(mel).float().to(y.device) + hann_window[str(sampling_rate) + "_" + str(y.device)] = torch.hann_window(win_size).to(y.device) + + y = torch.nn.functional.pad( + y.unsqueeze(1), (int((n_fft - hop_size) / 2), int((n_fft - hop_size) / 2)), mode="reflect" + ) + y = y.squeeze(1) + + spec = torch.view_as_real( + torch.stft( + y, + n_fft, + hop_length=hop_size, + win_length=win_size, + window=hann_window[str(sampling_rate) + "_" + str(y.device)], + center=center, + pad_mode="reflect", + normalized=False, + onesided=True, + return_complex=True, + ) + ) + + spec = torch.sqrt(spec.pow(2).sum(-1) + (1e-9)) + + spec = torch.matmul(mel_basis[str(sampling_rate) + "_" + str(fmax) + "_" + str(y.device)], spec) + spec = spectral_normalize_torch(spec) + + return spec diff --git a/modules/bigvgan/activations.py b/modules/bigvgan/activations.py new file mode 100644 index 0000000000000000000000000000000000000000..9d5420c1994f064635a3bb283d83f6cec762d542 --- /dev/null +++ b/modules/bigvgan/activations.py @@ -0,0 +1,120 @@ +# Implementation adapted from https://github.com/EdwardDixon/snake under the MIT license. +# LICENSE is in incl_licenses directory. + +import torch +from torch import nn, sin, pow +from torch.nn import Parameter + + +class Snake(nn.Module): + ''' + Implementation of a sine-based periodic activation function + Shape: + - Input: (B, C, T) + - Output: (B, C, T), same shape as the input + Parameters: + - alpha - trainable parameter + References: + - This activation function is from this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda: + https://arxiv.org/abs/2006.08195 + Examples: + >>> a1 = snake(256) + >>> x = torch.randn(256) + >>> x = a1(x) + ''' + def __init__(self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False): + ''' + Initialization. + INPUT: + - in_features: shape of the input + - alpha: trainable parameter + alpha is initialized to 1 by default, higher values = higher-frequency. + alpha will be trained along with the rest of your model. + ''' + super(Snake, self).__init__() + self.in_features = in_features + + # initialize alpha + self.alpha_logscale = alpha_logscale + if self.alpha_logscale: # log scale alphas initialized to zeros + self.alpha = Parameter(torch.zeros(in_features) * alpha) + else: # linear scale alphas initialized to ones + self.alpha = Parameter(torch.ones(in_features) * alpha) + + self.alpha.requires_grad = alpha_trainable + + self.no_div_by_zero = 0.000000001 + + def forward(self, x): + ''' + Forward pass of the function. + Applies the function to the input elementwise. + Snake ∶= x + 1/a * sin^2 (xa) + ''' + alpha = self.alpha.unsqueeze(0).unsqueeze(-1) # line up with x to [B, C, T] + if self.alpha_logscale: + alpha = torch.exp(alpha) + x = x + (1.0 / (alpha + self.no_div_by_zero)) * pow(sin(x * alpha), 2) + + return x + + +class SnakeBeta(nn.Module): + ''' + A modified Snake function which uses separate parameters for the magnitude of the periodic components + Shape: + - Input: (B, C, T) + - Output: (B, C, T), same shape as the input + Parameters: + - alpha - trainable parameter that controls frequency + - beta - trainable parameter that controls magnitude + References: + - This activation function is a modified version based on this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda: + https://arxiv.org/abs/2006.08195 + Examples: + >>> a1 = snakebeta(256) + >>> x = torch.randn(256) + >>> x = a1(x) + ''' + def __init__(self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False): + ''' + Initialization. + INPUT: + - in_features: shape of the input + - alpha - trainable parameter that controls frequency + - beta - trainable parameter that controls magnitude + alpha is initialized to 1 by default, higher values = higher-frequency. + beta is initialized to 1 by default, higher values = higher-magnitude. + alpha will be trained along with the rest of your model. + ''' + super(SnakeBeta, self).__init__() + self.in_features = in_features + + # initialize alpha + self.alpha_logscale = alpha_logscale + if self.alpha_logscale: # log scale alphas initialized to zeros + self.alpha = Parameter(torch.zeros(in_features) * alpha) + self.beta = Parameter(torch.zeros(in_features) * alpha) + else: # linear scale alphas initialized to ones + self.alpha = Parameter(torch.ones(in_features) * alpha) + self.beta = Parameter(torch.ones(in_features) * alpha) + + self.alpha.requires_grad = alpha_trainable + self.beta.requires_grad = alpha_trainable + + self.no_div_by_zero = 0.000000001 + + def forward(self, x): + ''' + Forward pass of the function. + Applies the function to the input elementwise. + SnakeBeta ∶= x + 1/b * sin^2 (xa) + ''' + alpha = self.alpha.unsqueeze(0).unsqueeze(-1) # line up with x to [B, C, T] + beta = self.beta.unsqueeze(0).unsqueeze(-1) + if self.alpha_logscale: + alpha = torch.exp(alpha) + beta = torch.exp(beta) + x = x + (1.0 / (beta + self.no_div_by_zero)) * pow(sin(x * alpha), 2) + + return x \ No newline at end of file diff --git a/modules/bigvgan/alias_free_activation/cuda/__init__.py b/modules/bigvgan/alias_free_activation/cuda/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/modules/bigvgan/alias_free_activation/cuda/activation1d.py b/modules/bigvgan/alias_free_activation/cuda/activation1d.py new file mode 100644 index 0000000000000000000000000000000000000000..76797ef1424b99a160803d53cb6c24fe20599bd2 --- /dev/null +++ b/modules/bigvgan/alias_free_activation/cuda/activation1d.py @@ -0,0 +1,77 @@ +# Copyright (c) 2024 NVIDIA CORPORATION. +# Licensed under the MIT license. + +import torch +import torch.nn as nn +from ..torch.resample import UpSample1d, DownSample1d + +# load fused CUDA kernel: this enables importing anti_alias_activation_cuda +from ..cuda import load + +anti_alias_activation_cuda = load.load() + + +class FusedAntiAliasActivation(torch.autograd.Function): + """ + Assumes filter size 12, replication padding on upsampling/downsampling, and logscale alpha/beta parameters as inputs. + The hyperparameters are hard-coded in the kernel to maximize speed. + NOTE: The fused kenrel is incorrect for Activation1d with different hyperparameters. + """ + + @staticmethod + def forward(ctx, inputs, up_ftr, down_ftr, alpha, beta): + activation_results = anti_alias_activation_cuda.forward( + inputs, up_ftr, down_ftr, alpha, beta + ) + + return activation_results + + @staticmethod + def backward(ctx, output_grads): + raise NotImplementedError + return output_grads, None, None + + +class Activation1d(nn.Module): + def __init__( + self, + activation, + up_ratio: int = 2, + down_ratio: int = 2, + up_kernel_size: int = 12, + down_kernel_size: int = 12, + fused: bool = True, + ): + super().__init__() + self.up_ratio = up_ratio + self.down_ratio = down_ratio + self.act = activation + self.upsample = UpSample1d(up_ratio, up_kernel_size) + self.downsample = DownSample1d(down_ratio, down_kernel_size) + + self.fused = fused # Whether to use fused CUDA kernel or not + + def forward(self, x): + if not self.fused: + x = self.upsample(x) + x = self.act(x) + x = self.downsample(x) + return x + else: + if self.act.__class__.__name__ == "Snake": + beta = self.act.alpha.data # Snake uses same params for alpha and beta + else: + beta = ( + self.act.beta.data + ) # Snakebeta uses different params for alpha and beta + alpha = self.act.alpha.data + if ( + not self.act.alpha_logscale + ): # Exp baked into cuda kernel, cancel it out with a log + alpha = torch.log(alpha) + beta = torch.log(beta) + + x = FusedAntiAliasActivation.apply( + x, self.upsample.filter, self.downsample.lowpass.filter, alpha, beta + ) + return x diff --git a/modules/bigvgan/alias_free_activation/cuda/anti_alias_activation.cpp b/modules/bigvgan/alias_free_activation/cuda/anti_alias_activation.cpp new file mode 100644 index 0000000000000000000000000000000000000000..94fd90da386e66ce12a64ef243e4d125926dfd2a --- /dev/null +++ b/modules/bigvgan/alias_free_activation/cuda/anti_alias_activation.cpp @@ -0,0 +1,23 @@ +/* coding=utf-8 + * Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + + #include
+
+| | | |
+| :-----------: | :-----------: | :-----------: |
+| **OpenSource Repo** | **Project Page** | **Join the Community** |
+| 
| [OpenVoice](https://research.myshell.ai/open-voice) | [](https://discord.gg/myshell) |
+
+
+"""
+
+markdown_table_v2 = """
+
+
+| | | | |
+| :-----------: | :-----------: | :-----------: | :-----------: |
+| **OpenSource Repo** | | **Project Page** | [OpenVoice](https://research.myshell.ai/open-voice) |
+
+| | |
+| :-----------: | :-----------: |
+**Join the Community** | [](https://discord.gg/myshell) |
+
+
+"""
+content = """
+
+ If the generated voice does not sound like the reference voice, please refer to this QnA. For multi-lingual & cross-lingual examples, please refer to this jupyter notebook.
+ This online demo mainly supports English. The default style also supports Chinese. But OpenVoice can adapt to any other language as long as a base speaker is provided.
+
+"""
+wrapped_markdown_content = f"{content}
"
+
+
+examples = [
+ [
+ "今天天气真好，我们一起出去吃饭吧。",
+ 'default',
+ "resources/demo_speaker1.mp3",
+ True,
+ ],[
+ "This audio is generated by open voice with a half-performance model.",
+ 'whispering',
+ "resources/demo_speaker2.mp3",
+ True,
+ ],
+ [
+ "He hoped there would be stew for dinner, turnips and carrots and bruised potatoes and fat mutton pieces to be ladled out in thick, peppered, flour-fattened sauce.",
+ 'sad',
+ "resources/demo_speaker0.mp3",
+ True,
+ ],
+]
+
+with gr.Blocks(analytics_enabled=False) as demo:
+
+ with gr.Row():
+ with gr.Column():
+ with gr.Row():
+ gr.Markdown(
+ """
+ ## 
+ """
+ )
+ with gr.Row():
+ gr.Markdown(markdown_table_v2)
+ with gr.Row():
+ gr.Markdown(description)
+ with gr.Column():
+ gr.Video('https://github.com/myshell-ai/OpenVoice/assets/40556743/3cba936f-82bf-476c-9e52-09f0f417bb2f', autoplay=True)
+
+ with gr.Row():
+ gr.HTML(wrapped_markdown_content)
+
+ with gr.Row():
+ with gr.Column():
+ input_text_gr = gr.Textbox(
+ label="Text Prompt",
+ info="One or two sentences at a time is better. Up to 200 text characters.",
+ value="He hoped there would be stew for dinner, turnips and carrots and bruised potatoes and fat mutton pieces to be ladled out in thick, peppered, flour-fattened sauce.",
+ )
+ style_gr = gr.Dropdown(
+ label="Style",
+ info="Select a style of output audio for the synthesised speech. (Chinese only support 'default' now)",
+ choices=['default', 'whispering', 'cheerful', 'terrified', 'angry', 'sad', 'friendly'],
+ max_choices=1,
+ value="default",
+ )
+ ref_gr = gr.Audio(
+ label="Reference Audio",
+ info="Click on the ✎ button to upload your own target speaker audio",
+ type="filepath",
+ value="resources/demo_speaker2.mp3",
+ )
+ tos_gr = gr.Checkbox(
+ label="Agree",
+ value=False,
+ info="I agree to the terms of the cc-by-nc-4.0 license-: https://github.com/myshell-ai/OpenVoice/blob/main/LICENSE",
+ )
+
+ tts_button = gr.Button("Send", elem_id="send-btn", visible=True)
+
+
+ with gr.Column():
+ out_text_gr = gr.Text(label="Info")
+ audio_gr = gr.Audio(label="Synthesised Audio", autoplay=True)
+ ref_audio_gr = gr.Audio(label="Reference Audio Used")
+
+ gr.Examples(examples,
+ label="Examples",
+ inputs=[input_text_gr, style_gr, ref_gr, tos_gr],
+ outputs=[out_text_gr, audio_gr, ref_audio_gr],
+ fn=predict,
+ cache_examples=False,)
+ tts_button.click(predict, [input_text_gr, style_gr, ref_gr, tos_gr], outputs=[out_text_gr, audio_gr, ref_audio_gr])
+
+demo.queue()
+demo.launch(debug=True, show_api=True, share=args.share)
diff --git a/modules/openvoice/se_extractor.py b/modules/openvoice/se_extractor.py
new file mode 100644
index 0000000000000000000000000000000000000000..ed9cfc7d7ed11deb8b879eae39543e798dc703b4
--- /dev/null
+++ b/modules/openvoice/se_extractor.py
@@ -0,0 +1,153 @@
+import os
+import glob
+import torch
+import hashlib
+import librosa
+import base64
+from glob import glob
+import numpy as np
+from pydub import AudioSegment
+from faster_whisper import WhisperModel
+import hashlib
+import base64
+import librosa
+# from whisper_timestamped.transcribe import get_audio_tensor, get_vad_segments
+
+model_size = "medium"
+# Run on GPU with FP16
+model = None
+def split_audio_whisper(audio_path, audio_name, target_dir='processed'):
+ global model
+ if model is None:
+ model = WhisperModel(model_size, device="cuda", compute_type="float16")
+ audio = AudioSegment.from_file(audio_path)
+ max_len = len(audio)
+
+ target_folder = os.path.join(target_dir, audio_name)
+
+ segments, info = model.transcribe(audio_path, beam_size=5, word_timestamps=True)
+ segments = list(segments)
+
+ # create directory
+ os.makedirs(target_folder, exist_ok=True)
+ wavs_folder = os.path.join(target_folder, 'wavs')
+ os.makedirs(wavs_folder, exist_ok=True)
+
+ # segments
+ s_ind = 0
+ start_time = None
+
+ for k, w in enumerate(segments):
+ # process with the time
+ if k == 0:
+ start_time = max(0, w.start)
+
+ end_time = w.end
+
+ # calculate confidence
+ if len(w.words) > 0:
+ confidence = sum([s.probability for s in w.words]) / len(w.words)
+ else:
+ confidence = 0.
+ # clean text
+ text = w.text.replace('...', '')
+
+ # left 0.08s for each audios
+ audio_seg = audio[int( start_time * 1000) : min(max_len, int(end_time * 1000) + 80)]
+
+ # segment file name
+ fname = f"{audio_name}_seg{s_ind}.wav"
+
+ # filter out the segment shorter than 1.5s and longer than 20s
+ save = audio_seg.duration_seconds > 1.5 and \
+ audio_seg.duration_seconds < 20. and \
+ len(text) >= 2 and len(text) < 200
+
+ if save:
+ output_file = os.path.join(wavs_folder, fname)
+ audio_seg.export(output_file, format='wav')
+
+ if k < len(segments) - 1:
+ start_time = max(0, segments[k+1].start - 0.08)
+
+ s_ind = s_ind + 1
+ return wavs_folder
+
+
+def split_audio_vad(audio_path, audio_name, target_dir, split_seconds=10.0):
+ SAMPLE_RATE = 16000
+ audio_vad = get_audio_tensor(audio_path)
+ segments = get_vad_segments(
+ audio_vad,
+ output_sample=True,
+ min_speech_duration=0.1,
+ min_silence_duration=1,
+ method="silero",
+ )
+ segments = [(seg["start"], seg["end"]) for seg in segments]
+ segments = [(float(s) / SAMPLE_RATE, float(e) / SAMPLE_RATE) for s,e in segments]
+ print(segments)
+ audio_active = AudioSegment.silent(duration=0)
+ audio = AudioSegment.from_file(audio_path)
+
+ for start_time, end_time in segments:
+ audio_active += audio[int( start_time * 1000) : int(end_time * 1000)]
+
+ audio_dur = audio_active.duration_seconds
+ print(f'after vad: dur = {audio_dur}')
+ target_folder = os.path.join(target_dir, audio_name)
+ wavs_folder = os.path.join(target_folder, 'wavs')
+ os.makedirs(wavs_folder, exist_ok=True)
+ start_time = 0.
+ count = 0
+ num_splits = int(np.round(audio_dur / split_seconds))
+ assert num_splits > 0, 'input audio is too short'
+ interval = audio_dur / num_splits
+
+ for i in range(num_splits):
+ end_time = min(start_time + interval, audio_dur)
+ if i == num_splits - 1:
+ end_time = audio_dur
+ output_file = f"{wavs_folder}/{audio_name}_seg{count}.wav"
+ audio_seg = audio_active[int(start_time * 1000): int(end_time * 1000)]
+ audio_seg.export(output_file, format='wav')
+ start_time = end_time
+ count += 1
+ return wavs_folder
+
+def hash_numpy_array(audio_path):
+ array, _ = librosa.load(audio_path, sr=None, mono=True)
+ # Convert the array to bytes
+ array_bytes = array.tobytes()
+ # Calculate the hash of the array bytes
+ hash_object = hashlib.sha256(array_bytes)
+ hash_value = hash_object.digest()
+ # Convert the hash value to base64
+ base64_value = base64.b64encode(hash_value)
+ return base64_value.decode('utf-8')[:16].replace('/', '_^')
+
+def get_se(audio_path, vc_model, target_dir='processed', vad=True):
+ device = vc_model.device
+ version = vc_model.version
+ print("OpenVoice version:", version)
+
+ audio_name = f"{os.path.basename(audio_path).rsplit('.', 1)[0]}_{version}_{hash_numpy_array(audio_path)}"
+ se_path = os.path.join(target_dir, audio_name, 'se.pth')
+
+ # if os.path.isfile(se_path):
+ # se = torch.load(se_path).to(device)
+ # return se, audio_name
+ # if os.path.isdir(audio_path):
+ # wavs_folder = audio_path
+
+ # if vad:
+ # wavs_folder = split_audio_vad(audio_path, target_dir=target_dir, audio_name=audio_name)
+ # else:
+ # wavs_folder = split_audio_whisper(audio_path, target_dir=target_dir, audio_name=audio_name)
+
+ # audio_segs = glob(f'{wavs_folder}/*.wav')
+ # if len(audio_segs) == 0:
+ # raise NotImplementedError('No audio segments found!')
+
+ return vc_model.extract_se([audio_path], se_save_path=se_path), audio_name
+
diff --git a/modules/openvoice/transforms.py b/modules/openvoice/transforms.py
new file mode 100644
index 0000000000000000000000000000000000000000..7d93c482d296f8bafd9cb5f2a6b1d7c7608cd17c
--- /dev/null
+++ b/modules/openvoice/transforms.py
@@ -0,0 +1,209 @@
+import torch
+from torch.nn import functional as F
+
+import numpy as np
+
+
+DEFAULT_MIN_BIN_WIDTH = 1e-3
+DEFAULT_MIN_BIN_HEIGHT = 1e-3
+DEFAULT_MIN_DERIVATIVE = 1e-3
+
+
+def piecewise_rational_quadratic_transform(
+ inputs,
+ unnormalized_widths,
+ unnormalized_heights,
+ unnormalized_derivatives,
+ inverse=False,
+ tails=None,
+ tail_bound=1.0,
+ min_bin_width=DEFAULT_MIN_BIN_WIDTH,
+ min_bin_height=DEFAULT_MIN_BIN_HEIGHT,
+ min_derivative=DEFAULT_MIN_DERIVATIVE,
+):
+ if tails is None:
+ spline_fn = rational_quadratic_spline
+ spline_kwargs = {}
+ else:
+ spline_fn = unconstrained_rational_quadratic_spline
+ spline_kwargs = {"tails": tails, "tail_bound": tail_bound}
+
+ outputs, logabsdet = spline_fn(
+ inputs=inputs,
+ unnormalized_widths=unnormalized_widths,
+ unnormalized_heights=unnormalized_heights,
+ unnormalized_derivatives=unnormalized_derivatives,
+ inverse=inverse,
+ min_bin_width=min_bin_width,
+ min_bin_height=min_bin_height,
+ min_derivative=min_derivative,
+ **spline_kwargs
+ )
+ return outputs, logabsdet
+
+
+def searchsorted(bin_locations, inputs, eps=1e-6):
+ bin_locations[..., -1] += eps
+ return torch.sum(inputs[..., None] >= bin_locations, dim=-1) - 1
+
+
+def unconstrained_rational_quadratic_spline(
+ inputs,
+ unnormalized_widths,
+ unnormalized_heights,
+ unnormalized_derivatives,
+ inverse=False,
+ tails="linear",
+ tail_bound=1.0,
+ min_bin_width=DEFAULT_MIN_BIN_WIDTH,
+ min_bin_height=DEFAULT_MIN_BIN_HEIGHT,
+ min_derivative=DEFAULT_MIN_DERIVATIVE,
+):
+ inside_interval_mask = (inputs >= -tail_bound) & (inputs <= tail_bound)
+ outside_interval_mask = ~inside_interval_mask
+
+ outputs = torch.zeros_like(inputs)
+ logabsdet = torch.zeros_like(inputs)
+
+ if tails == "linear":
+ unnormalized_derivatives = F.pad(unnormalized_derivatives, pad=(1, 1))
+ constant = np.log(np.exp(1 - min_derivative) - 1)
+ unnormalized_derivatives[..., 0] = constant
+ unnormalized_derivatives[..., -1] = constant
+
+ outputs[outside_interval_mask] = inputs[outside_interval_mask]
+ logabsdet[outside_interval_mask] = 0
+ else:
+ raise RuntimeError("{} tails are not implemented.".format(tails))
+
+ (
+ outputs[inside_interval_mask],
+ logabsdet[inside_interval_mask],
+ ) = rational_quadratic_spline(
+ inputs=inputs[inside_interval_mask],
+ unnormalized_widths=unnormalized_widths[inside_interval_mask, :],
+ unnormalized_heights=unnormalized_heights[inside_interval_mask, :],
+ unnormalized_derivatives=unnormalized_derivatives[inside_interval_mask, :],
+ inverse=inverse,
+ left=-tail_bound,
+ right=tail_bound,
+ bottom=-tail_bound,
+ top=tail_bound,
+ min_bin_width=min_bin_width,
+ min_bin_height=min_bin_height,
+ min_derivative=min_derivative,
+ )
+
+ return outputs, logabsdet
+
+
+def rational_quadratic_spline(
+ inputs,
+ unnormalized_widths,
+ unnormalized_heights,
+ unnormalized_derivatives,
+ inverse=False,
+ left=0.0,
+ right=1.0,
+ bottom=0.0,
+ top=1.0,
+ min_bin_width=DEFAULT_MIN_BIN_WIDTH,
+ min_bin_height=DEFAULT_MIN_BIN_HEIGHT,
+ min_derivative=DEFAULT_MIN_DERIVATIVE,
+):
+ if torch.min(inputs) < left or torch.max(inputs) > right:
+ raise ValueError("Input to a transform is not within its domain")
+
+ num_bins = unnormalized_widths.shape[-1]
+
+ if min_bin_width * num_bins > 1.0:
+ raise ValueError("Minimal bin width too large for the number of bins")
+ if min_bin_height * num_bins > 1.0:
+ raise ValueError("Minimal bin height too large for the number of bins")
+
+ widths = F.softmax(unnormalized_widths, dim=-1)
+ widths = min_bin_width + (1 - min_bin_width * num_bins) * widths
+ cumwidths = torch.cumsum(widths, dim=-1)
+ cumwidths = F.pad(cumwidths, pad=(1, 0), mode="constant", value=0.0)
+ cumwidths = (right - left) * cumwidths + left
+ cumwidths[..., 0] = left
+ cumwidths[..., -1] = right
+ widths = cumwidths[..., 1:] - cumwidths[..., :-1]
+
+ derivatives = min_derivative + F.softplus(unnormalized_derivatives)
+
+ heights = F.softmax(unnormalized_heights, dim=-1)
+ heights = min_bin_height + (1 - min_bin_height * num_bins) * heights
+ cumheights = torch.cumsum(heights, dim=-1)
+ cumheights = F.pad(cumheights, pad=(1, 0), mode="constant", value=0.0)
+ cumheights = (top - bottom) * cumheights + bottom
+ cumheights[..., 0] = bottom
+ cumheights[..., -1] = top
+ heights = cumheights[..., 1:] - cumheights[..., :-1]
+
+ if inverse:
+ bin_idx = searchsorted(cumheights, inputs)[..., None]
+ else:
+ bin_idx = searchsorted(cumwidths, inputs)[..., None]
+
+ input_cumwidths = cumwidths.gather(-1, bin_idx)[..., 0]
+ input_bin_widths = widths.gather(-1, bin_idx)[..., 0]
+
+ input_cumheights = cumheights.gather(-1, bin_idx)[..., 0]
+ delta = heights / widths
+ input_delta = delta.gather(-1, bin_idx)[..., 0]
+
+ input_derivatives = derivatives.gather(-1, bin_idx)[..., 0]
+ input_derivatives_plus_one = derivatives[..., 1:].gather(-1, bin_idx)[..., 0]
+
+ input_heights = heights.gather(-1, bin_idx)[..., 0]
+
+ if inverse:
+ a = (inputs - input_cumheights) * (
+ input_derivatives + input_derivatives_plus_one - 2 * input_delta
+ ) + input_heights * (input_delta - input_derivatives)
+ b = input_heights * input_derivatives - (inputs - input_cumheights) * (
+ input_derivatives + input_derivatives_plus_one - 2 * input_delta
+ )
+ c = -input_delta * (inputs - input_cumheights)
+
+ discriminant = b.pow(2) - 4 * a * c
+ assert (discriminant >= 0).all()
+
+ root = (2 * c) / (-b - torch.sqrt(discriminant))
+ outputs = root * input_bin_widths + input_cumwidths
+
+ theta_one_minus_theta = root * (1 - root)
+ denominator = input_delta + (
+ (input_derivatives + input_derivatives_plus_one - 2 * input_delta)
+ * theta_one_minus_theta
+ )
+ derivative_numerator = input_delta.pow(2) * (
+ input_derivatives_plus_one * root.pow(2)
+ + 2 * input_delta * theta_one_minus_theta
+ + input_derivatives * (1 - root).pow(2)
+ )
+ logabsdet = torch.log(derivative_numerator) - 2 * torch.log(denominator)
+
+ return outputs, -logabsdet
+ else:
+ theta = (inputs - input_cumwidths) / input_bin_widths
+ theta_one_minus_theta = theta * (1 - theta)
+
+ numerator = input_heights * (
+ input_delta * theta.pow(2) + input_derivatives * theta_one_minus_theta
+ )
+ denominator = input_delta + (
+ (input_derivatives + input_derivatives_plus_one - 2 * input_delta)
+ * theta_one_minus_theta
+ )
+ outputs = input_cumheights + numerator / denominator
+
+ derivative_numerator = input_delta.pow(2) * (
+ input_derivatives_plus_one * theta.pow(2)
+ + 2 * input_delta * theta_one_minus_theta
+ + input_derivatives * (1 - theta).pow(2)
+ )
+ logabsdet = torch.log(derivative_numerator) - 2 * torch.log(denominator)
+
+ return outputs, logabsdet
diff --git a/modules/openvoice/utils.py b/modules/openvoice/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..ec51484186b335d7ee43679a93c58780175dab0e
--- /dev/null
+++ b/modules/openvoice/utils.py
@@ -0,0 +1,194 @@
+import re
+import json
+import numpy as np
+
+
+def get_hparams_from_file(config_path):
+ with open(config_path, "r", encoding="utf-8") as f:
+ data = f.read()
+ config = json.loads(data)
+
+ hparams = HParams(**config)
+ return hparams
+
+class HParams:
+ def __init__(self, **kwargs):
+ for k, v in kwargs.items():
+ if type(v) == dict:
+ v = HParams(**v)
+ self[k] = v
+
+ def keys(self):
+ return self.__dict__.keys()
+
+ def items(self):
+ return self.__dict__.items()
+
+ def values(self):
+ return self.__dict__.values()
+
+ def __len__(self):
+ return len(self.__dict__)
+
+ def __getitem__(self, key):
+ return getattr(self, key)
+
+ def __setitem__(self, key, value):
+ return setattr(self, key, value)
+
+ def __contains__(self, key):
+ return key in self.__dict__
+
+ def __repr__(self):
+ return self.__dict__.__repr__()
+
+
+def string_to_bits(string, pad_len=8):
+ # Convert each character to its ASCII value
+ ascii_values = [ord(char) for char in string]
+
+ # Convert ASCII values to binary representation
+ binary_values = [bin(value)[2:].zfill(8) for value in ascii_values]
+
+ # Convert binary strings to integer arrays
+ bit_arrays = [[int(bit) for bit in binary] for binary in binary_values]
+
+ # Convert list of arrays to NumPy array
+ numpy_array = np.array(bit_arrays)
+ numpy_array_full = np.zeros((pad_len, 8), dtype=numpy_array.dtype)
+ numpy_array_full[:, 2] = 1
+ max_len = min(pad_len, len(numpy_array))
+ numpy_array_full[:max_len] = numpy_array[:max_len]
+ return numpy_array_full
+
+
+def bits_to_string(bits_array):
+ # Convert each row of the array to a binary string
+ binary_values = [''.join(str(bit) for bit in row) for row in bits_array]
+
+ # Convert binary strings to ASCII values
+ ascii_values = [int(binary, 2) for binary in binary_values]
+
+ # Convert ASCII values to characters
+ output_string = ''.join(chr(value) for value in ascii_values)
+
+ return output_string
+
+
+def split_sentence(text, min_len=10, language_str='[EN]'):
+ if language_str in ['EN']:
+ sentences = split_sentences_latin(text, min_len=min_len)
+ else:
+ sentences = split_sentences_zh(text, min_len=min_len)
+ return sentences
+
+def split_sentences_latin(text, min_len=10):
+ """Split Long sentences into list of short ones
+
+ Args:
+ str: Input sentences.
+
+ Returns:
+ List[str]: list of output sentences.
+ """
+ # deal with dirty sentences
+ text = re.sub('[。！？；]', '.', text)
+ text = re.sub('[，]', ',', text)
+ text = re.sub('[“”]', '"', text)
+ text = re.sub('[‘’]', "'", text)
+ text = re.sub(r"[\<\>\(\)\[\]\"\«\»]+", "", text)
+ text = re.sub('[\n\t ]+', ' ', text)
+ text = re.sub('([,.!?;])', r'\1 $#!', text)
+ # split
+ sentences = [s.strip() for s in text.split('$#!')]
+ if len(sentences[-1]) == 0: del sentences[-1]
+
+ new_sentences = []
+ new_sent = []
+ count_len = 0
+ for ind, sent in enumerate(sentences):
+ # print(sent)
+ new_sent.append(sent)
+ count_len += len(sent.split(" "))
+ if count_len > min_len or ind == len(sentences) - 1:
+ count_len = 0
+ new_sentences.append(' '.join(new_sent))
+ new_sent = []
+ return merge_short_sentences_latin(new_sentences)
+
+
+def merge_short_sentences_latin(sens):
+ """Avoid short sentences by merging them with the following sentence.
+
+ Args:
+ List[str]: list of input sentences.
+
+ Returns:
+ List[str]: list of output sentences.
+ """
+ sens_out = []
+ for s in sens:
+ # If the previous sentence is too short, merge them with
+ # the current sentence.
+ if len(sens_out) > 0 and len(sens_out[-1].split(" ")) <= 2:
+ sens_out[-1] = sens_out[-1] + " " + s
+ else:
+ sens_out.append(s)
+ try:
+ if len(sens_out[-1].split(" ")) <= 2:
+ sens_out[-2] = sens_out[-2] + " " + sens_out[-1]
+ sens_out.pop(-1)
+ except:
+ pass
+ return sens_out
+
+def split_sentences_zh(text, min_len=10):
+ text = re.sub('[。！？；]', '.', text)
+ text = re.sub('[，]', ',', text)
+ # 将文本中的换行符、空格和制表符替换为空格
+ text = re.sub('[\n\t ]+', ' ', text)
+ # 在标点符号后添加一个空格
+ text = re.sub('([,.!?;])', r'\1 $#!', text)
+ # 分隔句子并去除前后空格
+ # sentences = [s.strip() for s in re.split('(。|！|？|；)', text)]
+ sentences = [s.strip() for s in text.split('$#!')]
+ if len(sentences[-1]) == 0: del sentences[-1]
+
+ new_sentences = []
+ new_sent = []
+ count_len = 0
+ for ind, sent in enumerate(sentences):
+ new_sent.append(sent)
+ count_len += len(sent)
+ if count_len > min_len or ind == len(sentences) - 1:
+ count_len = 0
+ new_sentences.append(' '.join(new_sent))
+ new_sent = []
+ return merge_short_sentences_zh(new_sentences)
+
+
+def merge_short_sentences_zh(sens):
+ # return sens
+ """Avoid short sentences by merging them with the following sentence.
+
+ Args:
+ List[str]: list of input sentences.
+
+ Returns:
+ List[str]: list of output sentences.
+ """
+ sens_out = []
+ for s in sens:
+ # If the previous sentense is too short, merge them with
+ # the current sentence.
+ if len(sens_out) > 0 and len(sens_out[-1]) <= 2:
+ sens_out[-1] = sens_out[-1] + " " + s
+ else:
+ sens_out.append(s)
+ try:
+ if len(sens_out[-1]) <= 2:
+ sens_out[-2] = sens_out[-2] + " " + sens_out[-1]
+ sens_out.pop(-1)
+ except:
+ pass
+ return sens_out
\ No newline at end of file
diff --git a/modules/quantize.py b/modules/quantize.py
new file mode 100644
index 0000000000000000000000000000000000000000..1736ea33b72ce7292e068ad418e411c1b76f4b73
--- /dev/null
+++ b/modules/quantize.py
@@ -0,0 +1,229 @@
+from dac.nn.quantize import ResidualVectorQuantize
+from torch import nn
+from modules.wavenet import WN
+import torch
+import torchaudio
+import torchaudio.functional as audio_F
+import numpy as np
+from .alias_free_torch import *
+from torch.nn.utils import weight_norm
+from torch import nn, sin, pow
+from einops.layers.torch import Rearrange
+from dac.model.encodec import SConv1d
+
+def init_weights(m):
+ if isinstance(m, nn.Conv1d):
+ nn.init.trunc_normal_(m.weight, std=0.02)
+ nn.init.constant_(m.bias, 0)
+
+
+def WNConv1d(*args, **kwargs):
+ return weight_norm(nn.Conv1d(*args, **kwargs))
+
+
+def WNConvTranspose1d(*args, **kwargs):
+ return weight_norm(nn.ConvTranspose1d(*args, **kwargs))
+
+class SnakeBeta(nn.Module):
+ """
+ A modified Snake function which uses separate parameters for the magnitude of the periodic components
+ Shape:
+ - Input: (B, C, T)
+ - Output: (B, C, T), same shape as the input
+ Parameters:
+ - alpha - trainable parameter that controls frequency
+ - beta - trainable parameter that controls magnitude
+ References:
+ - This activation function is a modified version based on this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda:
+ https://arxiv.org/abs/2006.08195
+ Examples:
+ >>> a1 = snakebeta(256)
+ >>> x = torch.randn(256)
+ >>> x = a1(x)
+ """
+
+ def __init__(
+ self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False
+ ):
+ """
+ Initialization.
+ INPUT:
+ - in_features: shape of the input
+ - alpha - trainable parameter that controls frequency
+ - beta - trainable parameter that controls magnitude
+ alpha is initialized to 1 by default, higher values = higher-frequency.
+ beta is initialized to 1 by default, higher values = higher-magnitude.
+ alpha will be trained along with the rest of your model.
+ """
+ super(SnakeBeta, self).__init__()
+ self.in_features = in_features
+
+ # initialize alpha
+ self.alpha_logscale = alpha_logscale
+ if self.alpha_logscale: # log scale alphas initialized to zeros
+ self.alpha = nn.Parameter(torch.zeros(in_features) * alpha)
+ self.beta = nn.Parameter(torch.zeros(in_features) * alpha)
+ else: # linear scale alphas initialized to ones
+ self.alpha = nn.Parameter(torch.ones(in_features) * alpha)
+ self.beta = nn.Parameter(torch.ones(in_features) * alpha)
+
+ self.alpha.requires_grad = alpha_trainable
+ self.beta.requires_grad = alpha_trainable
+
+ self.no_div_by_zero = 0.000000001
+
+ def forward(self, x):
+ """
+ Forward pass of the function.
+ Applies the function to the input elementwise.
+ SnakeBeta := x + 1/b * sin^2 (xa)
+ """
+ alpha = self.alpha.unsqueeze(0).unsqueeze(-1) # line up with x to [B, C, T]
+ beta = self.beta.unsqueeze(0).unsqueeze(-1)
+ if self.alpha_logscale:
+ alpha = torch.exp(alpha)
+ beta = torch.exp(beta)
+ x = x + (1.0 / (beta + self.no_div_by_zero)) * pow(sin(x * alpha), 2)
+
+ return x
+
+class ResidualUnit(nn.Module):
+ def __init__(self, dim: int = 16, dilation: int = 1):
+ super().__init__()
+ pad = ((7 - 1) * dilation) // 2
+ self.block = nn.Sequential(
+ Activation1d(activation=SnakeBeta(dim, alpha_logscale=True)),
+ WNConv1d(dim, dim, kernel_size=7, dilation=dilation, padding=pad),
+ Activation1d(activation=SnakeBeta(dim, alpha_logscale=True)),
+ WNConv1d(dim, dim, kernel_size=1),
+ )
+
+ def forward(self, x):
+ return x + self.block(x)
+
+class CNNLSTM(nn.Module):
+ def __init__(self, indim, outdim, head, global_pred=False):
+ super().__init__()
+ self.global_pred = global_pred
+ self.model = nn.Sequential(
+ ResidualUnit(indim, dilation=1),
+ ResidualUnit(indim, dilation=2),
+ ResidualUnit(indim, dilation=3),
+ Activation1d(activation=SnakeBeta(indim, alpha_logscale=True)),
+ Rearrange("b c t -> b t c"),
+ )
+ self.heads = nn.ModuleList([nn.Linear(indim, outdim) for i in range(head)])
+
+ def forward(self, x):
+ # x: [B, C, T]
+ x = self.model(x)
+ if self.global_pred:
+ x = torch.mean(x, dim=1, keepdim=False)
+ outs = [head(x) for head in self.heads]
+ return outs
+
+def sequence_mask(length, max_length=None):
+ if max_length is None:
+ max_length = length.max()
+ x = torch.arange(max_length, dtype=length.dtype, device=length.device)
+ return x.unsqueeze(0) < length.unsqueeze(1)
+class FAquantizer(nn.Module):
+ def __init__(self, in_dim=1024,
+ n_p_codebooks=1,
+ n_c_codebooks=2,
+ n_t_codebooks=2,
+ n_r_codebooks=3,
+ codebook_size=1024,
+ codebook_dim=8,
+ quantizer_dropout=0.5,
+ causal=False,
+ separate_prosody_encoder=False,
+ timbre_norm=False,):
+ super(FAquantizer, self).__init__()
+ conv1d_type = SConv1d# if causal else nn.Conv1d
+ self.prosody_quantizer = ResidualVectorQuantize(
+ input_dim=in_dim,
+ n_codebooks=n_p_codebooks,
+ codebook_size=codebook_size,
+ codebook_dim=codebook_dim,
+ quantizer_dropout=quantizer_dropout,
+ )
+
+ self.content_quantizer = ResidualVectorQuantize(
+ input_dim=in_dim,
+ n_codebooks=n_c_codebooks,
+ codebook_size=codebook_size,
+ codebook_dim=codebook_dim,
+ quantizer_dropout=quantizer_dropout,
+ )
+
+ self.residual_quantizer = ResidualVectorQuantize(
+ input_dim=in_dim,
+ n_codebooks=n_r_codebooks,
+ codebook_size=codebook_size,
+ codebook_dim=codebook_dim,
+ quantizer_dropout=quantizer_dropout,
+ )
+
+ self.melspec_linear = conv1d_type(in_channels=20, out_channels=256, kernel_size=1, causal=causal)
+ self.melspec_encoder = WN(hidden_channels=256, kernel_size=5, dilation_rate=1, n_layers=8, gin_channels=0, p_dropout=0.2, causal=causal)
+ self.melspec_linear2 = conv1d_type(in_channels=256, out_channels=1024, kernel_size=1, causal=causal)
+
+ self.prob_random_mask_residual = 0.75
+
+ SPECT_PARAMS = {
+ "n_fft": 2048,
+ "win_length": 1200,
+ "hop_length": 300,
+ }
+ MEL_PARAMS = {
+ "n_mels": 80,
+ }
+
+ self.to_mel = torchaudio.transforms.MelSpectrogram(
+ n_mels=MEL_PARAMS["n_mels"], sample_rate=24000, **SPECT_PARAMS
+ )
+ self.mel_mean, self.mel_std = -4, 4
+ self.frame_rate = 24000 / 300
+ self.hop_length = 300
+
+ def preprocess(self, wave_tensor, n_bins=20):
+ mel_tensor = self.to_mel(wave_tensor.squeeze(1))
+ mel_tensor = (torch.log(1e-5 + mel_tensor) - self.mel_mean) / self.mel_std
+ return mel_tensor[:, :n_bins, :int(wave_tensor.size(-1) / self.hop_length)]
+
+ def forward(self, x, wave_segments):
+ outs = 0
+ prosody_feature = self.preprocess(wave_segments)
+
+ f0_input = prosody_feature # (B, T, 20)
+ f0_input = self.melspec_linear(f0_input)
+ f0_input = self.melspec_encoder(f0_input, torch.ones(f0_input.shape[0], 1, f0_input.shape[2]).to(
+ f0_input.device).bool())
+ f0_input = self.melspec_linear2(f0_input)
+
+ common_min_size = min(f0_input.size(2), x.size(2))
+ f0_input = f0_input[:, :, :common_min_size]
+
+ x = x[:, :, :common_min_size]
+
+ z_p, codes_p, latents_p, commitment_loss_p, codebook_loss_p = self.prosody_quantizer(
+ f0_input, 1
+ )
+ outs += z_p.detach()
+
+ z_c, codes_c, latents_c, commitment_loss_c, codebook_loss_c = self.content_quantizer(
+ x, 2
+ )
+ outs += z_c.detach()
+
+ residual_feature = x - z_p.detach() - z_c.detach()
+
+ z_r, codes_r, latents_r, commitment_loss_r, codebook_loss_r = self.residual_quantizer(
+ residual_feature, 3
+ )
+
+ quantized = [z_p, z_c, z_r]
+ codes = [codes_p, codes_c, codes_r]
+
+ return quantized, codes
\ No newline at end of file
diff --git a/modules/rmvpe.py b/modules/rmvpe.py
new file mode 100644
index 0000000000000000000000000000000000000000..8b4b98747f5cf76906c2e26c0883606f04d1c805
--- /dev/null
+++ b/modules/rmvpe.py
@@ -0,0 +1,631 @@
+from io import BytesIO
+import os
+from typing import List, Optional, Tuple
+import numpy as np
+import torch
+
+import torch.nn as nn
+import torch.nn.functional as F
+from librosa.util import normalize, pad_center, tiny
+from scipy.signal import get_window
+
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+class STFT(torch.nn.Module):
+ def __init__(
+ self, filter_length=1024, hop_length=512, win_length=None, window="hann"
+ ):
+ """
+ This module implements an STFT using 1D convolution and 1D transpose convolutions.
+ This is a bit tricky so there are some cases that probably won't work as working
+ out the same sizes before and after in all overlap add setups is tough. Right now,
+ this code should work with hop lengths that are half the filter length (50% overlap
+ between frames).
+
+ Keyword Arguments:
+ filter_length {int} -- Length of filters used (default: {1024})
+ hop_length {int} -- Hop length of STFT (restrict to 50% overlap between frames) (default: {512})
+ win_length {[type]} -- Length of the window function applied to each frame (if not specified, it
+ equals the filter length). (default: {None})
+ window {str} -- Type of window to use (options are bartlett, hann, hamming, blackman, blackmanharris)
+ (default: {'hann'})
+ """
+ super(STFT, self).__init__()
+ self.filter_length = filter_length
+ self.hop_length = hop_length
+ self.win_length = win_length if win_length else filter_length
+ self.window = window
+ self.forward_transform = None
+ self.pad_amount = int(self.filter_length / 2)
+ fourier_basis = np.fft.fft(np.eye(self.filter_length))
+
+ cutoff = int((self.filter_length / 2 + 1))
+ fourier_basis = np.vstack(
+ [np.real(fourier_basis[:cutoff, :]), np.imag(fourier_basis[:cutoff, :])]
+ )
+ forward_basis = torch.FloatTensor(fourier_basis)
+ inverse_basis = torch.FloatTensor(np.linalg.pinv(fourier_basis))
+
+ assert filter_length >= self.win_length
+ # get window and zero center pad it to filter_length
+ fft_window = get_window(window, self.win_length, fftbins=True)
+ fft_window = pad_center(fft_window, size=filter_length)
+ fft_window = torch.from_numpy(fft_window).float()
+
+ # window the bases
+ forward_basis *= fft_window
+ inverse_basis = (inverse_basis.T * fft_window).T
+
+ self.register_buffer("forward_basis", forward_basis.float())
+ self.register_buffer("inverse_basis", inverse_basis.float())
+ self.register_buffer("fft_window", fft_window.float())
+
+ def transform(self, input_data, return_phase=False):
+ """Take input data (audio) to STFT domain.
+
+ Arguments:
+ input_data {tensor} -- Tensor of floats, with shape (num_batch, num_samples)
+
+ Returns:
+ magnitude {tensor} -- Magnitude of STFT with shape (num_batch,
+ num_frequencies, num_frames)
+ phase {tensor} -- Phase of STFT with shape (num_batch,
+ num_frequencies, num_frames)
+ """
+ input_data = F.pad(
+ input_data,
+ (self.pad_amount, self.pad_amount),
+ mode="reflect",
+ )
+ forward_transform = input_data.unfold(
+ 1, self.filter_length, self.hop_length
+ ).permute(0, 2, 1)
+ forward_transform = torch.matmul(self.forward_basis, forward_transform)
+ cutoff = int((self.filter_length / 2) + 1)
+ real_part = forward_transform[:, :cutoff, :]
+ imag_part = forward_transform[:, cutoff:, :]
+ magnitude = torch.sqrt(real_part**2 + imag_part**2)
+ if return_phase:
+ phase = torch.atan2(imag_part.data, real_part.data)
+ return magnitude, phase
+ else:
+ return magnitude
+
+ def inverse(self, magnitude, phase):
+ """Call the inverse STFT (iSTFT), given magnitude and phase tensors produced
+ by the ```transform``` function.
+
+ Arguments:
+ magnitude {tensor} -- Magnitude of STFT with shape (num_batch,
+ num_frequencies, num_frames)
+ phase {tensor} -- Phase of STFT with shape (num_batch,
+ num_frequencies, num_frames)
+
+ Returns:
+ inverse_transform {tensor} -- Reconstructed audio given magnitude and phase. Of
+ shape (num_batch, num_samples)
+ """
+ cat = torch.cat(
+ [magnitude * torch.cos(phase), magnitude * torch.sin(phase)], dim=1
+ )
+ fold = torch.nn.Fold(
+ output_size=(1, (cat.size(-1) - 1) * self.hop_length + self.filter_length),
+ kernel_size=(1, self.filter_length),
+ stride=(1, self.hop_length),
+ )
+ inverse_transform = torch.matmul(self.inverse_basis, cat)
+ inverse_transform = fold(inverse_transform)[
+ :, 0, 0, self.pad_amount : -self.pad_amount
+ ]
+ window_square_sum = (
+ self.fft_window.pow(2).repeat(cat.size(-1), 1).T.unsqueeze(0)
+ )
+ window_square_sum = fold(window_square_sum)[
+ :, 0, 0, self.pad_amount : -self.pad_amount
+ ]
+ inverse_transform /= window_square_sum
+ return inverse_transform
+
+ def forward(self, input_data):
+ """Take input data (audio) to STFT domain and then back to audio.
+
+ Arguments:
+ input_data {tensor} -- Tensor of floats, with shape (num_batch, num_samples)
+
+ Returns:
+ reconstruction {tensor} -- Reconstructed audio given magnitude and phase. Of
+ shape (num_batch, num_samples)
+ """
+ self.magnitude, self.phase = self.transform(input_data, return_phase=True)
+ reconstruction = self.inverse(self.magnitude, self.phase)
+ return reconstruction
+
+
+from time import time as ttime
+
+
+class BiGRU(nn.Module):
+ def __init__(self, input_features, hidden_features, num_layers):
+ super(BiGRU, self).__init__()
+ self.gru = nn.GRU(
+ input_features,
+ hidden_features,
+ num_layers=num_layers,
+ batch_first=True,
+ bidirectional=True,
+ )
+
+ def forward(self, x):
+ return self.gru(x)[0]
+
+
+class ConvBlockRes(nn.Module):
+ def __init__(self, in_channels, out_channels, momentum=0.01):
+ super(ConvBlockRes, self).__init__()
+ self.conv = nn.Sequential(
+ nn.Conv2d(
+ in_channels=in_channels,
+ out_channels=out_channels,
+ kernel_size=(3, 3),
+ stride=(1, 1),
+ padding=(1, 1),
+ bias=False,
+ ),
+ nn.BatchNorm2d(out_channels, momentum=momentum),
+ nn.ReLU(),
+ nn.Conv2d(
+ in_channels=out_channels,
+ out_channels=out_channels,
+ kernel_size=(3, 3),
+ stride=(1, 1),
+ padding=(1, 1),
+ bias=False,
+ ),
+ nn.BatchNorm2d(out_channels, momentum=momentum),
+ nn.ReLU(),
+ )
+ # self.shortcut:Optional[nn.Module] = None
+ if in_channels != out_channels:
+ self.shortcut = nn.Conv2d(in_channels, out_channels, (1, 1))
+
+ def forward(self, x: torch.Tensor):
+ if not hasattr(self, "shortcut"):
+ return self.conv(x) + x
+ else:
+ return self.conv(x) + self.shortcut(x)
+
+
+class Encoder(nn.Module):
+ def __init__(
+ self,
+ in_channels,
+ in_size,
+ n_encoders,
+ kernel_size,
+ n_blocks,
+ out_channels=16,
+ momentum=0.01,
+ ):
+ super(Encoder, self).__init__()
+ self.n_encoders = n_encoders
+ self.bn = nn.BatchNorm2d(in_channels, momentum=momentum)
+ self.layers = nn.ModuleList()
+ self.latent_channels = []
+ for i in range(self.n_encoders):
+ self.layers.append(
+ ResEncoderBlock(
+ in_channels, out_channels, kernel_size, n_blocks, momentum=momentum
+ )
+ )
+ self.latent_channels.append([out_channels, in_size])
+ in_channels = out_channels
+ out_channels *= 2
+ in_size //= 2
+ self.out_size = in_size
+ self.out_channel = out_channels
+
+ def forward(self, x: torch.Tensor):
+ concat_tensors: List[torch.Tensor] = []
+ x = self.bn(x)
+ for i, layer in enumerate(self.layers):
+ t, x = layer(x)
+ concat_tensors.append(t)
+ return x, concat_tensors
+
+
+class ResEncoderBlock(nn.Module):
+ def __init__(
+ self, in_channels, out_channels, kernel_size, n_blocks=1, momentum=0.01
+ ):
+ super(ResEncoderBlock, self).__init__()
+ self.n_blocks = n_blocks
+ self.conv = nn.ModuleList()
+ self.conv.append(ConvBlockRes(in_channels, out_channels, momentum))
+ for i in range(n_blocks - 1):
+ self.conv.append(ConvBlockRes(out_channels, out_channels, momentum))
+ self.kernel_size = kernel_size
+ if self.kernel_size is not None:
+ self.pool = nn.AvgPool2d(kernel_size=kernel_size)
+
+ def forward(self, x):
+ for i, conv in enumerate(self.conv):
+ x = conv(x)
+ if self.kernel_size is not None:
+ return x, self.pool(x)
+ else:
+ return x
+
+
+class Intermediate(nn.Module): #
+ def __init__(self, in_channels, out_channels, n_inters, n_blocks, momentum=0.01):
+ super(Intermediate, self).__init__()
+ self.n_inters = n_inters
+ self.layers = nn.ModuleList()
+ self.layers.append(
+ ResEncoderBlock(in_channels, out_channels, None, n_blocks, momentum)
+ )
+ for i in range(self.n_inters - 1):
+ self.layers.append(
+ ResEncoderBlock(out_channels, out_channels, None, n_blocks, momentum)
+ )
+
+ def forward(self, x):
+ for i, layer in enumerate(self.layers):
+ x = layer(x)
+ return x
+
+
+class ResDecoderBlock(nn.Module):
+ def __init__(self, in_channels, out_channels, stride, n_blocks=1, momentum=0.01):
+ super(ResDecoderBlock, self).__init__()
+ out_padding = (0, 1) if stride == (1, 2) else (1, 1)
+ self.n_blocks = n_blocks
+ self.conv1 = nn.Sequential(
+ nn.ConvTranspose2d(
+ in_channels=in_channels,
+ out_channels=out_channels,
+ kernel_size=(3, 3),
+ stride=stride,
+ padding=(1, 1),
+ output_padding=out_padding,
+ bias=False,
+ ),
+ nn.BatchNorm2d(out_channels, momentum=momentum),
+ nn.ReLU(),
+ )
+ self.conv2 = nn.ModuleList()
+ self.conv2.append(ConvBlockRes(out_channels * 2, out_channels, momentum))
+ for i in range(n_blocks - 1):
+ self.conv2.append(ConvBlockRes(out_channels, out_channels, momentum))
+
+ def forward(self, x, concat_tensor):
+ x = self.conv1(x)
+ x = torch.cat((x, concat_tensor), dim=1)
+ for i, conv2 in enumerate(self.conv2):
+ x = conv2(x)
+ return x
+
+
+class Decoder(nn.Module):
+ def __init__(self, in_channels, n_decoders, stride, n_blocks, momentum=0.01):
+ super(Decoder, self).__init__()
+ self.layers = nn.ModuleList()
+ self.n_decoders = n_decoders
+ for i in range(self.n_decoders):
+ out_channels = in_channels // 2
+ self.layers.append(
+ ResDecoderBlock(in_channels, out_channels, stride, n_blocks, momentum)
+ )
+ in_channels = out_channels
+
+ def forward(self, x: torch.Tensor, concat_tensors: List[torch.Tensor]):
+ for i, layer in enumerate(self.layers):
+ x = layer(x, concat_tensors[-1 - i])
+ return x
+
+
+class DeepUnet(nn.Module):
+ def __init__(
+ self,
+ kernel_size,
+ n_blocks,
+ en_de_layers=5,
+ inter_layers=4,
+ in_channels=1,
+ en_out_channels=16,
+ ):
+ super(DeepUnet, self).__init__()
+ self.encoder = Encoder(
+ in_channels, 128, en_de_layers, kernel_size, n_blocks, en_out_channels
+ )
+ self.intermediate = Intermediate(
+ self.encoder.out_channel // 2,
+ self.encoder.out_channel,
+ inter_layers,
+ n_blocks,
+ )
+ self.decoder = Decoder(
+ self.encoder.out_channel, en_de_layers, kernel_size, n_blocks
+ )
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ x, concat_tensors = self.encoder(x)
+ x = self.intermediate(x)
+ x = self.decoder(x, concat_tensors)
+ return x
+
+
+class E2E(nn.Module):
+ def __init__(
+ self,
+ n_blocks,
+ n_gru,
+ kernel_size,
+ en_de_layers=5,
+ inter_layers=4,
+ in_channels=1,
+ en_out_channels=16,
+ ):
+ super(E2E, self).__init__()
+ self.unet = DeepUnet(
+ kernel_size,
+ n_blocks,
+ en_de_layers,
+ inter_layers,
+ in_channels,
+ en_out_channels,
+ )
+ self.cnn = nn.Conv2d(en_out_channels, 3, (3, 3), padding=(1, 1))
+ if n_gru:
+ self.fc = nn.Sequential(
+ BiGRU(3 * 128, 256, n_gru),
+ nn.Linear(512, 360),
+ nn.Dropout(0.25),
+ nn.Sigmoid(),
+ )
+ else:
+ self.fc = nn.Sequential(
+ nn.Linear(3 * nn.N_MELS, nn.N_CLASS), nn.Dropout(0.25), nn.Sigmoid()
+ )
+
+ def forward(self, mel):
+ # print(mel.shape)
+ mel = mel.transpose(-1, -2).unsqueeze(1)
+ x = self.cnn(self.unet(mel)).transpose(1, 2).flatten(-2)
+ x = self.fc(x)
+ # print(x.shape)
+ return x
+
+
+from librosa.filters import mel
+
+
+class MelSpectrogram(torch.nn.Module):
+ def __init__(
+ self,
+ is_half,
+ n_mel_channels,
+ sampling_rate,
+ win_length,
+ hop_length,
+ n_fft=None,
+ mel_fmin=0,
+ mel_fmax=None,
+ clamp=1e-5,
+ ):
+ super().__init__()
+ n_fft = win_length if n_fft is None else n_fft
+ self.hann_window = {}
+ mel_basis = mel(
+ sr=sampling_rate,
+ n_fft=n_fft,
+ n_mels=n_mel_channels,
+ fmin=mel_fmin,
+ fmax=mel_fmax,
+ htk=True,
+ )
+ mel_basis = torch.from_numpy(mel_basis).float()
+ self.register_buffer("mel_basis", mel_basis)
+ self.n_fft = win_length if n_fft is None else n_fft
+ self.hop_length = hop_length
+ self.win_length = win_length
+ self.sampling_rate = sampling_rate
+ self.n_mel_channels = n_mel_channels
+ self.clamp = clamp
+ self.is_half = is_half
+
+ def forward(self, audio, keyshift=0, speed=1, center=True):
+ factor = 2 ** (keyshift / 12)
+ n_fft_new = int(np.round(self.n_fft * factor))
+ win_length_new = int(np.round(self.win_length * factor))
+ hop_length_new = int(np.round(self.hop_length * speed))
+ keyshift_key = str(keyshift) + "_" + str(audio.device)
+ if keyshift_key not in self.hann_window:
+ self.hann_window[keyshift_key] = torch.hann_window(win_length_new).to(
+ audio.device
+ )
+ if "privateuseone" in str(audio.device):
+ if not hasattr(self, "stft"):
+ self.stft = STFT(
+ filter_length=n_fft_new,
+ hop_length=hop_length_new,
+ win_length=win_length_new,
+ window="hann",
+ ).to(audio.device)
+ magnitude = self.stft.transform(audio)
+ else:
+ fft = torch.stft(
+ audio,
+ n_fft=n_fft_new,
+ hop_length=hop_length_new,
+ win_length=win_length_new,
+ window=self.hann_window[keyshift_key],
+ center=center,
+ return_complex=True,
+ )
+ magnitude = torch.sqrt(fft.real.pow(2) + fft.imag.pow(2))
+ if keyshift != 0:
+ size = self.n_fft // 2 + 1
+ resize = magnitude.size(1)
+ if resize < size:
+ magnitude = F.pad(magnitude, (0, 0, 0, size - resize))
+ magnitude = magnitude[:, :size, :] * self.win_length / win_length_new
+ mel_output = torch.matmul(self.mel_basis, magnitude)
+ if self.is_half == True:
+ mel_output = mel_output.half()
+ log_mel_spec = torch.log(torch.clamp(mel_output, min=self.clamp))
+ return log_mel_spec
+
+
+class RMVPE:
+ def __init__(self, model_path: str, is_half, device=None, use_jit=False):
+ self.resample_kernel = {}
+ self.resample_kernel = {}
+ self.is_half = is_half
+ if device is None:
+ device = "cuda:0" if torch.cuda.is_available() else "cpu"
+ self.device = device
+ self.mel_extractor = MelSpectrogram(
+ is_half, 128, 16000, 1024, 160, None, 30, 8000
+ ).to(device)
+ if "privateuseone" in str(device):
+ import onnxruntime as ort
+
+ ort_session = ort.InferenceSession(
+ "%s/rmvpe.onnx" % os.environ["rmvpe_root"],
+ providers=["DmlExecutionProvider"],
+ )
+ self.model = ort_session
+ else:
+ if str(self.device) == "cuda":
+ self.device = torch.device("cuda:0")
+
+ def get_default_model():
+ model = E2E(4, 1, (2, 2))
+ ckpt = torch.load(model_path, map_location="cpu")
+ model.load_state_dict(ckpt)
+ model.eval()
+ if is_half:
+ model = model.half()
+ else:
+ model = model.float()
+ return model
+
+ self.model = get_default_model()
+
+ self.model = self.model.to(device)
+ cents_mapping = 20 * np.arange(360) + 1997.3794084376191
+ self.cents_mapping = np.pad(cents_mapping, (4, 4)) # 368
+
+ def mel2hidden(self, mel):
+ with torch.no_grad():
+ n_frames = mel.shape[-1]
+ n_pad = 32 * ((n_frames - 1) // 32 + 1) - n_frames
+ if n_pad > 0:
+ mel = F.pad(mel, (0, n_pad), mode="constant")
+ if "privateuseone" in str(self.device):
+ onnx_input_name = self.model.get_inputs()[0].name
+ onnx_outputs_names = self.model.get_outputs()[0].name
+ hidden = self.model.run(
+ [onnx_outputs_names],
+ input_feed={onnx_input_name: mel.cpu().numpy()},
+ )[0]
+ else:
+ mel = mel.half() if self.is_half else mel.float()
+ hidden = self.model(mel)
+ return hidden[:, :n_frames]
+
+ def decode(self, hidden, thred=0.03):
+ cents_pred = self.to_local_average_cents(hidden, thred=thred)
+ f0 = 10 * (2 ** (cents_pred / 1200))
+ f0[f0 == 10] = 0
+ # f0 = np.array([10 * (2 ** (cent_pred / 1200)) if cent_pred else 0 for cent_pred in cents_pred])
+ return f0
+
+ def infer_from_audio(self, audio, thred=0.03):
+ # torch.cuda.synchronize()
+ # t0 = ttime()
+ if not torch.is_tensor(audio):
+ audio = torch.from_numpy(audio)
+ mel = self.mel_extractor(
+ audio.float().to(self.device).unsqueeze(0), center=True
+ )
+ # print(123123123,mel.device.type)
+ # torch.cuda.synchronize()
+ # t1 = ttime()
+ hidden = self.mel2hidden(mel)
+ # torch.cuda.synchronize()
+ # t2 = ttime()
+ # print(234234,hidden.device.type)
+ if "privateuseone" not in str(self.device):
+ hidden = hidden.squeeze(0).cpu().numpy()
+ else:
+ hidden = hidden[0]
+ if self.is_half == True:
+ hidden = hidden.astype("float32")
+
+ f0 = self.decode(hidden, thred=thred)
+ # torch.cuda.synchronize()
+ # t3 = ttime()
+ # print("hmvpe:%s\t%s\t%s\t%s"%(t1-t0,t2-t1,t3-t2,t3-t0))
+ return f0
+ def infer_from_audio_batch(self, audio, thred=0.03):
+ # torch.cuda.synchronize()
+ # t0 = ttime()
+ if not torch.is_tensor(audio):
+ audio = torch.from_numpy(audio)
+ mel = self.mel_extractor(
+ audio.float().to(self.device), center=True
+ )
+ # print(123123123,mel.device.type)
+ # torch.cuda.synchronize()
+ # t1 = ttime()
+ hidden = self.mel2hidden(mel)
+ # torch.cuda.synchronize()
+ # t2 = ttime()
+ # print(234234,hidden.device.type)
+ if "privateuseone" not in str(self.device):
+ hidden = hidden.cpu().numpy()
+ else:
+ pass
+ if self.is_half == True:
+ hidden = hidden.astype("float32")
+
+ f0s = []
+ for bib in range(hidden.shape[0]):
+ f0s.append(self.decode(hidden[bib], thred=thred))
+ f0s = np.stack(f0s)
+ f0s = torch.from_numpy(f0s).to(self.device)
+ # torch.cuda.synchronize()
+ # t3 = ttime()
+ # print("hmvpe:%s\t%s\t%s\t%s"%(t1-t0,t2-t1,t3-t2,t3-t0))
+ return f0s
+
+ def to_local_average_cents(self, salience, thred=0.05):
+ # t0 = ttime()
+ center = np.argmax(salience, axis=1) # 帧长#index
+ salience = np.pad(salience, ((0, 0), (4, 4))) # 帧长,368
+ # t1 = ttime()
+ center += 4
+ todo_salience = []
+ todo_cents_mapping = []
+ starts = center - 4
+ ends = center + 5
+ for idx in range(salience.shape[0]):
+ todo_salience.append(salience[:, starts[idx] : ends[idx]][idx])
+ todo_cents_mapping.append(self.cents_mapping[starts[idx] : ends[idx]])
+ # t2 = ttime()
+ todo_salience = np.array(todo_salience) # 帧长，9
+ todo_cents_mapping = np.array(todo_cents_mapping) # 帧长，9
+ product_sum = np.sum(todo_salience * todo_cents_mapping, 1)
+ weight_sum = np.sum(todo_salience, 1) # 帧长
+ devided = product_sum / weight_sum # 帧长
+ # t3 = ttime()
+ maxx = np.max(salience, axis=1) # 帧长
+ devided[maxx <= thred] = 0
+ # t4 = ttime()
+ # print("decode:%s\t%s\t%s\t%s" % (t1 - t0, t2 - t1, t3 - t2, t4 - t3))
+ return devided
diff --git a/modules/vocos/__init__.py b/modules/vocos/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..232ea57fa79422f5a99800d05ea7e85bae29b46e
--- /dev/null
+++ b/modules/vocos/__init__.py
@@ -0,0 +1,4 @@
+from .pretrained import Vocos
+
+
+__version__ = "0.1.0"
diff --git a/modules/vocos/heads.py b/modules/vocos/heads.py
new file mode 100644
index 0000000000000000000000000000000000000000..cb60c7ef0051db4d8b98f5ae94c7280eff62c26f
--- /dev/null
+++ b/modules/vocos/heads.py
@@ -0,0 +1,164 @@
+from typing import Optional
+
+import torch
+from torch import nn
+from torchaudio.functional.functional import _hz_to_mel, _mel_to_hz
+
+from .spectral_ops import IMDCT, ISTFT
+from .modules import symexp
+
+
+class FourierHead(nn.Module):
+ """Base class for inverse fourier modules."""
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ """
+ Args:
+ x (Tensor): Input tensor of shape (B, L, H), where B is the batch size,
+ L is the sequence length, and H denotes the model dimension.
+
+ Returns:
+ Tensor: Reconstructed time-domain audio signal of shape (B, T), where T is the length of the output signal.
+ """
+ raise NotImplementedError("Subclasses must implement the forward method.")
+
+
+class ISTFTHead(FourierHead):
+ """
+ ISTFT Head module for predicting STFT complex coefficients.
+
+ Args:
+ dim (int): Hidden dimension of the model.
+ n_fft (int): Size of Fourier transform.
+ hop_length (int): The distance between neighboring sliding window frames, which should align with
+ the resolution of the input features.
+ padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same".
+ """
+
+ def __init__(self, dim: int, n_fft: int, hop_length: int, padding: str = "same"):
+ super().__init__()
+ out_dim = n_fft + 2
+ self.out = torch.nn.Linear(dim, out_dim)
+ self.istft = ISTFT(n_fft=n_fft, hop_length=hop_length, win_length=n_fft, padding=padding)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ """
+ Forward pass of the ISTFTHead module.
+
+ Args:
+ x (Tensor): Input tensor of shape (B, L, H), where B is the batch size,
+ L is the sequence length, and H denotes the model dimension.
+
+ Returns:
+ Tensor: Reconstructed time-domain audio signal of shape (B, T), where T is the length of the output signal.
+ """
+ x = self.out(x).transpose(1, 2)
+ mag, p = x.chunk(2, dim=1)
+ mag = torch.exp(mag)
+ mag = torch.clip(mag, max=1e2) # safeguard to prevent excessively large magnitudes
+ # wrapping happens here. These two lines produce real and imaginary value
+ x = torch.cos(p)
+ y = torch.sin(p)
+ # recalculating phase here does not produce anything new
+ # only costs time
+ # phase = torch.atan2(y, x)
+ # S = mag * torch.exp(phase * 1j)
+ # better directly produce the complex value
+ S = mag * (x + 1j * y)
+ audio = self.istft(S)
+ return audio
+
+
+class IMDCTSymExpHead(FourierHead):
+ """
+ IMDCT Head module for predicting MDCT coefficients with symmetric exponential function
+
+ Args:
+ dim (int): Hidden dimension of the model.
+ mdct_frame_len (int): Length of the MDCT frame.
+ padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same".
+ sample_rate (int, optional): The sample rate of the audio. If provided, the last layer will be initialized
+ based on perceptual scaling. Defaults to None.
+ clip_audio (bool, optional): Whether to clip the audio output within the range of [-1.0, 1.0]. Defaults to False.
+ """
+
+ def __init__(
+ self,
+ dim: int,
+ mdct_frame_len: int,
+ padding: str = "same",
+ sample_rate: Optional[int] = None,
+ clip_audio: bool = False,
+ ):
+ super().__init__()
+ out_dim = mdct_frame_len // 2
+ self.out = nn.Linear(dim, out_dim)
+ self.imdct = IMDCT(frame_len=mdct_frame_len, padding=padding)
+ self.clip_audio = clip_audio
+
+ if sample_rate is not None:
+ # optionally init the last layer following mel-scale
+ m_max = _hz_to_mel(sample_rate // 2)
+ m_pts = torch.linspace(0, m_max, out_dim)
+ f_pts = _mel_to_hz(m_pts)
+ scale = 1 - (f_pts / f_pts.max())
+
+ with torch.no_grad():
+ self.out.weight.mul_(scale.view(-1, 1))
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ """
+ Forward pass of the IMDCTSymExpHead module.
+
+ Args:
+ x (Tensor): Input tensor of shape (B, L, H), where B is the batch size,
+ L is the sequence length, and H denotes the model dimension.
+
+ Returns:
+ Tensor: Reconstructed time-domain audio signal of shape (B, T), where T is the length of the output signal.
+ """
+ x = self.out(x)
+ x = symexp(x)
+ x = torch.clip(x, min=-1e2, max=1e2) # safeguard to prevent excessively large magnitudes
+ audio = self.imdct(x)
+ if self.clip_audio:
+ audio = torch.clip(x, min=-1.0, max=1.0)
+
+ return audio
+
+
+class IMDCTCosHead(FourierHead):
+ """
+ IMDCT Head module for predicting MDCT coefficients with parametrizing MDCT = exp(m) · cos(p)
+
+ Args:
+ dim (int): Hidden dimension of the model.
+ mdct_frame_len (int): Length of the MDCT frame.
+ padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same".
+ clip_audio (bool, optional): Whether to clip the audio output within the range of [-1.0, 1.0]. Defaults to False.
+ """
+
+ def __init__(self, dim: int, mdct_frame_len: int, padding: str = "same", clip_audio: bool = False):
+ super().__init__()
+ self.clip_audio = clip_audio
+ self.out = nn.Linear(dim, mdct_frame_len)
+ self.imdct = IMDCT(frame_len=mdct_frame_len, padding=padding)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ """
+ Forward pass of the IMDCTCosHead module.
+
+ Args:
+ x (Tensor): Input tensor of shape (B, L, H), where B is the batch size,
+ L is the sequence length, and H denotes the model dimension.
+
+ Returns:
+ Tensor: Reconstructed time-domain audio signal of shape (B, T), where T is the length of the output signal.
+ """
+ x = self.out(x)
+ m, p = x.chunk(2, dim=2)
+ m = torch.exp(m).clip(max=1e2) # safeguard to prevent excessively large magnitudes
+ audio = self.imdct(m * torch.cos(p))
+ if self.clip_audio:
+ audio = torch.clip(x, min=-1.0, max=1.0)
+ return audio
diff --git a/modules/vocos/helpers.py b/modules/vocos/helpers.py
new file mode 100644
index 0000000000000000000000000000000000000000..d1c769faeef696feb36a8cd7bb3ba6be8efafc4f
--- /dev/null
+++ b/modules/vocos/helpers.py
@@ -0,0 +1,71 @@
+import matplotlib
+import numpy as np
+import torch
+from matplotlib import pyplot as plt
+from pytorch_lightning import Callback
+
+matplotlib.use("Agg")
+
+
+def save_figure_to_numpy(fig: plt.Figure) -> np.ndarray:
+ """
+ Save a matplotlib figure to a numpy array.
+
+ Args:
+ fig (Figure): Matplotlib figure object.
+
+ Returns:
+ ndarray: Numpy array representing the figure.
+ """
+ data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep="")
+ data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
+ return data
+
+
+def plot_spectrogram_to_numpy(spectrogram: np.ndarray) -> np.ndarray:
+ """
+ Plot a spectrogram and convert it to a numpy array.
+
+ Args:
+ spectrogram (ndarray): Spectrogram data.
+
+ Returns:
+ ndarray: Numpy array representing the plotted spectrogram.
+ """
+ spectrogram = spectrogram.astype(np.float32)
+ fig, ax = plt.subplots(figsize=(12, 3))
+ im = ax.imshow(spectrogram, aspect="auto", origin="lower", interpolation="none")
+ plt.colorbar(im, ax=ax)
+ plt.xlabel("Frames")
+ plt.ylabel("Channels")
+ plt.tight_layout()
+
+ fig.canvas.draw()
+ data = save_figure_to_numpy(fig)
+ plt.close()
+ return data
+
+
+class GradNormCallback(Callback):
+ """
+ Callback to log the gradient norm.
+ """
+
+ def on_after_backward(self, trainer, model):
+ model.log("grad_norm", gradient_norm(model))
+
+
+def gradient_norm(model: torch.nn.Module, norm_type: float = 2.0) -> torch.Tensor:
+ """
+ Compute the gradient norm.
+
+ Args:
+ model (Module): PyTorch model.
+ norm_type (float, optional): Type of the norm. Defaults to 2.0.
+
+ Returns:
+ Tensor: Gradient norm.
+ """
+ grads = [p.grad for p in model.parameters() if p.grad is not None]
+ total_norm = torch.norm(torch.stack([torch.norm(g.detach(), norm_type) for g in grads]), norm_type)
+ return total_norm
diff --git a/modules/vocos/loss.py b/modules/vocos/loss.py
new file mode 100644
index 0000000000000000000000000000000000000000..91c15590646bf8743bbc428cf7c41caff8f19dac
--- /dev/null
+++ b/modules/vocos/loss.py
@@ -0,0 +1,114 @@
+from typing import List, Tuple
+
+import torch
+import torchaudio
+from torch import nn
+
+from vocos.modules import safe_log
+
+
+class MelSpecReconstructionLoss(nn.Module):
+ """
+ L1 distance between the mel-scaled magnitude spectrograms of the ground truth sample and the generated sample
+ """
+
+ def __init__(
+ self, sample_rate: int = 24000, n_fft: int = 1024, hop_length: int = 256, n_mels: int = 100,
+ ):
+ super().__init__()
+ self.mel_spec = torchaudio.transforms.MelSpectrogram(
+ sample_rate=sample_rate, n_fft=n_fft, hop_length=hop_length, n_mels=n_mels, center=True, power=1,
+ )
+
+ def forward(self, y_hat, y) -> torch.Tensor:
+ """
+ Args:
+ y_hat (Tensor): Predicted audio waveform.
+ y (Tensor): Ground truth audio waveform.
+
+ Returns:
+ Tensor: L1 loss between the mel-scaled magnitude spectrograms.
+ """
+ mel_hat = safe_log(self.mel_spec(y_hat))
+ mel = safe_log(self.mel_spec(y))
+
+ loss = torch.nn.functional.l1_loss(mel, mel_hat)
+
+ return loss
+
+
+class GeneratorLoss(nn.Module):
+ """
+ Generator Loss module. Calculates the loss for the generator based on discriminator outputs.
+ """
+
+ def forward(self, disc_outputs: List[torch.Tensor]) -> Tuple[torch.Tensor, List[torch.Tensor]]:
+ """
+ Args:
+ disc_outputs (List[Tensor]): List of discriminator outputs.
+
+ Returns:
+ Tuple[Tensor, List[Tensor]]: Tuple containing the total loss and a list of loss values from
+ the sub-discriminators
+ """
+ loss = torch.zeros(1, device=disc_outputs[0].device, dtype=disc_outputs[0].dtype)
+ gen_losses = []
+ for dg in disc_outputs:
+ l = torch.mean(torch.clamp(1 - dg, min=0))
+ gen_losses.append(l)
+ loss += l
+
+ return loss, gen_losses
+
+
+class DiscriminatorLoss(nn.Module):
+ """
+ Discriminator Loss module. Calculates the loss for the discriminator based on real and generated outputs.
+ """
+
+ def forward(
+ self, disc_real_outputs: List[torch.Tensor], disc_generated_outputs: List[torch.Tensor]
+ ) -> Tuple[torch.Tensor, List[torch.Tensor], List[torch.Tensor]]:
+ """
+ Args:
+ disc_real_outputs (List[Tensor]): List of discriminator outputs for real samples.
+ disc_generated_outputs (List[Tensor]): List of discriminator outputs for generated samples.
+
+ Returns:
+ Tuple[Tensor, List[Tensor], List[Tensor]]: A tuple containing the total loss, a list of loss values from
+ the sub-discriminators for real outputs, and a list of
+ loss values for generated outputs.
+ """
+ loss = torch.zeros(1, device=disc_real_outputs[0].device, dtype=disc_real_outputs[0].dtype)
+ r_losses = []
+ g_losses = []
+ for dr, dg in zip(disc_real_outputs, disc_generated_outputs):
+ r_loss = torch.mean(torch.clamp(1 - dr, min=0))
+ g_loss = torch.mean(torch.clamp(1 + dg, min=0))
+ loss += r_loss + g_loss
+ r_losses.append(r_loss)
+ g_losses.append(g_loss)
+
+ return loss, r_losses, g_losses
+
+
+class FeatureMatchingLoss(nn.Module):
+ """
+ Feature Matching Loss module. Calculates the feature matching loss between feature maps of the sub-discriminators.
+ """
+
+ def forward(self, fmap_r: List[List[torch.Tensor]], fmap_g: List[List[torch.Tensor]]) -> torch.Tensor:
+ """
+ Args:
+ fmap_r (List[List[Tensor]]): List of feature maps from real samples.
+ fmap_g (List[List[Tensor]]): List of feature maps from generated samples.
+
+ Returns:
+ Tensor: The calculated feature matching loss.
+ """
+ loss = torch.zeros(1, device=fmap_r[0][0].device, dtype=fmap_r[0][0].dtype)
+ for dr, dg in zip(fmap_r, fmap_g):
+ for rl, gl in zip(dr, dg):
+ loss += torch.mean(torch.abs(rl - gl))
+
+ return loss
diff --git a/modules/vocos/models.py b/modules/vocos/models.py
new file mode 100644
index 0000000000000000000000000000000000000000..7e5f4922aeecad7f110a8191191043b2ec9baabb
--- /dev/null
+++ b/modules/vocos/models.py
@@ -0,0 +1,118 @@
+from typing import Optional
+
+import torch
+from torch import nn
+from torch.nn.utils import weight_norm
+
+from .modules import ConvNeXtBlock, ResBlock1, AdaLayerNorm
+
+
+class Backbone(nn.Module):
+ """Base class for the generator's backbone. It preserves the same temporal resolution across all layers."""
+
+ def forward(self, x: torch.Tensor, **kwargs) -> torch.Tensor:
+ """
+ Args:
+ x (Tensor): Input tensor of shape (B, C, L), where B is the batch size,
+ C denotes output features, and L is the sequence length.
+
+ Returns:
+ Tensor: Output of shape (B, L, H), where B is the batch size, L is the sequence length,
+ and H denotes the model dimension.
+ """
+ raise NotImplementedError("Subclasses must implement the forward method.")
+
+
+class VocosBackbone(Backbone):
+ """
+ Vocos backbone module built with ConvNeXt blocks. Supports additional conditioning with Adaptive Layer Normalization
+
+ Args:
+ input_channels (int): Number of input features channels.
+ dim (int): Hidden dimension of the model.
+ intermediate_dim (int): Intermediate dimension used in ConvNeXtBlock.
+ num_layers (int): Number of ConvNeXtBlock layers.
+ layer_scale_init_value (float, optional): Initial value for layer scaling. Defaults to `1 / num_layers`.
+ adanorm_num_embeddings (int, optional): Number of embeddings for AdaLayerNorm.
+ None means non-conditional model. Defaults to None.
+ """
+
+ def __init__(
+ self,
+ input_channels: int,
+ dim: int,
+ intermediate_dim: int,
+ num_layers: int,
+ layer_scale_init_value: Optional[float] = None,
+ adanorm_num_embeddings: Optional[int] = None,
+ ):
+ super().__init__()
+ self.input_channels = input_channels
+ self.embed = nn.Conv1d(input_channels, dim, kernel_size=7, padding=3)
+ self.adanorm = adanorm_num_embeddings is not None
+ if adanorm_num_embeddings:
+ self.norm = AdaLayerNorm(adanorm_num_embeddings, dim, eps=1e-6)
+ else:
+ self.norm = nn.LayerNorm(dim, eps=1e-6)
+ layer_scale_init_value = layer_scale_init_value or 1 / num_layers
+ self.convnext = nn.ModuleList(
+ [
+ ConvNeXtBlock(
+ dim=dim,
+ intermediate_dim=intermediate_dim,
+ layer_scale_init_value=layer_scale_init_value,
+ adanorm_num_embeddings=adanorm_num_embeddings,
+ )
+ for _ in range(num_layers)
+ ]
+ )
+ self.final_layer_norm = nn.LayerNorm(dim, eps=1e-6)
+ self.apply(self._init_weights)
+
+ def _init_weights(self, m):
+ if isinstance(m, (nn.Conv1d, nn.Linear)):
+ nn.init.trunc_normal_(m.weight, std=0.02)
+ nn.init.constant_(m.bias, 0)
+
+ def forward(self, x: torch.Tensor, **kwargs) -> torch.Tensor:
+ bandwidth_id = kwargs.get('bandwidth_id', None)
+ x = self.embed(x)
+ if self.adanorm:
+ assert bandwidth_id is not None
+ x = self.norm(x.transpose(1, 2), cond_embedding_id=bandwidth_id)
+ else:
+ x = self.norm(x.transpose(1, 2))
+ x = x.transpose(1, 2)
+ for conv_block in self.convnext:
+ x = conv_block(x, cond_embedding_id=bandwidth_id)
+ x = self.final_layer_norm(x.transpose(1, 2))
+ return x
+
+
+class VocosResNetBackbone(Backbone):
+ """
+ Vocos backbone module built with ResBlocks.
+
+ Args:
+ input_channels (int): Number of input features channels.
+ dim (int): Hidden dimension of the model.
+ num_blocks (int): Number of ResBlock1 blocks.
+ layer_scale_init_value (float, optional): Initial value for layer scaling. Defaults to None.
+ """
+
+ def __init__(
+ self, input_channels, dim, num_blocks, layer_scale_init_value=None,
+ ):
+ super().__init__()
+ self.input_channels = input_channels
+ self.embed = weight_norm(nn.Conv1d(input_channels, dim, kernel_size=3, padding=1))
+ layer_scale_init_value = layer_scale_init_value or 1 / num_blocks / 3
+ self.resnet = nn.Sequential(
+ *[ResBlock1(dim=dim, layer_scale_init_value=layer_scale_init_value) for _ in range(num_blocks)]
+ )
+
+ def forward(self, x: torch.Tensor, **kwargs) -> torch.Tensor:
+ x = self.embed(x)
+ x = self.resnet(x)
+ x = x.transpose(1, 2)
+ return x
diff --git a/modules/vocos/modules.py b/modules/vocos/modules.py
new file mode 100644
index 0000000000000000000000000000000000000000..2797667528bfa7b4453a6fc1d477f1fb065dd855
--- /dev/null
+++ b/modules/vocos/modules.py
@@ -0,0 +1,213 @@
+from typing import Optional, Tuple
+
+import torch
+from torch import nn
+from torch.nn.utils import weight_norm, remove_weight_norm
+
+
+class ConvNeXtBlock(nn.Module):
+ """ConvNeXt Block adapted from https://github.com/facebookresearch/ConvNeXt to 1D audio signal.
+
+ Args:
+ dim (int): Number of input channels.
+ intermediate_dim (int): Dimensionality of the intermediate layer.
+ layer_scale_init_value (float, optional): Initial value for the layer scale. None means no scaling.
+ Defaults to None.
+ adanorm_num_embeddings (int, optional): Number of embeddings for AdaLayerNorm.
+ None means non-conditional LayerNorm. Defaults to None.
+ """
+
+ def __init__(
+ self,
+ dim: int,
+ intermediate_dim: int,
+ layer_scale_init_value: float,
+ adanorm_num_embeddings: Optional[int] = None,
+ ):
+ super().__init__()
+ self.dwconv = nn.Conv1d(dim, dim, kernel_size=7, padding=3, groups=dim) # depthwise conv
+ self.adanorm = adanorm_num_embeddings is not None
+ if adanorm_num_embeddings:
+ self.norm = AdaLayerNorm(adanorm_num_embeddings, dim, eps=1e-6)
+ else:
+ self.norm = nn.LayerNorm(dim, eps=1e-6)
+ self.pwconv1 = nn.Linear(dim, intermediate_dim) # pointwise/1x1 convs, implemented with linear layers
+ self.act = nn.GELU()
+ self.pwconv2 = nn.Linear(intermediate_dim, dim)
+ self.gamma = (
+ nn.Parameter(layer_scale_init_value * torch.ones(dim), requires_grad=True)
+ if layer_scale_init_value > 0
+ else None
+ )
+
+ def forward(self, x: torch.Tensor, cond_embedding_id: Optional[torch.Tensor] = None) -> torch.Tensor:
+ residual = x
+ x = self.dwconv(x)
+ x = x.transpose(1, 2) # (B, C, T) -> (B, T, C)
+ if self.adanorm:
+ assert cond_embedding_id is not None
+ x = self.norm(x, cond_embedding_id)
+ else:
+ x = self.norm(x)
+ x = self.pwconv1(x)
+ x = self.act(x)
+ x = self.pwconv2(x)
+ if self.gamma is not None:
+ x = self.gamma * x
+ x = x.transpose(1, 2) # (B, T, C) -> (B, C, T)
+
+ x = residual + x
+ return x
+
+
+class AdaLayerNorm(nn.Module):
+ """
+ Adaptive Layer Normalization module with learnable embeddings per `num_embeddings` classes
+
+ Args:
+ num_embeddings (int): Number of embeddings.
+ embedding_dim (int): Dimension of the embeddings.
+ """
+
+ def __init__(self, num_embeddings: int, embedding_dim: int, eps: float = 1e-6):
+ super().__init__()
+ self.eps = eps
+ self.dim = embedding_dim
+ self.scale = nn.Embedding(num_embeddings=num_embeddings, embedding_dim=embedding_dim)
+ self.shift = nn.Embedding(num_embeddings=num_embeddings, embedding_dim=embedding_dim)
+ torch.nn.init.ones_(self.scale.weight)
+ torch.nn.init.zeros_(self.shift.weight)
+
+ def forward(self, x: torch.Tensor, cond_embedding_id: torch.Tensor) -> torch.Tensor:
+ scale = self.scale(cond_embedding_id)
+ shift = self.shift(cond_embedding_id)
+ x = nn.functional.layer_norm(x, (self.dim,), eps=self.eps)
+ x = x * scale + shift
+ return x
+
+
+class ResBlock1(nn.Module):
+ """
+ ResBlock adapted from HiFi-GAN V1 (https://github.com/jik876/hifi-gan) with dilated 1D convolutions,
+ but without upsampling layers.
+
+ Args:
+ dim (int): Number of input channels.
+ kernel_size (int, optional): Size of the convolutional kernel. Defaults to 3.
+ dilation (tuple[int], optional): Dilation factors for the dilated convolutions.
+ Defaults to (1, 3, 5).
+ lrelu_slope (float, optional): Negative slope of the LeakyReLU activation function.
+ Defaults to 0.1.
+ layer_scale_init_value (float, optional): Initial value for the layer scale. None means no scaling.
+ Defaults to None.
+ """
+
+ def __init__(
+ self,
+ dim: int,
+ kernel_size: int = 3,
+ dilation: Tuple[int, int, int] = (1, 3, 5),
+ lrelu_slope: float = 0.1,
+ layer_scale_init_value: Optional[float] = None,
+ ):
+ super().__init__()
+ self.lrelu_slope = lrelu_slope
+ self.convs1 = nn.ModuleList(
+ [
+ weight_norm(
+ nn.Conv1d(
+ dim,
+ dim,
+ kernel_size,
+ 1,
+ dilation=dilation[0],
+ padding=self.get_padding(kernel_size, dilation[0]),
+ )
+ ),
+ weight_norm(
+ nn.Conv1d(
+ dim,
+ dim,
+ kernel_size,
+ 1,
+ dilation=dilation[1],
+ padding=self.get_padding(kernel_size, dilation[1]),
+ )
+ ),
+ weight_norm(
+ nn.Conv1d(
+ dim,
+ dim,
+ kernel_size,
+ 1,
+ dilation=dilation[2],
+ padding=self.get_padding(kernel_size, dilation[2]),
+ )
+ ),
+ ]
+ )
+
+ self.convs2 = nn.ModuleList(
+ [
+ weight_norm(nn.Conv1d(dim, dim, kernel_size, 1, dilation=1, padding=self.get_padding(kernel_size, 1))),
+ weight_norm(nn.Conv1d(dim, dim, kernel_size, 1, dilation=1, padding=self.get_padding(kernel_size, 1))),
+ weight_norm(nn.Conv1d(dim, dim, kernel_size, 1, dilation=1, padding=self.get_padding(kernel_size, 1))),
+ ]
+ )
+
+ self.gamma = nn.ParameterList(
+ [
+ nn.Parameter(layer_scale_init_value * torch.ones(dim, 1), requires_grad=True)
+ if layer_scale_init_value is not None
+ else None,
+ nn.Parameter(layer_scale_init_value * torch.ones(dim, 1), requires_grad=True)
+ if layer_scale_init_value is not None
+ else None,
+ nn.Parameter(layer_scale_init_value * torch.ones(dim, 1), requires_grad=True)
+ if layer_scale_init_value is not None
+ else None,
+ ]
+ )
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ for c1, c2, gamma in zip(self.convs1, self.convs2, self.gamma):
+ xt = torch.nn.functional.leaky_relu(x, negative_slope=self.lrelu_slope)
+ xt = c1(xt)
+ xt = torch.nn.functional.leaky_relu(xt, negative_slope=self.lrelu_slope)
+ xt = c2(xt)
+ if gamma is not None:
+ xt = gamma * xt
+ x = xt + x
+ return x
+
+ def remove_weight_norm(self):
+ for l in self.convs1:
+ remove_weight_norm(l)
+ for l in self.convs2:
+ remove_weight_norm(l)
+
+ @staticmethod
+ def get_padding(kernel_size: int, dilation: int = 1) -> int:
+ return int((kernel_size * dilation - dilation) / 2)
+
+
+def safe_log(x: torch.Tensor, clip_val: float = 1e-7) -> torch.Tensor:
+ """
+ Computes the element-wise logarithm of the input tensor with clipping to avoid near-zero values.
+
+ Args:
+ x (Tensor): Input tensor.
+ clip_val (float, optional): Minimum value to clip the input tensor. Defaults to 1e-7.
+
+ Returns:
+ Tensor: Element-wise logarithm of the input tensor with clipping applied.
+ """
+ return torch.log(torch.clip(x, min=clip_val))
+
+
+def symlog(x: torch.Tensor) -> torch.Tensor:
+ return torch.sign(x) * torch.log1p(x.abs())
+
+
+def symexp(x: torch.Tensor) -> torch.Tensor:
+ return torch.sign(x) * (torch.exp(x.abs()) - 1)
diff --git a/modules/vocos/pretrained.py b/modules/vocos/pretrained.py
new file mode 100644
index 0000000000000000000000000000000000000000..2f5292acb7c565ffbf14ae681653846fa9ca96b6
--- /dev/null
+++ b/modules/vocos/pretrained.py
@@ -0,0 +1,51 @@
+from __future__ import annotations
+
+from typing import Any, Dict, Tuple, Union, Optional
+
+import torch
+import yaml
+from torch import nn
+from .heads import ISTFTHead
+from .models import VocosBackbone
+
+
+class Vocos(nn.Module):
+ """
+ The Vocos class represents a Fourier-based neural vocoder for audio synthesis.
+ This class is primarily designed for inference, with support for loading from pretrained
+ model checkpoints. It consists of three main components: a feature extractor,
+ a backbone, and a head.
+ """
+
+ def __init__(
+ self, args,
+ ):
+ super().__init__()
+ self.backbone = VocosBackbone(
+ input_channels=args.vocos.backbone.input_channels,
+ dim=args.vocos.backbone.dim,
+ intermediate_dim=args.vocos.backbone.intermediate_dim,
+ num_layers=args.vocos.backbone.num_layers,
+ )
+ self.head = ISTFTHead(
+ dim=args.vocos.head.dim,
+ n_fft=args.vocos.head.n_fft,
+ hop_length=args.vocos.head.hop_length,
+ padding=args.vocos.head.padding,
+ )
+
+ def forward(self, features_input: torch.Tensor, **kwargs: Any) -> torch.Tensor:
+ """
+ Method to decode audio waveform from already calculated features. The features input is passed through
+ the backbone and the head to reconstruct the audio output.
+
+ Args:
+ features_input (Tensor): The input tensor of features of shape (B, C, L), where B is the batch size,
+ C denotes the feature dimension, and L is the sequence length.
+
+ Returns:
+ Tensor: The output tensor representing the reconstructed audio waveform of shape (B, T).
+ """
+ x = self.backbone(features_input, **kwargs)
+ audio_output = self.head(x)
+ return audio_output
diff --git a/modules/vocos/spectral_ops.py b/modules/vocos/spectral_ops.py
new file mode 100644
index 0000000000000000000000000000000000000000..b130ed342cc638f260f15134b671e325f2cdf454
--- /dev/null
+++ b/modules/vocos/spectral_ops.py
@@ -0,0 +1,192 @@
+import numpy as np
+import scipy
+import torch
+from torch import nn, view_as_real, view_as_complex
+
+
+class ISTFT(nn.Module):
+ """
+ Custom implementation of ISTFT since torch.istft doesn't allow custom padding (other than `center=True`) with
+ windowing. This is because the NOLA (Nonzero Overlap Add) check fails at the edges.
+ See issue: https://github.com/pytorch/pytorch/issues/62323
+ Specifically, in the context of neural vocoding we are interested in "same" padding analogous to CNNs.
+ The NOLA constraint is met as we trim padded samples anyway.
+
+ Args:
+ n_fft (int): Size of Fourier transform.
+ hop_length (int): The distance between neighboring sliding window frames.
+ win_length (int): The size of window frame and STFT filter.
+ padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same".
+ """
+
+ def __init__(self, n_fft: int, hop_length: int, win_length: int, padding: str = "same"):
+ super().__init__()
+ if padding not in ["center", "same"]:
+ raise ValueError("Padding must be 'center' or 'same'.")
+ self.padding = padding
+ self.n_fft = n_fft
+ self.hop_length = hop_length
+ self.win_length = win_length
+ window = torch.hann_window(win_length)
+ self.register_buffer("window", window)
+
+ def forward(self, spec: torch.Tensor) -> torch.Tensor:
+ """
+ Compute the Inverse Short Time Fourier Transform (ISTFT) of a complex spectrogram.
+
+ Args:
+ spec (Tensor): Input complex spectrogram of shape (B, N, T), where B is the batch size,
+ N is the number of frequency bins, and T is the number of time frames.
+
+ Returns:
+ Tensor: Reconstructed time-domain signal of shape (B, L), where L is the length of the output signal.
+ """
+ if self.padding == "center":
+ # Fallback to pytorch native implementation
+ return torch.istft(spec, self.n_fft, self.hop_length, self.win_length, self.window, center=True)
+ elif self.padding == "same":
+ pad = (self.win_length - self.hop_length) // 2
+ else:
+ raise ValueError("Padding must be 'center' or 'same'.")
+
+ assert spec.dim() == 3, "Expected a 3D tensor as input"
+ B, N, T = spec.shape
+
+ # Inverse FFT
+ ifft = torch.fft.irfft(spec, self.n_fft, dim=1, norm="backward")
+ ifft = ifft * self.window[None, :, None]
+
+ # Overlap and Add
+ output_size = (T - 1) * self.hop_length + self.win_length
+ y = torch.nn.functional.fold(
+ ifft, output_size=(1, output_size), kernel_size=(1, self.win_length), stride=(1, self.hop_length),
+ )[:, 0, 0, pad:-pad]
+
+ # Window envelope
+ window_sq = self.window.square().expand(1, T, -1).transpose(1, 2)
+ window_envelope = torch.nn.functional.fold(
+ window_sq, output_size=(1, output_size), kernel_size=(1, self.win_length), stride=(1, self.hop_length),
+ ).squeeze()[pad:-pad]
+
+ # Normalize
+ assert (window_envelope > 1e-11).all()
+ y = y / window_envelope
+
+ return y
+
+
+class MDCT(nn.Module):
+ """
+ Modified Discrete Cosine Transform (MDCT) module.
+
+ Args:
+ frame_len (int): Length of the MDCT frame.
+ padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same".
+ """
+
+ def __init__(self, frame_len: int, padding: str = "same"):
+ super().__init__()
+ if padding not in ["center", "same"]:
+ raise ValueError("Padding must be 'center' or 'same'.")
+ self.padding = padding
+ self.frame_len = frame_len
+ N = frame_len // 2
+ n0 = (N + 1) / 2
+ window = torch.from_numpy(scipy.signal.cosine(frame_len)).float()
+ self.register_buffer("window", window)
+
+ pre_twiddle = torch.exp(-1j * torch.pi * torch.arange(frame_len) / frame_len)
+ post_twiddle = torch.exp(-1j * torch.pi * n0 * (torch.arange(N) + 0.5) / N)
+ # view_as_real: NCCL Backend does not support ComplexFloat data type
+ # https://github.com/pytorch/pytorch/issues/71613
+ self.register_buffer("pre_twiddle", view_as_real(pre_twiddle))
+ self.register_buffer("post_twiddle", view_as_real(post_twiddle))
+
+ def forward(self, audio: torch.Tensor) -> torch.Tensor:
+ """
+ Apply the Modified Discrete Cosine Transform (MDCT) to the input audio.
+
+ Args:
+ audio (Tensor): Input audio waveform of shape (B, T), where B is the batch size
+ and T is the length of the audio.
+
+ Returns:
+ Tensor: MDCT coefficients of shape (B, L, N), where L is the number of output frames
+ and N is the number of frequency bins.
+ """
+ if self.padding == "center":
+ audio = torch.nn.functional.pad(audio, (self.frame_len // 2, self.frame_len // 2))
+ elif self.padding == "same":
+ # hop_length is 1/2 frame_len
+ audio = torch.nn.functional.pad(audio, (self.frame_len // 4, self.frame_len // 4))
+ else:
+ raise ValueError("Padding must be 'center' or 'same'.")
+
+ x = audio.unfold(-1, self.frame_len, self.frame_len // 2)
+ N = self.frame_len // 2
+ x = x * self.window.expand(x.shape)
+ X = torch.fft.fft(x * view_as_complex(self.pre_twiddle).expand(x.shape), dim=-1)[..., :N]
+ res = X * view_as_complex(self.post_twiddle).expand(X.shape) * np.sqrt(1 / N)
+ return torch.real(res) * np.sqrt(2)
+
+
+class IMDCT(nn.Module):
+ """
+ Inverse Modified Discrete Cosine Transform (IMDCT) module.
+
+ Args:
+ frame_len (int): Length of the MDCT frame.
+ padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same".
+ """
+
+ def __init__(self, frame_len: int, padding: str = "same"):
+ super().__init__()
+ if padding not in ["center", "same"]:
+ raise ValueError("Padding must be 'center' or 'same'.")
+ self.padding = padding
+ self.frame_len = frame_len
+ N = frame_len // 2
+ n0 = (N + 1) / 2
+ window = torch.from_numpy(scipy.signal.cosine(frame_len)).float()
+ self.register_buffer("window", window)
+
+ pre_twiddle = torch.exp(1j * torch.pi * n0 * torch.arange(N * 2) / N)
+ post_twiddle = torch.exp(1j * torch.pi * (torch.arange(N * 2) + n0) / (N * 2))
+ self.register_buffer("pre_twiddle", view_as_real(pre_twiddle))
+ self.register_buffer("post_twiddle", view_as_real(post_twiddle))
+
+ def forward(self, X: torch.Tensor) -> torch.Tensor:
+ """
+ Apply the Inverse Modified Discrete Cosine Transform (IMDCT) to the input MDCT coefficients.
+
+ Args:
+ X (Tensor): Input MDCT coefficients of shape (B, L, N), where B is the batch size,
+ L is the number of frames, and N is the number of frequency bins.
+
+ Returns:
+ Tensor: Reconstructed audio waveform of shape (B, T), where T is the length of the audio.
+ """
+ B, L, N = X.shape
+ Y = torch.zeros((B, L, N * 2), dtype=X.dtype, device=X.device)
+ Y[..., :N] = X
+ Y[..., N:] = -1 * torch.conj(torch.flip(X, dims=(-1,)))
+ y = torch.fft.ifft(Y * view_as_complex(self.pre_twiddle).expand(Y.shape), dim=-1)
+ y = torch.real(y * view_as_complex(self.post_twiddle).expand(y.shape)) * np.sqrt(N) * np.sqrt(2)
+ result = y * self.window.expand(y.shape)
+ output_size = (1, (L + 1) * N)
+ audio = torch.nn.functional.fold(
+ result.transpose(1, 2),
+ output_size=output_size,
+ kernel_size=(1, self.frame_len),
+ stride=(1, self.frame_len // 2),
+ )[:, 0, 0, :]
+
+ if self.padding == "center":
+ pad = self.frame_len // 2
+ elif self.padding == "same":
+ pad = self.frame_len // 4
+ else:
+ raise ValueError("Padding must be 'center' or 'same'.")
+
+ audio = audio[:, pad:-pad]
+ return audio
diff --git a/modules/wavenet.py b/modules/wavenet.py
new file mode 100644
index 0000000000000000000000000000000000000000..6c81ffbe58c99fdfb30da9567d609cf193350d8d
--- /dev/null
+++ b/modules/wavenet.py
@@ -0,0 +1,174 @@
+import math
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from modules.encodec import SConv1d
+
+from . import commons
+LRELU_SLOPE = 0.1
+
+class LayerNorm(nn.Module):
+ def __init__(self, channels, eps=1e-5):
+ super().__init__()
+ self.channels = channels
+ self.eps = eps
+
+ self.gamma = nn.Parameter(torch.ones(channels))
+ self.beta = nn.Parameter(torch.zeros(channels))
+
+ def forward(self, x):
+ x = x.transpose(1, -1)
+ x = F.layer_norm(x, (self.channels,), self.gamma, self.beta, self.eps)
+ return x.transpose(1, -1)
+
+
+class ConvReluNorm(nn.Module):
+ def __init__(self, in_channels, hidden_channels, out_channels, kernel_size, n_layers, p_dropout):
+ super().__init__()
+ self.in_channels = in_channels
+ self.hidden_channels = hidden_channels
+ self.out_channels = out_channels
+ self.kernel_size = kernel_size
+ self.n_layers = n_layers
+ self.p_dropout = p_dropout
+ assert n_layers > 1, "Number of layers should be larger than 0."
+
+ self.conv_layers = nn.ModuleList()
+ self.norm_layers = nn.ModuleList()
+ self.conv_layers.append(nn.Conv1d(in_channels, hidden_channels, kernel_size, padding=kernel_size // 2))
+ self.norm_layers.append(LayerNorm(hidden_channels))
+ self.relu_drop = nn.Sequential(
+ nn.ReLU(),
+ nn.Dropout(p_dropout))
+ for _ in range(n_layers - 1):
+ self.conv_layers.append(nn.Conv1d(hidden_channels, hidden_channels, kernel_size, padding=kernel_size // 2))
+ self.norm_layers.append(LayerNorm(hidden_channels))
+ self.proj = nn.Conv1d(hidden_channels, out_channels, 1)
+ self.proj.weight.data.zero_()
+ self.proj.bias.data.zero_()
+
+ def forward(self, x, x_mask):
+ x_org = x
+ for i in range(self.n_layers):
+ x = self.conv_layers[i](x * x_mask)
+ x = self.norm_layers[i](x)
+ x = self.relu_drop(x)
+ x = x_org + self.proj(x)
+ return x * x_mask
+
+
+class DDSConv(nn.Module):
+ """
+ Dialted and Depth-Separable Convolution
+ """
+
+ def __init__(self, channels, kernel_size, n_layers, p_dropout=0.):
+ super().__init__()
+ self.channels = channels
+ self.kernel_size = kernel_size
+ self.n_layers = n_layers
+ self.p_dropout = p_dropout
+
+ self.drop = nn.Dropout(p_dropout)
+ self.convs_sep = nn.ModuleList()
+ self.convs_1x1 = nn.ModuleList()
+ self.norms_1 = nn.ModuleList()
+ self.norms_2 = nn.ModuleList()
+ for i in range(n_layers):
+ dilation = kernel_size ** i
+ padding = (kernel_size * dilation - dilation) // 2
+ self.convs_sep.append(nn.Conv1d(channels, channels, kernel_size,
+ groups=channels, dilation=dilation, padding=padding
+ ))
+ self.convs_1x1.append(nn.Conv1d(channels, channels, 1))
+ self.norms_1.append(LayerNorm(channels))
+ self.norms_2.append(LayerNorm(channels))
+
+ def forward(self, x, x_mask, g=None):
+ if g is not None:
+ x = x + g
+ for i in range(self.n_layers):
+ y = self.convs_sep[i](x * x_mask)
+ y = self.norms_1[i](y)
+ y = F.gelu(y)
+ y = self.convs_1x1[i](y)
+ y = self.norms_2[i](y)
+ y = F.gelu(y)
+ y = self.drop(y)
+ x = x + y
+ return x * x_mask
+
+
+class WN(torch.nn.Module):
+ def __init__(self, hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=0, p_dropout=0, causal=False):
+ super(WN, self).__init__()
+ conv1d_type = SConv1d
+ assert (kernel_size % 2 == 1)
+ self.hidden_channels = hidden_channels
+ self.kernel_size = kernel_size,
+ self.dilation_rate = dilation_rate
+ self.n_layers = n_layers
+ self.gin_channels = gin_channels
+ self.p_dropout = p_dropout
+
+ self.in_layers = torch.nn.ModuleList()
+ self.res_skip_layers = torch.nn.ModuleList()
+ self.drop = nn.Dropout(p_dropout)
+
+ if gin_channels != 0:
+ self.cond_layer = conv1d_type(gin_channels, 2 * hidden_channels * n_layers, 1, norm='weight_norm')
+
+ for i in range(n_layers):
+ dilation = dilation_rate ** i
+ padding = int((kernel_size * dilation - dilation) / 2)
+ in_layer = conv1d_type(hidden_channels, 2 * hidden_channels, kernel_size, dilation=dilation,
+ padding=padding, norm='weight_norm', causal=causal)
+ self.in_layers.append(in_layer)
+
+ # last one is not necessary
+ if i < n_layers - 1:
+ res_skip_channels = 2 * hidden_channels
+ else:
+ res_skip_channels = hidden_channels
+
+ res_skip_layer = conv1d_type(hidden_channels, res_skip_channels, 1, norm='weight_norm', causal=causal)
+ self.res_skip_layers.append(res_skip_layer)
+
+ def forward(self, x, x_mask, g=None, **kwargs):
+ output = torch.zeros_like(x)
+ n_channels_tensor = torch.IntTensor([self.hidden_channels])
+
+ if g is not None:
+ g = self.cond_layer(g)
+
+ for i in range(self.n_layers):
+ x_in = self.in_layers[i](x)
+ if g is not None:
+ cond_offset = i * 2 * self.hidden_channels
+ g_l = g[:, cond_offset:cond_offset + 2 * self.hidden_channels, :]
+ else:
+ g_l = torch.zeros_like(x_in)
+
+ acts = commons.fused_add_tanh_sigmoid_multiply(
+ x_in,
+ g_l,
+ n_channels_tensor)
+ acts = self.drop(acts)
+
+ res_skip_acts = self.res_skip_layers[i](acts)
+ if i < self.n_layers - 1:
+ res_acts = res_skip_acts[:, :self.hidden_channels, :]
+ x = (x + res_acts) * x_mask
+ output = output + res_skip_acts[:, self.hidden_channels:, :]
+ else:
+ output = output + res_skip_acts
+ return output * x_mask
+
+ def remove_weight_norm(self):
+ if self.gin_channels != 0:
+ torch.nn.utils.remove_weight_norm(self.cond_layer)
+ for l in self.in_layers:
+ torch.nn.utils.remove_weight_norm(l)
+ for l in self.res_skip_layers:
+ torch.nn.utils.remove_weight_norm(l)
\ No newline at end of file
diff --git a/optimizers.py b/optimizers.py
new file mode 100644
index 0000000000000000000000000000000000000000..47b81fa30c4d86266c166780a632e46e5c877649
--- /dev/null
+++ b/optimizers.py
@@ -0,0 +1,96 @@
+#coding:utf-8
+import os, sys
+import os.path as osp
+import numpy as np
+import torch
+from torch import nn
+from torch.optim import Optimizer
+from functools import reduce
+from torch.optim import AdamW
+
+class MultiOptimizer:
+ def __init__(self, optimizers={}, schedulers={}):
+ self.optimizers = optimizers
+ self.schedulers = schedulers
+ self.keys = list(optimizers.keys())
+ self.param_groups = reduce(lambda x,y: x+y, [v.param_groups for v in self.optimizers.values()])
+
+ def state_dict(self):
+ state_dicts = [(key, self.optimizers[key].state_dict())\
+ for key in self.keys]
+ return state_dicts
+
+ def scheduler_state_dict(self):
+ state_dicts = [(key, self.schedulers[key].state_dict())\
+ for key in self.keys]
+ return state_dicts
+
+ def load_state_dict(self, state_dict):
+ for key, val in state_dict:
+ try:
+ self.optimizers[key].load_state_dict(val)
+ except:
+ print("Unloaded %s" % key)
+
+ def load_scheduler_state_dict(self, state_dict):
+ for key, val in state_dict:
+ try:
+ self.schedulers[key].load_state_dict(val)
+ except:
+ print("Unloaded %s" % key)
+
+ def step(self, key=None, scaler=None):
+ keys = [key] if key is not None else self.keys
+ _ = [self._step(key, scaler) for key in keys]
+
+ def _step(self, key, scaler=None):
+ if scaler is not None:
+ scaler.step(self.optimizers[key])
+ scaler.update()
+ else:
+ self.optimizers[key].step()
+
+ def zero_grad(self, key=None):
+ if key is not None:
+ self.optimizers[key].zero_grad()
+ else:
+ _ = [self.optimizers[key].zero_grad() for key in self.keys]
+
+ def scheduler(self, *args, key=None):
+ if key is not None:
+ self.schedulers[key].step(*args)
+ else:
+ _ = [self.schedulers[key].step_batch(*args) for key in self.keys]
+
+def define_scheduler(optimizer, params):
+ scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=params['gamma'])
+
+ return scheduler
+
+def build_optimizer(model_dict, lr, type='AdamW'):
+ optim = {}
+ for key, model in model_dict.items():
+ model_parameters = model.parameters()
+ parameters_names = []
+ parameters_names.append(
+ [
+ name_param_pair[0]
+ for name_param_pair in model.named_parameters()
+ ]
+ )
+ if type == 'AdamW':
+ optim[key] = AdamW(
+ model_parameters,
+ lr=lr,
+ betas=(0.9, 0.98),
+ eps=1e-9,
+ weight_decay=0.1,
+ )
+ else:
+ raise ValueError('Unknown optimizer type: %s' % type)
+
+ schedulers = dict([(key, torch.optim.lr_scheduler.ExponentialLR(opt, gamma=0.999996))
+ for key, opt in optim.items()])
+
+ multi_optim = MultiOptimizer(optim, schedulers)
+ return multi_optim
\ No newline at end of file
diff --git a/real-time-gui.py b/real-time-gui.py
new file mode 100644
index 0000000000000000000000000000000000000000..ce43e07e77ddd569981618aa59d6b14c1b86ada9
--- /dev/null
+++ b/real-time-gui.py
@@ -0,0 +1,1141 @@
+import os
+import sys
+from dotenv import load_dotenv
+import shutil
+
+load_dotenv()
+
+os.environ["OMP_NUM_THREADS"] = "4"
+if sys.platform == "darwin":
+ os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
+
+now_dir = os.getcwd()
+sys.path.append(now_dir)
+import multiprocessing
+import warnings
+import yaml
+
+warnings.simplefilter("ignore")
+
+from tqdm import tqdm
+from modules.commons import *
+import librosa
+import torchaudio
+import torchaudio.compliance.kaldi as kaldi
+
+from hf_utils import load_custom_model_from_hf
+
+import os
+import sys
+import torch
+from modules.commons import str2bool
+# Load model and configuration
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+flag_vc = False
+
+prompt_condition, mel2, style2 = None, None, None
+reference_wav_name = ""
+
+prompt_len = 3 # in seconds
+ce_dit_difference = 2.0 # 2 seconds
+fp16 = False
+@torch.no_grad()
+def custom_infer(model_set,
+ reference_wav,
+ new_reference_wav_name,
+ input_wav_res,
+ block_frame_16k,
+ skip_head,
+ skip_tail,
+ return_length,
+ diffusion_steps,
+ inference_cfg_rate,
+ max_prompt_length,
+ cd_difference=2.0,
+ ):
+ global prompt_condition, mel2, style2
+ global reference_wav_name
+ global prompt_len
+ global ce_dit_difference
+ (
+ model,
+ semantic_fn,
+ vocoder_fn,
+ campplus_model,
+ to_mel,
+ mel_fn_args,
+ ) = model_set
+ sr = mel_fn_args["sampling_rate"]
+ hop_length = mel_fn_args["hop_size"]
+ if ce_dit_difference != cd_difference:
+ ce_dit_difference = cd_difference
+ print(f"Setting ce_dit_difference to {cd_difference} seconds.")
+ if prompt_condition is None or reference_wav_name != new_reference_wav_name or prompt_len != max_prompt_length:
+ prompt_len = max_prompt_length
+ print(f"Setting max prompt length to {max_prompt_length} seconds.")
+ reference_wav = reference_wav[:int(sr * prompt_len)]
+ reference_wav_tensor = torch.from_numpy(reference_wav).to(device)
+
+ ori_waves_16k = torchaudio.functional.resample(reference_wav_tensor, sr, 16000)
+ S_ori = semantic_fn(ori_waves_16k.unsqueeze(0))
+ feat2 = torchaudio.compliance.kaldi.fbank(
+ ori_waves_16k.unsqueeze(0), num_mel_bins=80, dither=0, sample_frequency=16000
+ )
+ feat2 = feat2 - feat2.mean(dim=0, keepdim=True)
+ style2 = campplus_model(feat2.unsqueeze(0))
+
+ mel2 = to_mel(reference_wav_tensor.unsqueeze(0))
+ target2_lengths = torch.LongTensor([mel2.size(2)]).to(mel2.device)
+ prompt_condition = model.length_regulator(
+ S_ori, ylens=target2_lengths, n_quantizers=3, f0=None
+ )[0]
+
+ reference_wav_name = new_reference_wav_name
+
+ converted_waves_16k = input_wav_res
+ start_event = torch.cuda.Event(enable_timing=True)
+ end_event = torch.cuda.Event(enable_timing=True)
+ torch.cuda.synchronize()
+ start_event.record()
+ S_alt = semantic_fn(converted_waves_16k.unsqueeze(0))
+ end_event.record()
+ torch.cuda.synchronize() # Wait for the events to be recorded!
+ elapsed_time_ms = start_event.elapsed_time(end_event)
+ print(f"Time taken for semantic_fn: {elapsed_time_ms}ms")
+
+ ce_dit_frame_difference = int(ce_dit_difference * 50)
+ S_alt = S_alt[:, ce_dit_frame_difference:]
+ target_lengths = torch.LongTensor([(skip_head + return_length + skip_tail - ce_dit_frame_difference) / 50 * sr // hop_length]).to(S_alt.device)
+ print(f"target_lengths: {target_lengths}")
+ cond = model.length_regulator(
+ S_alt, ylens=target_lengths , n_quantizers=3, f0=None
+ )[0]
+ cat_condition = torch.cat([prompt_condition, cond], dim=1)
+ with torch.autocast(device_type=device.type, dtype=torch.float16 if fp16 else torch.float32):
+ vc_target = model.cfm.inference(
+ cat_condition,
+ torch.LongTensor([cat_condition.size(1)]).to(mel2.device),
+ mel2,
+ style2,
+ None,
+ n_timesteps=diffusion_steps,
+ inference_cfg_rate=inference_cfg_rate,
+ )
+ vc_target = vc_target[:, :, mel2.size(-1) :]
+ print(f"vc_target.shape: {vc_target.shape}")
+ vc_wave = vocoder_fn(vc_target).squeeze()
+ output_len = return_length * sr // 50
+ tail_len = skip_tail * sr // 50
+ output = vc_wave[-output_len - tail_len: -tail_len]
+
+ return output
+
+def load_models(args):
+ global fp16
+ fp16 = args.fp16
+ print(f"Using fp16: {fp16}")
+ if args.checkpoint_path is None or args.checkpoint_path == "":
+ dit_checkpoint_path, dit_config_path = load_custom_model_from_hf("Plachta/Seed-VC",
+ "DiT_uvit_tat_xlsr_ema.pth",
+ "config_dit_mel_seed_uvit_xlsr_tiny.yml")
+ else:
+ dit_checkpoint_path = args.checkpoint_path
+ dit_config_path = args.config_path
+ config = yaml.safe_load(open(dit_config_path, "r"))
+ model_params = recursive_munch(config["model_params"])
+ model_params.dit_type = 'DiT'
+ model = build_model(model_params, stage="DiT")
+ hop_length = config["preprocess_params"]["spect_params"]["hop_length"]
+ sr = config["preprocess_params"]["sr"]
+
+ # Load checkpoints
+ model, _, _, _ = load_checkpoint(
+ model,
+ None,
+ dit_checkpoint_path,
+ load_only_params=True,
+ ignore_modules=[],
+ is_distributed=False,
+ )
+ for key in model:
+ model[key].eval()
+ model[key].to(device)
+ model.cfm.estimator.setup_caches(max_batch_size=1, max_seq_length=8192)
+
+ # Load additional modules
+ from modules.campplus.DTDNN import CAMPPlus
+
+ campplus_ckpt_path = load_custom_model_from_hf(
+ "funasr/campplus", "campplus_cn_common.bin", config_filename=None
+ )
+ campplus_model = CAMPPlus(feat_dim=80, embedding_size=192)
+ campplus_model.load_state_dict(torch.load(campplus_ckpt_path, map_location="cpu"))
+ campplus_model.eval()
+ campplus_model.to(device)
+
+ vocoder_type = model_params.vocoder.type
+
+ if vocoder_type == 'bigvgan':
+ from modules.bigvgan import bigvgan
+ bigvgan_name = model_params.vocoder.name
+ bigvgan_model = bigvgan.BigVGAN.from_pretrained(bigvgan_name, use_cuda_kernel=False)
+ # remove weight norm in the model and set to eval mode
+ bigvgan_model.remove_weight_norm()
+ bigvgan_model = bigvgan_model.eval().to(device)
+ vocoder_fn = bigvgan_model
+ elif vocoder_type == 'hifigan':
+ from modules.hifigan.generator import HiFTGenerator
+ from modules.hifigan.f0_predictor import ConvRNNF0Predictor
+ hift_config = yaml.safe_load(open('configs/hifigan.yml', 'r'))
+ hift_gen = HiFTGenerator(**hift_config['hift'], f0_predictor=ConvRNNF0Predictor(**hift_config['f0_predictor']))
+ hift_path = load_custom_model_from_hf("FunAudioLLM/CosyVoice-300M", 'hift.pt', None)
+ hift_gen.load_state_dict(torch.load(hift_path, map_location='cpu'))
+ hift_gen.eval()
+ hift_gen.to(device)
+ vocoder_fn = hift_gen
+ elif vocoder_type == "vocos":
+ vocos_config = yaml.safe_load(open(model_params.vocoder.vocos.config, 'r'))
+ vocos_path = model_params.vocoder.vocos.path
+ vocos_model_params = recursive_munch(vocos_config['model_params'])
+ vocos = build_model(vocos_model_params, stage='mel_vocos')
+ vocos_checkpoint_path = vocos_path
+ vocos, _, _, _ = load_checkpoint(vocos, None, vocos_checkpoint_path,
+ load_only_params=True, ignore_modules=[], is_distributed=False)
+ _ = [vocos[key].eval().to(device) for key in vocos]
+ _ = [vocos[key].to(device) for key in vocos]
+ total_params = sum(sum(p.numel() for p in vocos[key].parameters() if p.requires_grad) for key in vocos.keys())
+ print(f"Vocoder model total parameters: {total_params / 1_000_000:.2f}M")
+ vocoder_fn = vocos.decoder
+ else:
+ raise ValueError(f"Unknown vocoder type: {vocoder_type}")
+
+ speech_tokenizer_type = model_params.speech_tokenizer.type
+ if speech_tokenizer_type == 'whisper':
+ # whisper
+ from transformers import AutoFeatureExtractor, WhisperModel
+ whisper_name = model_params.speech_tokenizer.name
+ whisper_model = WhisperModel.from_pretrained(whisper_name, torch_dtype=torch.float16).to(device)
+ del whisper_model.decoder
+ whisper_feature_extractor = AutoFeatureExtractor.from_pretrained(whisper_name)
+
+ def semantic_fn(waves_16k):
+ ori_inputs = whisper_feature_extractor([waves_16k.squeeze(0).cpu().numpy()],
+ return_tensors="pt",
+ return_attention_mask=True)
+ ori_input_features = whisper_model._mask_input_features(
+ ori_inputs.input_features, attention_mask=ori_inputs.attention_mask).to(device)
+ with torch.no_grad():
+ ori_outputs = whisper_model.encoder(
+ ori_input_features.to(whisper_model.encoder.dtype),
+ head_mask=None,
+ output_attentions=False,
+ output_hidden_states=False,
+ return_dict=True,
+ )
+ S_ori = ori_outputs.last_hidden_state.to(torch.float32)
+ S_ori = S_ori[:, :waves_16k.size(-1) // 320 + 1]
+ return S_ori
+ elif speech_tokenizer_type == 'cnhubert':
+ from transformers import (
+ Wav2Vec2FeatureExtractor,
+ HubertModel,
+ )
+ hubert_model_name = config['model_params']['speech_tokenizer']['name']
+ hubert_feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(hubert_model_name)
+ hubert_model = HubertModel.from_pretrained(hubert_model_name)
+ hubert_model = hubert_model.to(device)
+ hubert_model = hubert_model.eval()
+ hubert_model = hubert_model.half()
+
+ def semantic_fn(waves_16k):
+ ori_waves_16k_input_list = [
+ waves_16k[bib].cpu().numpy()
+ for bib in range(len(waves_16k))
+ ]
+ ori_inputs = hubert_feature_extractor(ori_waves_16k_input_list,
+ return_tensors="pt",
+ return_attention_mask=True,
+ padding=True,
+ sampling_rate=16000).to(device)
+ with torch.no_grad():
+ ori_outputs = hubert_model(
+ ori_inputs.input_values.half(),
+ )
+ S_ori = ori_outputs.last_hidden_state.float()
+ return S_ori
+ elif speech_tokenizer_type == 'xlsr':
+ from transformers import (
+ Wav2Vec2FeatureExtractor,
+ Wav2Vec2Model,
+ )
+ model_name = config['model_params']['speech_tokenizer']['name']
+ output_layer = config['model_params']['speech_tokenizer']['output_layer']
+ wav2vec_feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(model_name)
+ wav2vec_model = Wav2Vec2Model.from_pretrained(model_name)
+ wav2vec_model.encoder.layers = wav2vec_model.encoder.layers[:output_layer]
+ wav2vec_model = wav2vec_model.to(device)
+ wav2vec_model = wav2vec_model.eval()
+ wav2vec_model = wav2vec_model.half()
+
+ def semantic_fn(waves_16k):
+ ori_waves_16k_input_list = [
+ waves_16k[bib].cpu().numpy()
+ for bib in range(len(waves_16k))
+ ]
+ ori_inputs = wav2vec_feature_extractor(ori_waves_16k_input_list,
+ return_tensors="pt",
+ return_attention_mask=True,
+ padding=True,
+ sampling_rate=16000).to(device)
+ with torch.no_grad():
+ ori_outputs = wav2vec_model(
+ ori_inputs.input_values.half(),
+ )
+ S_ori = ori_outputs.last_hidden_state.float()
+ return S_ori
+ else:
+ raise ValueError(f"Unknown speech tokenizer type: {speech_tokenizer_type}")
+ # Generate mel spectrograms
+ mel_fn_args = {
+ "n_fft": config['preprocess_params']['spect_params']['n_fft'],
+ "win_size": config['preprocess_params']['spect_params']['win_length'],
+ "hop_size": config['preprocess_params']['spect_params']['hop_length'],
+ "num_mels": config['preprocess_params']['spect_params']['n_mels'],
+ "sampling_rate": sr,
+ "fmin": config['preprocess_params']['spect_params'].get('fmin', 0),
+ "fmax": None if config['preprocess_params']['spect_params'].get('fmax', "None") == "None" else 8000,
+ "center": False
+ }
+ from modules.audio import mel_spectrogram
+
+ to_mel = lambda x: mel_spectrogram(x, **mel_fn_args)
+
+ return (
+ model,
+ semantic_fn,
+ vocoder_fn,
+ campplus_model,
+ to_mel,
+ mel_fn_args,
+ )
+
+def printt(strr, *args):
+ if len(args) == 0:
+ print(strr)
+ else:
+ print(strr % args)
+
+class Config:
+ def __init__(self):
+ self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+if __name__ == "__main__":
+ import json
+ import multiprocessing
+ import re
+ import threading
+ import time
+ import traceback
+ from multiprocessing import Queue, cpu_count
+ import argparse
+
+ import librosa
+ import numpy as np
+ import FreeSimpleGUI as sg
+ import sounddevice as sd
+ import torch
+ import torch.nn.functional as F
+ import torchaudio.transforms as tat
+
+
+ current_dir = os.getcwd()
+ n_cpu = cpu_count()
+ class GUIConfig:
+ def __init__(self) -> None:
+ self.reference_audio_path: str = ""
+ # self.index_path: str = ""
+ self.diffusion_steps: int = 10
+ self.sr_type: str = "sr_model"
+ self.block_time: float = 0.25 # s
+ self.threhold: int = -60
+ self.crossfade_time: float = 0.05
+ self.extra_time_ce: float = 2.5
+ self.extra_time: float = 0.5
+ self.extra_time_right: float = 2.0
+ self.I_noise_reduce: bool = False
+ self.O_noise_reduce: bool = False
+ self.inference_cfg_rate: float = 0.7
+ self.sg_hostapi: str = ""
+ self.wasapi_exclusive: bool = False
+ self.sg_input_device: str = ""
+ self.sg_output_device: str = ""
+
+
+ class GUI:
+ def __init__(self, args) -> None:
+ self.gui_config = GUIConfig()
+ self.config = Config()
+ self.function = "vc"
+ self.delay_time = 0
+ self.hostapis = None
+ self.input_devices = None
+ self.output_devices = None
+ self.input_devices_indices = None
+ self.output_devices_indices = None
+ self.stream = None
+ self.model_set = load_models(args)
+ from funasr import AutoModel
+ self.vad_model = AutoModel(model="fsmn-vad", model_revision="v2.0.4")
+ self.update_devices()
+ self.launcher()
+
+ def load(self):
+ try:
+ os.makedirs("configs/inuse", exist_ok=True)
+ if not os.path.exists("configs/inuse/config.json"):
+ shutil.copy("configs/config.json", "configs/inuse/config.json")
+ with open("configs/inuse/config.json", "r") as j:
+ data = json.load(j)
+ data["sr_model"] = data["sr_type"] == "sr_model"
+ data["sr_device"] = data["sr_type"] == "sr_device"
+ if data["sg_hostapi"] in self.hostapis:
+ self.update_devices(hostapi_name=data["sg_hostapi"])
+ if (
+ data["sg_input_device"] not in self.input_devices
+ or data["sg_output_device"] not in self.output_devices
+ ):
+ self.update_devices()
+ data["sg_hostapi"] = self.hostapis[0]
+ data["sg_input_device"] = self.input_devices[
+ self.input_devices_indices.index(sd.default.device[0])
+ ]
+ data["sg_output_device"] = self.output_devices[
+ self.output_devices_indices.index(sd.default.device[1])
+ ]
+ else:
+ data["sg_hostapi"] = self.hostapis[0]
+ data["sg_input_device"] = self.input_devices[
+ self.input_devices_indices.index(sd.default.device[0])
+ ]
+ data["sg_output_device"] = self.output_devices[
+ self.output_devices_indices.index(sd.default.device[1])
+ ]
+ except:
+ with open("configs/inuse/config.json", "w") as j:
+ data = {
+ "sg_hostapi": self.hostapis[0],
+ "sg_wasapi_exclusive": False,
+ "sg_input_device": self.input_devices[
+ self.input_devices_indices.index(sd.default.device[0])
+ ],
+ "sg_output_device": self.output_devices[
+ self.output_devices_indices.index(sd.default.device[1])
+ ],
+ "sr_type": "sr_model",
+ "block_time": 0.3,
+ "crossfade_length": 0.04,
+ "extra_time_ce": 2.5,
+ "extra_time": 0.5,
+ "extra_time_right": 0.02,
+ "diffusion_steps": 10,
+ "inference_cfg_rate": 0.7,
+ "max_prompt_length": 3.0,
+ }
+ data["sr_model"] = data["sr_type"] == "sr_model"
+ data["sr_device"] = data["sr_type"] == "sr_device"
+ return data
+
+ def launcher(self):
+ self.config = Config()
+ data = self.load()
+ sg.theme("LightBlue3")
+ layout = [
+ [
+ sg.Frame(
+ title="Load reference audio",
+ layout=[
+ [
+ sg.Input(
+ default_text=data.get("reference_audio_path", ""),
+ key="reference_audio_path",
+ ),
+ sg.FileBrowse(
+ "choose an audio file",
+ initial_folder=os.path.join(
+ os.getcwd(), "examples/reference"
+ ),
+ file_types=((". wav"), (". mp3"), (". flac"), (". m4a"), (". ogg"), (". opus")),
+ ),
+ ],
+ ],
+ )
+ ],
+ [
+ sg.Frame(
+ layout=[
+ [
+ sg.Text("Device type"),
+ sg.Combo(
+ self.hostapis,
+ key="sg_hostapi",
+ default_value=data.get("sg_hostapi", ""),
+ enable_events=True,
+ size=(20, 1),
+ ),
+ sg.Checkbox(
+ "WASAPI Exclusive Device",
+ key="sg_wasapi_exclusive",
+ default=data.get("sg_wasapi_exclusive", False),
+ enable_events=True,
+ ),
+ ],
+ [
+ sg.Text("Input Device"),
+ sg.Combo(
+ self.input_devices,
+ key="sg_input_device",
+ default_value=data.get("sg_input_device", ""),
+ enable_events=True,
+ size=(45, 1),
+ ),
+ ],
+ [
+ sg.Text("Output Device"),
+ sg.Combo(
+ self.output_devices,
+ key="sg_output_device",
+ default_value=data.get("sg_output_device", ""),
+ enable_events=True,
+ size=(45, 1),
+ ),
+ ],
+ [
+ sg.Button("Reload devices", key="reload_devices"),
+ sg.Radio(
+ "Use model sampling rate",
+ "sr_type",
+ key="sr_model",
+ default=data.get("sr_model", True),
+ enable_events=True,
+ ),
+ sg.Radio(
+ "Use device sampling rate",
+ "sr_type",
+ key="sr_device",
+ default=data.get("sr_device", False),
+ enable_events=True,
+ ),
+ sg.Text("Sampling rate:"),
+ sg.Text("", key="sr_stream"),
+ ],
+ ],
+ title="Sound Device",
+ )
+ ],
+ [
+ sg.Frame(
+ layout=[
+ # [
+ # sg.Text("Activation threshold"),
+ # sg.Slider(
+ # range=(-60, 0),
+ # key="threhold",
+ # resolution=1,
+ # orientation="h",
+ # default_value=data.get("threhold", -60),
+ # enable_events=True,
+ # ),
+ # ],
+ [
+ sg.Text("Diffusion steps"),
+ sg.Slider(
+ range=(1, 30),
+ key="diffusion_steps",
+ resolution=1,
+ orientation="h",
+ default_value=data.get("diffusion_steps", 10),
+ enable_events=True,
+ ),
+ ],
+ [
+ sg.Text("Inference cfg rate"),
+ sg.Slider(
+ range=(0.0, 1.0),
+ key="inference_cfg_rate",
+ resolution=0.1,
+ orientation="h",
+ default_value=data.get("inference_cfg_rate", 0.7),
+ enable_events=True,
+ ),
+ ],
+ [
+ sg.Text("Max prompt length (s)"),
+ sg.Slider(
+ range=(1.0, 20.0),
+ key="max_prompt_length",
+ resolution=0.5,
+ orientation="h",
+ default_value=data.get("max_prompt_length", 3.0),
+ enable_events=True,
+ ),
+ ],
+ ],
+ title="Regular settings",
+ ),
+ sg.Frame(
+ layout=[
+ [
+ sg.Text("Block time"),
+ sg.Slider(
+ range=(0.04, 3.0),
+ key="block_time",
+ resolution=0.02,
+ orientation="h",
+ default_value=data.get("block_time", 1.0),
+ enable_events=True,
+ ),
+ ],
+ [
+ sg.Text("Crossfade length"),
+ sg.Slider(
+ range=(0.02, 0.5),
+ key="crossfade_length",
+ resolution=0.02,
+ orientation="h",
+ default_value=data.get("crossfade_length", 0.1),
+ enable_events=True,
+ ),
+ ],
+ [
+ sg.Text("Extra content encoder context time (left)"),
+ sg.Slider(
+ range=(0.5, 10.0),
+ key="extra_time_ce",
+ resolution=0.1,
+ orientation="h",
+ default_value=data.get("extra_time_ce", 5.0),
+ enable_events=True,
+ ),
+ ],
+ [
+ sg.Text("Extra DiT context time (left)"),
+ sg.Slider(
+ range=(0.5, 10.0),
+ key="extra_time",
+ resolution=0.1,
+ orientation="h",
+ default_value=data.get("extra_time", 5.0),
+ enable_events=True,
+ ),
+ ],
+ [
+ sg.Text("Extra context time (right)"),
+ sg.Slider(
+ range=(0.02, 10.0),
+ key="extra_time_right",
+ resolution=0.02,
+ orientation="h",
+ default_value=data.get("extra_time_right", 2.0),
+ enable_events=True,
+ ),
+ ],
+ ],
+ title="Performance settings",
+ ),
+ ],
+ [
+ sg.Button("Start Voice Conversion", key="start_vc"),
+ sg.Button("Stop Voice Conversion", key="stop_vc"),
+ sg.Radio(
+ "Input listening",
+ "function",
+ key="im",
+ default=False,
+ enable_events=True,
+ ),
+ sg.Radio(
+ "Voice Conversion",
+ "function",
+ key="vc",
+ default=True,
+ enable_events=True,
+ ),
+ sg.Text("Algorithm delay (ms):"),
+ sg.Text("0", key="delay_time"),
+ sg.Text("Inference time (ms):"),
+ sg.Text("0", key="infer_time"),
+ ],
+ ]
+ self.window = sg.Window("Seed-VC - GUI", layout=layout, finalize=True)
+ self.event_handler()
+
+ def event_handler(self):
+ global flag_vc
+ while True:
+ event, values = self.window.read()
+ if event == sg.WINDOW_CLOSED:
+ self.stop_stream()
+ exit()
+ if event == "reload_devices" or event == "sg_hostapi":
+ self.gui_config.sg_hostapi = values["sg_hostapi"]
+ self.update_devices(hostapi_name=values["sg_hostapi"])
+ if self.gui_config.sg_hostapi not in self.hostapis:
+ self.gui_config.sg_hostapi = self.hostapis[0]
+ self.window["sg_hostapi"].Update(values=self.hostapis)
+ self.window["sg_hostapi"].Update(value=self.gui_config.sg_hostapi)
+ if (
+ self.gui_config.sg_input_device not in self.input_devices
+ and len(self.input_devices) > 0
+ ):
+ self.gui_config.sg_input_device = self.input_devices[0]
+ self.window["sg_input_device"].Update(values=self.input_devices)
+ self.window["sg_input_device"].Update(
+ value=self.gui_config.sg_input_device
+ )
+ if self.gui_config.sg_output_device not in self.output_devices:
+ self.gui_config.sg_output_device = self.output_devices[0]
+ self.window["sg_output_device"].Update(values=self.output_devices)
+ self.window["sg_output_device"].Update(
+ value=self.gui_config.sg_output_device
+ )
+ if event == "start_vc" and not flag_vc:
+ if self.set_values(values) == True:
+ printt("cuda_is_available: %s", torch.cuda.is_available())
+ self.start_vc()
+ settings = {
+ "reference_audio_path": values["reference_audio_path"],
+ # "index_path": values["index_path"],
+ "sg_hostapi": values["sg_hostapi"],
+ "sg_wasapi_exclusive": values["sg_wasapi_exclusive"],
+ "sg_input_device": values["sg_input_device"],
+ "sg_output_device": values["sg_output_device"],
+ "sr_type": ["sr_model", "sr_device"][
+ [
+ values["sr_model"],
+ values["sr_device"],
+ ].index(True)
+ ],
+ # "threhold": values["threhold"],
+ "diffusion_steps": values["diffusion_steps"],
+ "inference_cfg_rate": values["inference_cfg_rate"],
+ "max_prompt_length": values["max_prompt_length"],
+ "block_time": values["block_time"],
+ "crossfade_length": values["crossfade_length"],
+ "extra_time_ce": values["extra_time_ce"],
+ "extra_time": values["extra_time"],
+ "extra_time_right": values["extra_time_right"],
+ }
+ with open("configs/inuse/config.json", "w") as j:
+ json.dump(settings, j)
+ if self.stream is not None:
+ self.delay_time = (
+ self.stream.latency[-1]
+ + values["block_time"]
+ + values["crossfade_length"]
+ + values["extra_time_right"]
+ + 0.01
+ )
+ self.window["sr_stream"].update(self.gui_config.samplerate)
+ self.window["delay_time"].update(
+ int(np.round(self.delay_time * 1000))
+ )
+ # Parameter hot update
+ # if event == "threhold":
+ # self.gui_config.threhold = values["threhold"]
+ elif event == "diffusion_steps":
+ self.gui_config.diffusion_steps = values["diffusion_steps"]
+ elif event == "inference_cfg_rate":
+ self.gui_config.inference_cfg_rate = values["inference_cfg_rate"]
+ elif event in ["vc", "im"]:
+ self.function = event
+ elif event == "stop_vc" or event != "start_vc":
+ # Other parameters do not support hot update
+ self.stop_stream()
+
+ def set_values(self, values):
+ if len(values["reference_audio_path"].strip()) == 0:
+ sg.popup("Choose an audio file")
+ return False
+ pattern = re.compile("[^\x00-\x7F]+")
+ if pattern.findall(values["reference_audio_path"]):
+ sg.popup("audio file path contains non-ascii characters")
+ return False
+ self.set_devices(values["sg_input_device"], values["sg_output_device"])
+ self.gui_config.sg_hostapi = values["sg_hostapi"]
+ self.gui_config.sg_wasapi_exclusive = values["sg_wasapi_exclusive"]
+ self.gui_config.sg_input_device = values["sg_input_device"]
+ self.gui_config.sg_output_device = values["sg_output_device"]
+ self.gui_config.reference_audio_path = values["reference_audio_path"]
+ self.gui_config.sr_type = ["sr_model", "sr_device"][
+ [
+ values["sr_model"],
+ values["sr_device"],
+ ].index(True)
+ ]
+ # self.gui_config.threhold = values["threhold"]
+ self.gui_config.diffusion_steps = values["diffusion_steps"]
+ self.gui_config.inference_cfg_rate = values["inference_cfg_rate"]
+ self.gui_config.max_prompt_length = values["max_prompt_length"]
+ self.gui_config.block_time = values["block_time"]
+ self.gui_config.crossfade_time = values["crossfade_length"]
+ self.gui_config.extra_time_ce = values["extra_time_ce"]
+ self.gui_config.extra_time = values["extra_time"]
+ self.gui_config.extra_time_right = values["extra_time_right"]
+ return True
+
+ def start_vc(self):
+ torch.cuda.empty_cache()
+ self.reference_wav, _ = librosa.load(
+ self.gui_config.reference_audio_path, sr=self.model_set[-1]["sampling_rate"]
+ )
+ self.gui_config.samplerate = (
+ self.model_set[-1]["sampling_rate"]
+ if self.gui_config.sr_type == "sr_model"
+ else self.get_device_samplerate()
+ )
+ self.gui_config.channels = self.get_device_channels()
+ self.zc = self.gui_config.samplerate // 50 # 44100 // 100 = 441
+ self.block_frame = (
+ int(
+ np.round(
+ self.gui_config.block_time
+ * self.gui_config.samplerate
+ / self.zc
+ )
+ )
+ * self.zc
+ )
+ self.block_frame_16k = 320 * self.block_frame // self.zc
+ self.crossfade_frame = (
+ int(
+ np.round(
+ self.gui_config.crossfade_time
+ * self.gui_config.samplerate
+ / self.zc
+ )
+ )
+ * self.zc
+ )
+ self.sola_buffer_frame = min(self.crossfade_frame, 4 * self.zc)
+ self.sola_search_frame = self.zc
+ self.extra_frame = (
+ int(
+ np.round(
+ self.gui_config.extra_time_ce
+ * self.gui_config.samplerate
+ / self.zc
+ )
+ )
+ * self.zc
+ )
+ self.extra_frame_right = (
+ int(
+ np.round(
+ self.gui_config.extra_time_right
+ * self.gui_config.samplerate
+ / self.zc
+ )
+ )
+ * self.zc
+ )
+ self.input_wav: torch.Tensor = torch.zeros(
+ self.extra_frame
+ + self.crossfade_frame
+ + self.sola_search_frame
+ + self.block_frame
+ + self.extra_frame_right,
+ device=self.config.device,
+ dtype=torch.float32,
+ ) # 2 * 44100 + 0.08 * 44100 + 0.01 * 44100 + 0.25 * 44100
+ self.input_wav_denoise: torch.Tensor = self.input_wav.clone()
+ self.input_wav_res: torch.Tensor = torch.zeros(
+ 320 * self.input_wav.shape[0] // self.zc,
+ device=self.config.device,
+ dtype=torch.float32,
+ ) # input wave 44100 -> 16000
+ self.rms_buffer: np.ndarray = np.zeros(4 * self.zc, dtype="float32")
+ self.sola_buffer: torch.Tensor = torch.zeros(
+ self.sola_buffer_frame, device=self.config.device, dtype=torch.float32
+ )
+ self.nr_buffer: torch.Tensor = self.sola_buffer.clone()
+ self.output_buffer: torch.Tensor = self.input_wav.clone()
+ self.skip_head = self.extra_frame // self.zc
+ self.skip_tail = self.extra_frame_right // self.zc
+ self.return_length = (
+ self.block_frame + self.sola_buffer_frame + self.sola_search_frame
+ ) // self.zc
+ self.fade_in_window: torch.Tensor = (
+ torch.sin(
+ 0.5
+ * np.pi
+ * torch.linspace(
+ 0.0,
+ 1.0,
+ steps=self.sola_buffer_frame,
+ device=self.config.device,
+ dtype=torch.float32,
+ )
+ )
+ ** 2
+ )
+ self.fade_out_window: torch.Tensor = 1 - self.fade_in_window
+ self.resampler = tat.Resample(
+ orig_freq=self.gui_config.samplerate,
+ new_freq=16000,
+ dtype=torch.float32,
+ ).to(self.config.device)
+ if self.model_set[-1]["sampling_rate"] != self.gui_config.samplerate:
+ self.resampler2 = tat.Resample(
+ orig_freq=self.model_set[-1]["sampling_rate"],
+ new_freq=self.gui_config.samplerate,
+ dtype=torch.float32,
+ ).to(self.config.device)
+ else:
+ self.resampler2 = None
+ self.vad_cache = {}
+ self.vad_chunk_size = 1000 * self.gui_config.block_time
+ self.vad_speech_detected = False
+ self.set_speech_detected_false_at_end_flag = False
+ self.start_stream()
+
+ def start_stream(self):
+ global flag_vc
+ if not flag_vc:
+ flag_vc = True
+ if (
+ "WASAPI" in self.gui_config.sg_hostapi
+ and self.gui_config.sg_wasapi_exclusive
+ ):
+ extra_settings = sd.WasapiSettings(exclusive=True)
+ else:
+ extra_settings = None
+ self.stream = sd.Stream(
+ callback=self.audio_callback,
+ blocksize=self.block_frame,
+ samplerate=self.gui_config.samplerate,
+ channels=self.gui_config.channels,
+ dtype="float32",
+ extra_settings=extra_settings,
+ )
+ self.stream.start()
+
+ def stop_stream(self):
+ global flag_vc
+ if flag_vc:
+ flag_vc = False
+ if self.stream is not None:
+ self.stream.abort()
+ self.stream.close()
+ self.stream = None
+
+ def audio_callback(
+ self, indata: np.ndarray, outdata: np.ndarray, frames, times, status
+ ):
+ """
+ Audio block callback function
+ """
+ global flag_vc
+ print(indata.shape)
+ start_time = time.perf_counter()
+ indata = librosa.to_mono(indata.T)
+
+ # VAD first
+ start_event = torch.cuda.Event(enable_timing=True)
+ end_event = torch.cuda.Event(enable_timing=True)
+ torch.cuda.synchronize()
+ start_event.record()
+ indata_16k = librosa.resample(indata, orig_sr=self.gui_config.samplerate, target_sr=16000)
+ res = self.vad_model.generate(input=indata_16k, cache=self.vad_cache, is_final=False, chunk_size=self.vad_chunk_size)
+ res_value = res[0]["value"]
+ print(res_value)
+ if len(res_value) % 2 == 1 and not self.vad_speech_detected:
+ self.vad_speech_detected = True
+ elif len(res_value) % 2 == 1 and self.vad_speech_detected:
+ self.set_speech_detected_false_at_end_flag = True
+ end_event.record()
+ torch.cuda.synchronize() # Wait for the events to be recorded!
+ elapsed_time_ms = start_event.elapsed_time(end_event)
+ print(f"Time taken for VAD: {elapsed_time_ms}ms")
+
+ # if self.gui_config.threhold > -60:
+ # indata = np.append(self.rms_buffer, indata)
+ # rms = librosa.feature.rms(
+ # y=indata, frame_length=4 * self.zc, hop_length=self.zc
+ # )[:, 2:]
+ # self.rms_buffer[:] = indata[-4 * self.zc :]
+ # indata = indata[2 * self.zc - self.zc // 2 :]
+ # db_threhold = (
+ # librosa.amplitude_to_db(rms, ref=1.0)[0] < self.gui_config.threhold
+ # )
+ # for i in range(db_threhold.shape[0]):
+ # if db_threhold[i]:
+ # indata[i * self.zc : (i + 1) * self.zc] = 0
+ # indata = indata[self.zc // 2 :]
+ self.input_wav[: -self.block_frame] = self.input_wav[
+ self.block_frame :
+ ].clone()
+ self.input_wav[-indata.shape[0] :] = torch.from_numpy(indata).to(
+ self.config.device
+ )
+ self.input_wav_res[: -self.block_frame_16k] = self.input_wav_res[
+ self.block_frame_16k :
+ ].clone()
+ self.input_wav_res[-320 * (indata.shape[0] // self.zc + 1) :] = (
+ self.resampler(self.input_wav[-indata.shape[0] - 2 * self.zc :])[
+ 320:
+ ]
+ )
+ print(f"preprocess time: {time.perf_counter() - start_time:.2f}")
+ # infer
+ if self.function == "vc":
+ if self.gui_config.extra_time_ce - self.gui_config.extra_time < 0:
+ raise ValueError("Content encoder extra context must be greater than DiT extra context!")
+ start_event = torch.cuda.Event(enable_timing=True)
+ end_event = torch.cuda.Event(enable_timing=True)
+ torch.cuda.synchronize()
+ start_event.record()
+ infer_wav = custom_infer(
+ self.model_set,
+ self.reference_wav,
+ self.gui_config.reference_audio_path,
+ self.input_wav_res,
+ self.block_frame_16k,
+ self.skip_head,
+ self.skip_tail,
+ self.return_length,
+ int(self.gui_config.diffusion_steps),
+ self.gui_config.inference_cfg_rate,
+ self.gui_config.max_prompt_length,
+ self.gui_config.extra_time_ce - self.gui_config.extra_time,
+ )
+ if self.resampler2 is not None:
+ infer_wav = self.resampler2(infer_wav)
+ end_event.record()
+ torch.cuda.synchronize() # Wait for the events to be recorded!
+ elapsed_time_ms = start_event.elapsed_time(end_event)
+ print(f"Time taken for VC: {elapsed_time_ms}ms")
+ if not self.vad_speech_detected:
+ infer_wav = torch.zeros_like(self.input_wav[self.extra_frame :])
+ elif self.gui_config.I_noise_reduce:
+ infer_wav = self.input_wav_denoise[self.extra_frame :].clone()
+ else:
+ infer_wav = self.input_wav[self.extra_frame :].clone()
+
+ # SOLA algorithm from https://github.com/yxlllc/DDSP-SVC
+ conv_input = infer_wav[
+ None, None, : self.sola_buffer_frame + self.sola_search_frame
+ ]
+
+ cor_nom = F.conv1d(conv_input, self.sola_buffer[None, None, :])
+ cor_den = torch.sqrt(
+ F.conv1d(
+ conv_input**2,
+ torch.ones(1, 1, self.sola_buffer_frame, device=self.config.device),
+ )
+ + 1e-8
+ )
+ if sys.platform == "darwin":
+ _, sola_offset = torch.max(cor_nom[0, 0] / cor_den[0, 0])
+ sola_offset = sola_offset.item()
+ else:
+
+ sola_offset = torch.argmax(cor_nom[0, 0] / cor_den[0, 0])
+
+ print(f"sola_offset = {int(sola_offset)}")
+
+ #post_process_start = time.perf_counter()
+ infer_wav = infer_wav[sola_offset:]
+ infer_wav[: self.sola_buffer_frame] *= self.fade_in_window
+ infer_wav[: self.sola_buffer_frame] += (
+ self.sola_buffer * self.fade_out_window
+ )
+ self.sola_buffer[:] = infer_wav[
+ self.block_frame : self.block_frame + self.sola_buffer_frame
+ ]
+ outdata[:] = (
+ infer_wav[: self.block_frame]
+ .repeat(self.gui_config.channels, 1)
+ .t()
+ .cpu()
+ .numpy()
+ )
+
+ total_time = time.perf_counter() - start_time
+ if flag_vc:
+ self.window["infer_time"].update(int(total_time * 1000))
+
+ if self.set_speech_detected_false_at_end_flag:
+ self.vad_speech_detected = False
+ self.set_speech_detected_false_at_end_flag = False
+
+ print(f"Infer time: {total_time:.2f}")
+
+ def update_devices(self, hostapi_name=None):
+ """Get input and output devices."""
+ global flag_vc
+ flag_vc = False
+ sd._terminate()
+ sd._initialize()
+ devices = sd.query_devices()
+ hostapis = sd.query_hostapis()
+ for hostapi in hostapis:
+ for device_idx in hostapi["devices"]:
+ devices[device_idx]["hostapi_name"] = hostapi["name"]
+ self.hostapis = [hostapi["name"] for hostapi in hostapis]
+ if hostapi_name not in self.hostapis:
+ hostapi_name = self.hostapis[0]
+ self.input_devices = [
+ d["name"]
+ for d in devices
+ if d["max_input_channels"] > 0 and d["hostapi_name"] == hostapi_name
+ ]
+ self.output_devices = [
+ d["name"]
+ for d in devices
+ if d["max_output_channels"] > 0 and d["hostapi_name"] == hostapi_name
+ ]
+ self.input_devices_indices = [
+ d["index"] if "index" in d else d["name"]
+ for d in devices
+ if d["max_input_channels"] > 0 and d["hostapi_name"] == hostapi_name
+ ]
+ self.output_devices_indices = [
+ d["index"] if "index" in d else d["name"]
+ for d in devices
+ if d["max_output_channels"] > 0 and d["hostapi_name"] == hostapi_name
+ ]
+
+ def set_devices(self, input_device, output_device):
+ """set input and output devices."""
+ sd.default.device[0] = self.input_devices_indices[
+ self.input_devices.index(input_device)
+ ]
+ sd.default.device[1] = self.output_devices_indices[
+ self.output_devices.index(output_device)
+ ]
+ printt("Input device: %s:%s", str(sd.default.device[0]), input_device)
+ printt("Output device: %s:%s", str(sd.default.device[1]), output_device)
+
+ def get_device_samplerate(self):
+ return int(
+ sd.query_devices(device=sd.default.device[0])["default_samplerate"]
+ )
+
+ def get_device_channels(self):
+ max_input_channels = sd.query_devices(device=sd.default.device[0])[
+ "max_input_channels"
+ ]
+ max_output_channels = sd.query_devices(device=sd.default.device[1])[
+ "max_output_channels"
+ ]
+ return min(max_input_channels, max_output_channels, 2)
+
+
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--checkpoint-path", type=str, default=None, help="Path to the model checkpoint")
+ parser.add_argument("--config-path", type=str, default=None, help="Path to the vocoder checkpoint")
+ parser.add_argument("--fp16", type=str2bool, nargs="?", const=True, help="Whether to use fp16", default=True)
+ args = parser.parse_args()
+ gui = GUI(args)
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000000000000000000000000000000000000..ce4775d72b90e4de8778bed7324d7d1c4ec3a102
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1,20 @@
+--extra-index-url https://download.pytorch.org/whl/cu113
+torch
+torchvision
+torchaudio
+scipy==1.13.1
+librosa==0.10.2
+huggingface-hub
+munch
+einops
+descript-audio-codec
+gradio
+pydub
+resemblyzer
+jiwer
+transformers
+FreeSimpleGUI
+soundfile
+sounddevice
+modelscope
+funasr
\ No newline at end of file
diff --git a/ruff.toml b/ruff.toml
new file mode 100644
index 0000000000000000000000000000000000000000..1925fcfae9cbb3b3693bf2d609a1d1ea4c25850e
--- /dev/null
+++ b/ruff.toml
@@ -0,0 +1,70 @@
+# Exclude a variety of commonly ignored directories.
+exclude = [
+ ".bzr",
+ ".direnv",
+ ".eggs",
+ ".git",
+ ".hg",
+ ".mypy_cache",
+ ".nox",
+ ".pants.d",
+ ".ruff_cache",
+ ".svn",
+ ".tox",
+ ".venv",
+ "__pypackages__",
+ "_build",
+ "buck-out",
+ "build",
+ "dist",
+ "node_modules",
+ "venv",
+]
+extend-exclude = []
+line-length = 88
+indent-width = 4
+target-version = "py310"
+show-fixes = true
+src = [".", "modules"]
+
+[lint]
+select = [
+ "E", "F", "B", "Q", "I", "C90", "N", "D", "UP", "YTT", "ANN", "S", "BLE",
+ "FBT", "A", "COM", "C4", "DTZ", "T10", "EM", "EXE", "ISC", "ICN", "INP",
+ "PIE", "T20", "PT", "Q", "RET", "SIM", "ARG", "ERA", "PD", "PGH", "PL",
+ "TRY", "RUF",
+]
+ignore = [
+ "D105",
+ "D107",
+ "D203",
+ "D213",
+ "S101", # assert-used
+ "INP001", # implicit-namespace-package
+ "ANN101", # missing-type-self
+ "ANN102", # missing-type-cls
+ "ANN204", # missing-return-type-special-method
+ "ERA001", # commented-out-code
+ "ANN002", # missing-type-args
+ "ANN003", # missing-type-kwargs
+ "RET504", # unnecessary-assign
+ "COM812", # TBD: some conflict
+ "ISC001", # TBD: some conflict
+]
+fixable = ["ALL"]
+unfixable = []
+
+[format]
+quote-style = "double"
+indent-style = "space"
+
+[lint.isort]
+# force-sort-within-sections and lines-between-types should be incompatible
+force-sort-within-sections = false
+lines-between-types = 1
+force-single-line = true
+no-sections = false
+from-first = false
+
+[lint.pydocstyle]
+convention = "google"
diff --git a/train.py b/train.py
new file mode 100644
index 0000000000000000000000000000000000000000..ec64e8e4b7045e214beb1b2c1eb1d80992e2a3c5
--- /dev/null
+++ b/train.py
@@ -0,0 +1,408 @@
+import os
+os.environ['HF_HUB_CACHE'] = './checkpoints/hf_cache'
+import torch
+import random
+import librosa
+import yaml
+import argparse
+import torchaudio
+import torchaudio.compliance.kaldi as kaldi
+import glob
+from tqdm import tqdm
+
+from modules.commons import recursive_munch, build_model, load_checkpoint
+from optimizers import build_optimizer
+from data.ft_dataset import build_ft_dataloader
+from hf_utils import load_custom_model_from_hf
+
+
+
+
+class Trainer:
+ def __init__(self,
+ config_path,
+ pretrained_ckpt_path,
+ data_dir,
+ run_name,
+ batch_size=0,
+ num_workers=0,
+ steps=1000,
+ save_interval=500,
+ max_epochs=1000,
+ device="cuda:0",
+ ):
+ self.device = device
+ config = yaml.safe_load(open(config_path))
+ self.log_dir = os.path.join(config['log_dir'], run_name)
+ os.makedirs(self.log_dir, exist_ok=True)
+ # copy config file to log dir
+ os.system(f'cp {config_path} {self.log_dir}')
+ batch_size = config.get('batch_size', 10) if batch_size == 0 else batch_size
+ self.max_steps = steps
+
+ self.n_epochs = max_epochs
+ self.log_interval = config.get('log_interval', 10)
+ self.save_interval = save_interval
+
+ self.sr = config['preprocess_params'].get('sr', 22050)
+ self.hop_length = config['preprocess_params']['spect_params'].get('hop_length', 256)
+ self.win_length = config['preprocess_params']['spect_params'].get('win_length', 1024)
+ self.n_fft = config['preprocess_params']['spect_params'].get('n_fft', 1024)
+ preprocess_params = config['preprocess_params']
+
+ self.train_dataloader = build_ft_dataloader(
+ data_dir,
+ preprocess_params['spect_params'],
+ self.sr,
+ batch_size=batch_size,
+ num_workers=num_workers,
+ )
+ self.f0_condition = config['model_params']['DiT'].get('f0_condition', False)
+ self.build_sv_model(device, config)
+ self.build_semantic_fn(device, config)
+ if self.f0_condition:
+ self.build_f0_fn(device, config)
+ self.build_converter(device, config)
+ self.build_vocoder(device, config)
+
+ scheduler_params = {
+ "warmup_steps": 0,
+ "base_lr": 0.00001,
+ }
+
+ self.model_params = recursive_munch(config['model_params'])
+ self.model = build_model(self.model_params, stage='DiT')
+
+ _ = [self.model[key].to(device) for key in self.model]
+ self.model.cfm.estimator.setup_caches(max_batch_size=batch_size, max_seq_length=8192)
+
+ # initialize optimizers after preparing models for compatibility with FSDP
+ self.optimizer = build_optimizer({key: self.model[key] for key in self.model},
+ lr=float(scheduler_params['base_lr']))
+
+ if pretrained_ckpt_path is None:
+ # find latest checkpoint with name pattern of 'T2V_epoch_*_step_*.pth'
+ available_checkpoints = glob.glob(os.path.join(self.log_dir, "DiT_epoch_*_step_*.pth"))
+ if len(available_checkpoints) > 0:
+ # find the checkpoint that has the highest step number
+ latest_checkpoint = max(
+ available_checkpoints, key=lambda x: int(x.split("_")[-1].split(".")[0])
+ )
+ earliest_checkpoint = min(
+ available_checkpoints, key=lambda x: int(x.split("_")[-1].split(".")[0])
+ )
+ # delete the earliest checkpoint
+ if (
+ earliest_checkpoint != latest_checkpoint
+ and len(available_checkpoints) > 2
+ ):
+ os.remove(earliest_checkpoint)
+ print(f"Removed {earliest_checkpoint}")
+ elif config.get('pretrained_model', ''):
+ latest_checkpoint = load_custom_model_from_hf("Plachta/Seed-VC", config['pretrained_model'], None)
+ else:
+ latest_checkpoint = ""
+ else:
+ assert os.path.exists(pretrained_ckpt_path), f"Pretrained checkpoint {pretrained_ckpt_path} not found"
+ latest_checkpoint = pretrained_ckpt_path
+
+ if os.path.exists(latest_checkpoint):
+ self.model, self.optimizer, self.epoch, self.iters = load_checkpoint(self.model, self.optimizer, latest_checkpoint,
+ load_only_params=True,
+ ignore_modules=[],
+ is_distributed=False)
+ print(f"Loaded checkpoint from {latest_checkpoint}")
+ else:
+ self.epoch, self.iters = 0, 0
+ print("Failed to load any checkpoint, this implies you are training from scratch.")
+ def build_sv_model(self, device, config):
+ # speaker verification model
+ from modules.campplus.DTDNN import CAMPPlus
+ self.campplus_model = CAMPPlus(feat_dim=80, embedding_size=192)
+ campplus_sd_path = load_custom_model_from_hf("funasr/campplus", "campplus_cn_common.bin", config_filename=None)
+ campplus_sd = torch.load(campplus_sd_path, map_location='cpu')
+ self.campplus_model.load_state_dict(campplus_sd)
+ self.campplus_model.eval()
+ self.campplus_model.to(device)
+ self.sv_fn = self.campplus_model
+ def build_f0_fn(self, device, config):
+ from modules.rmvpe import RMVPE
+ model_path = load_custom_model_from_hf("lj1995/VoiceConversionWebUI", "rmvpe.pt", None)
+ self.rmvpe = RMVPE(model_path, is_half=False, device=device)
+ self.f0_fn = self.rmvpe
+ def build_converter(self, device, config):
+ # speaker perturbation model
+ from modules.openvoice.api import ToneColorConverter
+ ckpt_converter, config_converter = load_custom_model_from_hf("myshell-ai/OpenVoiceV2", "converter/checkpoint.pth", "converter/config.json")
+ self.tone_color_converter = ToneColorConverter(config_converter, device=device,)
+ self.tone_color_converter.load_ckpt(ckpt_converter)
+ self.tone_color_converter.model.eval()
+ se_db_path = load_custom_model_from_hf("Plachta/Seed-VC", "se_db.pt", None)
+ self.se_db = torch.load(se_db_path, map_location='cpu')
+
+ def build_vocoder(self, device, config):
+ vocoder_type = config['model_params']['vocoder']['type']
+ vocoder_name = config['model_params']['vocoder'].get('name', None)
+ if vocoder_type == 'bigvgan':
+ from modules.bigvgan import bigvgan
+ bigvgan_name = vocoder_name
+ self.bigvgan_model = bigvgan.BigVGAN.from_pretrained(bigvgan_name, use_cuda_kernel=False)
+ # remove weight norm in the model and set to eval mode
+ self.bigvgan_model.remove_weight_norm()
+ self.bigvgan_model = self.bigvgan_model.eval().to(device)
+ vocoder_fn = self.bigvgan_model
+ elif vocoder_type == 'hifigan':
+ from modules.hifigan.generator import HiFTGenerator
+ from modules.hifigan.f0_predictor import ConvRNNF0Predictor
+ hift_config = yaml.safe_load(open('configs/hifigan.yml', 'r'))
+ hift_path = load_custom_model_from_hf("FunAudioLLM/CosyVoice-300M", 'hift.pt', None)
+ self.hift_gen = HiFTGenerator(**hift_config['hift'],
+ f0_predictor=ConvRNNF0Predictor(**hift_config['f0_predictor']))
+ self.hift_gen.load_state_dict(torch.load(hift_path, map_location='cpu'))
+ self.hift_gen.eval()
+ self.hift_gen.to(device)
+ vocoder_fn = self.hift_gen
+ else:
+ raise ValueError(f"Unsupported vocoder type: {vocoder_type}")
+ self.vocoder_fn = vocoder_fn
+
+ def build_semantic_fn(self, device, config):
+ # speech tokenizer
+ speech_tokenizer_type = config['model_params']['speech_tokenizer'].get('type', 'cosyvoice')
+ if speech_tokenizer_type == 'whisper':
+ from transformers import AutoFeatureExtractor, WhisperModel
+ whisper_model_name = config['model_params']['speech_tokenizer']['name']
+ self.whisper_model = WhisperModel.from_pretrained(whisper_model_name).to(device)
+ self.whisper_feature_extractor = AutoFeatureExtractor.from_pretrained(whisper_model_name)
+ del self.whisper_model.decoder
+ def semantic_fn(waves_16k):
+ ori_inputs = self.whisper_feature_extractor([w16k.cpu().numpy() for w16k in waves_16k],
+ return_tensors="pt",
+ return_attention_mask=True,
+ sampling_rate=16000,)
+ ori_input_features = self.whisper_model._mask_input_features(
+ ori_inputs.input_features, attention_mask=ori_inputs.attention_mask).to(device)
+ with torch.no_grad():
+ ori_outputs = self.whisper_model.encoder(
+ ori_input_features.to(self.whisper_model.encoder.dtype),
+ head_mask=None,
+ output_attentions=False,
+ output_hidden_states=False,
+ return_dict=True,
+ )
+ S_ori = ori_outputs.last_hidden_state.to(torch.float32)
+ S_ori = S_ori[:, :waves_16k.size(-1) // 320 + 1]
+ return S_ori
+ elif speech_tokenizer_type == 'xlsr':
+ from transformers import (
+ Wav2Vec2FeatureExtractor,
+ Wav2Vec2Model,
+ )
+ model_name = config['model_params']['speech_tokenizer']['name']
+ output_layer = config['model_params']['speech_tokenizer']['output_layer']
+ self.wav2vec_feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(model_name)
+ self.wav2vec_model = Wav2Vec2Model.from_pretrained(model_name)
+ self.wav2vec_model.encoder.layers = self.wav2vec_model.encoder.layers[:output_layer]
+ self.wav2vec_model = self.wav2vec_model.to(device)
+ self.wav2vec_model = self.wav2vec_model.eval()
+ self.wav2vec_model = self.wav2vec_model.half()
+
+ def semantic_fn(waves_16k):
+ ori_waves_16k_input_list = [
+ waves_16k[bib].cpu().numpy()
+ for bib in range(len(waves_16k))
+ ]
+ ori_inputs = self.wav2vec_feature_extractor(ori_waves_16k_input_list,
+ return_tensors="pt",
+ return_attention_mask=True,
+ padding=True,
+ sampling_rate=16000).to(device)
+ with torch.no_grad():
+ ori_outputs = self.wav2vec_model(
+ ori_inputs.input_values.half(),
+ )
+ S_ori = ori_outputs.last_hidden_state.float()
+ return S_ori
+ else:
+ raise ValueError(f"Unsupported speech tokenizer type: {speech_tokenizer_type}")
+ self.semantic_fn = semantic_fn
+
+ def train_one_step(self, batch):
+ waves, mels, wave_lengths, mel_input_length = batch
+
+ B = waves.size(0)
+ target_size = mels.size(2)
+ target = mels
+ target_lengths = mel_input_length
+
+ # get speaker embedding
+ if self.sr != 22050:
+ waves_22k = torchaudio.functional.resample(waves, self.sr, 22050)
+ wave_lengths_22k = (wave_lengths.float() * 22050 / self.sr).long()
+ else:
+ waves_22k = waves
+ wave_lengths_22k = wave_lengths
+ se_batch = self.tone_color_converter.extract_se(waves_22k, wave_lengths_22k)
+
+ ref_se_idx = torch.randint(0, len(self.se_db), (B,))
+ ref_se = self.se_db[ref_se_idx]
+ ref_se = ref_se.to(self.device)
+
+ # convert
+ converted_waves_22k = self.tone_color_converter.convert(waves_22k, wave_lengths_22k, se_batch, ref_se).squeeze(1)
+
+ if self.sr != 22050:
+ converted_waves = torchaudio.functional.resample(converted_waves_22k, 22050, self.sr)
+ else:
+ converted_waves = converted_waves_22k
+
+ waves_16k = torchaudio.functional.resample(waves, self.sr, 16000)
+ wave_lengths_16k = (wave_lengths.float() * 16000 / self.sr).long()
+ converted_waves_16k = torchaudio.functional.resample(converted_waves, self.sr, 16000)
+ # extract S_alt (perturbed speech tokens)
+ S_ori = self.semantic_fn(waves_16k)
+ S_alt = self.semantic_fn(converted_waves_16k)
+
+ if self.f0_condition:
+ F0_ori = self.rmvpe.infer_from_audio_batch(waves_16k)
+ else:
+ F0_ori = None
+ # interpolate speech token to match acoustic feature length
+ alt_cond, _, alt_codes, alt_commitment_loss, alt_codebook_loss = (
+ self.model.length_regulator(S_alt, ylens=target_lengths, f0=F0_ori))
+ ori_cond, _, ori_codes, ori_commitment_loss, ori_codebook_loss = (
+ self.model.length_regulator(S_ori, ylens=target_lengths, f0=F0_ori))
+ if alt_commitment_loss is None:
+ alt_commitment_loss = 0
+ alt_codebook_loss = 0
+ ori_commitment_loss = 0
+ ori_codebook_loss = 0
+
+ # randomly set a length as prompt
+ prompt_len_max = target_lengths - 1
+ prompt_len = (torch.rand([B], device=alt_cond.device) * prompt_len_max).floor().to(dtype=torch.long)
+ prompt_len[torch.rand([B], device=alt_cond.device) < 0.1] = 0
+
+ # for prompt cond token, it must be from ori_cond instead of alt_cond
+ cond = alt_cond.clone()
+ for bib in range(B):
+ cond[bib, :prompt_len[bib]] = ori_cond[bib, :prompt_len[bib]]
+
+ # diffusion target
+ common_min_len = min(target_size, cond.size(1))
+ target = target[:, :, :common_min_len]
+ cond = cond[:, :common_min_len]
+ target_lengths = torch.clamp(target_lengths, max=common_min_len)
+ x = target
+ # style vectors are extracted from prompt only to avoid inference time OOD
+ feat_list = []
+ for bib in range(B):
+ feat = kaldi.fbank(waves_16k[bib:bib + 1, :wave_lengths_16k[bib]],
+ num_mel_bins=80,
+ dither=0,
+ sample_frequency=16000)
+ feat = feat - feat.mean(dim=0, keepdim=True)
+ feat_list.append(feat)
+ max_feat_len = max([feat.size(0) for feat in feat_list])
+ feat_lens = torch.tensor([feat.size(0) for feat in feat_list], dtype=torch.int32).to(self.device) // 2
+ feat_list = [
+ torch.nn.functional.pad(feat, (0, 0, 0, max_feat_len - feat.size(0)), value=float(feat.min().item()))
+ for feat in feat_list
+ ]
+ y_list = []
+ with torch.no_grad():
+ for feat in feat_list:
+ y = self.sv_fn(feat.unsqueeze(0))
+ y_list.append(y)
+ y = torch.cat(y_list, dim=0)
+
+ loss, _ = self.model.cfm(x, target_lengths, prompt_len, cond, y)
+
+ loss_total = (loss +
+ (alt_commitment_loss + ori_commitment_loss) * 0.05 +
+ (ori_codebook_loss + alt_codebook_loss) * 0.15)
+
+ self.optimizer.zero_grad()
+ loss_total.backward()
+ grad_norm_g = torch.nn.utils.clip_grad_norm_(self.model.cfm.parameters(), 10.0)
+ grad_norm_g2 = torch.nn.utils.clip_grad_norm_(self.model.length_regulator.parameters(), 10.0)
+ self.optimizer.step('cfm')
+ self.optimizer.step('length_regulator')
+ self.optimizer.scheduler(key='cfm')
+ self.optimizer.scheduler(key='length_regulator')
+
+ return loss.detach().item()
+ def train_one_epoch(self):
+ _ = [self.model[key].train() for key in self.model]
+ for i, batch in enumerate(tqdm(self.train_dataloader)):
+ batch = [b.to(self.device) for b in batch]
+ loss = self.train_one_step(batch)
+ self.ema_loss = self.ema_loss * self.loss_smoothing_rate + loss * (1 - self.loss_smoothing_rate) if self.iters > 0 else loss
+ if self.iters % self.log_interval == 0:
+ print(f"epoch {self.epoch}, step {self.iters}, loss: {self.ema_loss}")
+ self.iters += 1
+ if self.iters >= self.max_steps:
+ break
+ if self.iters % self.save_interval == 0:
+ print('Saving..')
+ state = {
+ 'net': {key: self.model[key].state_dict() for key in self.model},
+ 'optimizer': self.optimizer.state_dict(),
+ 'scheduler': self.optimizer.scheduler_state_dict(),
+ 'iters': self.iters,
+ 'epoch': self.epoch,
+ }
+ save_path = os.path.join(self.log_dir, 'DiT_epoch_%05d_step_%05d.pth' % (self.epoch, self.iters))
+ torch.save(state, save_path)
+
+ # find all checkpoints and remove old ones
+ checkpoints = glob.glob(os.path.join(self.log_dir, 'DiT_epoch_*.pth'))
+ if len(checkpoints) > 2:
+ # sort by step
+ checkpoints.sort(key=lambda x: int(x.split('_')[-1].split('.')[0]))
+ for cp in checkpoints[:-2]:
+ os.remove(cp)
+
+ def train(self):
+ self.ema_loss = 0
+ self.loss_smoothing_rate = 0.99
+ for epoch in range(self.n_epochs):
+ self.epoch = epoch
+ self.train_one_epoch()
+ if self.iters >= self.max_steps:
+ break
+ print('Saving..')
+ state = {
+ 'net': {key: self.model[key].state_dict() for key in self.model},
+ }
+ os.makedirs(self.log_dir, exist_ok=True)
+ save_path = os.path.join(self.log_dir, 'ft_model.pth')
+ torch.save(state, save_path)
+
+def main(args):
+ trainer = Trainer(
+ config_path=args.config,
+ pretrained_ckpt_path=args.pretrained_ckpt,
+ data_dir=args.dataset_dir,
+ run_name=args.run_name,
+ batch_size=args.batch_size,
+ steps=args.max_steps,
+ max_epochs=args.max_epochs,
+ save_interval=args.save_every,
+ num_workers=args.num_workers,
+ )
+ trainer.train()
+
+if __name__ == '__main__':
+ parser = argparse.ArgumentParser()
+ parser.add_argument('--config', type=str, default='./configs/presets/config_dit_mel_seed_uvit_xlsr_tiny.yml')
+ parser.add_argument('--pretrained-ckpt', type=str, default=None)
+ parser.add_argument('--dataset-dir', type=str, default='/path/to/dataset')
+ parser.add_argument('--run-name', type=str, default='my_run')
+ parser.add_argument('--batch-size', type=int, default=2)
+ parser.add_argument('--max-steps', type=int, default=1000)
+ parser.add_argument('--max-epochs', type=int, default=1000)
+ parser.add_argument('--save-every', type=int, default=500)
+ parser.add_argument('--num-workers', type=int, default=0)
+ args = parser.parse_args()
+ main(args)
\ No newline at end of file