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') + ] + + return gr.Interface( + fn=convert_voice_v1_wrapper, + description=description, + inputs=inputs, + outputs=outputs, + title="Seed Voice Conversion V1 (Voice & Singing Voice Conversion)", + examples=examples, + cache_examples=False, + ) + + +def create_v2_interface(): + # Set up Gradio interface + description = ( + "Zero-shot voice/style 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.
" + "Please click the 'convert style/emotion/accent' checkbox to convert the style, emotion, or accent of the source audio, or else only timbre conversion will be performed.
" + "Click the 'anonymization only' checkbox will ignore reference audio but convert source to an 'average voice' determined by model itself.
" + "无需训练的 zero-shot 语音/口音转换模型，若需本地部署查看[GitHub页面](https://github.com/Plachtaa/seed-vc)
" + "请注意，参考音频若超过 25 秒，则会被自动裁剪至此长度。
若源音频和参考音频的总时长超过 30 秒，源音频将被分段处理。" + "
请勾选 'convert style/emotion/accent' 以转换源音频的风格、情感或口音，否则仅执行音色转换。
" + "勾选 'anonymization only' 会无视参考音频而将源音频转换为某种由模型自身决定的 '平均音色'。
" + + "Credits to [Vevo](https://github.com/open-mmlab/Amphion/tree/main/models/vc/vevo)" + ) + inputs = [ + gr.Audio(type="filepath", label="Source Audio / 源音频"), + gr.Audio(type="filepath", label="Reference Audio / 参考音频"), + gr.Slider(minimum=1, maximum=200, value=30, step=1, label="Diffusion Steps / 扩散步数", + info="30 by default, 50~100 for best quality / 默认为 30，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.0, label="Intelligibility CFG Rate", + info="controls pronunciation intelligibility / 控制发音清晰度"), + gr.Slider(minimum=0.0, maximum=1.0, step=0.1, value=0.7, label="Similarity CFG Rate", + info="controls similarity to reference audio / 控制与参考音频的相似度"), + gr.Slider(minimum=0.1, maximum=1.0, step=0.1, value=0.9, label="Top-p", + info="AR model sampling top P"), + gr.Slider(minimum=0.1, maximum=2.0, step=0.1, value=1.0, label="Temperature", + info="AR model sampling temperature"), + gr.Slider(minimum=1.0, maximum=3.0, step=0.1, value=1.0, label="Repetition Penalty", + info="AR model sampling repetition penalty"), + gr.Checkbox(label="convert style/emotion/accent", value=False), + gr.Checkbox(label="anonymization only", value=False), + ] + + examples = [ + ["examples/source/yae_0.wav", "examples/reference/dingzhen_0.wav", 50, 1.0, 0.0, 0.7, 0.9, 1.0, 1.0, False, + False], + ["examples/source/jay_0.wav", "examples/reference/azuma_0.wav", 50, 1.0, 0.0, 0.7, 0.9, 1.0, 1.0, False, False], + ] + + outputs = [ + gr.Audio(label="Stream Output Audio / 流式输出", streaming=True, format='mp3'), + gr.Audio(label="Full Output Audio / 完整输出", streaming=False, format='wav') + ] + + return gr.Interface( + fn=convert_voice_v2_wrapper, + description=description, + inputs=inputs, + outputs=outputs, + title="Seed Voice Conversion V2 (Voice & Style Conversion)", + examples=examples, + cache_examples=False, + ) + + +def main(args): + global vc_wrapper_v1, vc_wrapper_v2 + # Create interfaces based on enabled versions + interfaces = [] + + # Load V2 models if enabled + if args.enable_v2: + print("Loading V2 models...") + vc_wrapper_v2 = load_v2_models(args) + v2_interface = create_v2_interface() + interfaces.append(("V2 - Voice & Style Conversion", v2_interface)) + + # Create V1 interface if enabled + if args.enable_v1: + print("Creating V1 interface...") + v1_interface = create_v1_interface() + interfaces.append(("V1 - Voice & Singing Voice Conversion", v1_interface)) + + # Check if at least one version is enabled + if not interfaces: + print("Error: At least one version (V1 or V2) must be enabled.") + return + + # Create tabs + with gr.Blocks(title="Seed Voice Conversion") as demo: + gr.Markdown("# Seed Voice Conversion") + + if len(interfaces) > 1: + gr.Markdown("Choose between V1 (Voice & Singing Voice Conversion) or V2 (Voice & Style Conversion)") + + with gr.Tabs(): + for tab_name, interface in interfaces: + with gr.TabItem(tab_name): + interface.render() + else: + # If only one version is enabled, don't use tabs + for _, interface in interfaces: + interface.render() + + # Launch the combined interface + demo.launch() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--compile", action="store_true", help="Compile the model using torch.compile") + parser.add_argument("--enable-v1", action="store_true", + help="Enable V1 (Voice & Singing Voice Conversion)") + parser.add_argument("--enable-v2", action="store_true", + help="Enable V2 (Voice & Style Conversion)") + args = parser.parse_args() + main(args) \ No newline at end of file diff --git a/seed-vc/app_svc.py b/seed-vc/app_svc.py new file mode 100644 index 0000000000000000000000000000000000000000..c34d8bc3e832c8869691ef17adc042680a879fb1 --- /dev/null +++ b/seed-vc/app_svc.py @@ -0,0 +1,450 @@ +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 + +fp16 = False +device = None +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 is None or args.checkpoint == "": + 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}") + dit_checkpoint_path = args.checkpoint + dit_config_path = args.config + 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) + + if device.type == "mps": + F0_ori = torch.from_numpy(F0_ori).float().to(device)[None] + F0_alt = torch.from_numpy(F0_alt).float().to(device)[None] + else: + 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", type=str, help="Path to the checkpoint file", default=None) + parser.add_argument("--config", 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) + parser.add_argument("--gpu", type=int, help="Which GPU id to use", default=0) + args = parser.parse_args() + cuda_target = f"cuda:{args.gpu}" if args.gpu else "cuda" + + if torch.cuda.is_available(): + device = torch.device(cuda_target) + elif torch.backends.mps.is_available(): + device = torch.device("mps") + else: + device = torch.device("cpu") + main(args) diff --git a/seed-vc/app_vc.py b/seed-vc/app_vc.py new file mode 100644 index 0000000000000000000000000000000000000000..92e82d6181e8a5362407612f571a8679e6720876 --- /dev/null +++ b/seed-vc/app_vc.py @@ -0,0 +1,399 @@ +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 +fp16 = False +device = None +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 is None or args.checkpoint == "": + 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 + dit_config_path = args.config + 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", type=str, help="Path to the checkpoint file", default=None) + parser.add_argument("--config", 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) + parser.add_argument("--gpu", type=int, help="Which GPU id to use", default=0) + args = parser.parse_args() + cuda_target = f"cuda:{args.gpu}" if args.gpu else "cuda" + + if torch.cuda.is_available(): + device = torch.device(cuda_target) + elif torch.backends.mps.is_available(): + device = torch.device("mps") + else: + device = torch.device("cpu") + main(args) \ No newline at end of file diff --git a/seed-vc/app_vc_v2.py b/seed-vc/app_vc_v2.py new file mode 100644 index 0000000000000000000000000000000000000000..4922d93f41da860d5b538d135a2104d9e150afec --- /dev/null +++ b/seed-vc/app_vc_v2.py @@ -0,0 +1,99 @@ +import gradio as gr +import torch +import yaml + +if torch.cuda.is_available(): + device = torch.device("cuda") +elif torch.backends.mps.is_available(): + device = torch.device("mps") +else: + device = torch.device("cpu") + +dtype = torch.float16 +def load_models(args): + from hydra.utils import instantiate + from omegaconf import DictConfig + cfg = DictConfig(yaml.safe_load(open("configs/v2/vc_wrapper.yaml", "r"))) + vc_wrapper = instantiate(cfg) + vc_wrapper.load_checkpoints(ar_checkpoint_path=args.ar_checkpoint_path, + cfm_checkpoint_path=args.cfm_checkpoint_path) + vc_wrapper.to(device) + vc_wrapper.eval() + + vc_wrapper.setup_ar_caches(max_batch_size=1, max_seq_len=4096, dtype=dtype, device=device) + + if args.compile: + torch._inductor.config.coordinate_descent_tuning = True + torch._inductor.config.triton.unique_kernel_names = True + + if hasattr(torch._inductor.config, "fx_graph_cache"): + # Experimental feature to reduce compilation times, will be on by default in future + torch._inductor.config.fx_graph_cache = True + vc_wrapper.compile_ar() + # vc_wrapper.compile_cfm() + + return vc_wrapper + +def main(args): + vc_wrapper = load_models(args) + + # Set up Gradio interface + 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=30, step=1, label="Diffusion Steps / 扩散步数", + info="30 by default, 50~100 for best quality / 默认为 30，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.5, label="Intelligibility CFG Rate", + info="has subtle influence / 有微小影响"), + gr.Slider(minimum=0.0, maximum=1.0, step=0.1, value=0.5, label="Similarity CFG Rate", + info="has subtle influence / 有微小影响"), + gr.Slider(minimum=0.1, maximum=1.0, step=0.1, value=0.9, label="Top-p", + info="Controls diversity of generated audio / 控制生成音频的多样性"), + gr.Slider(minimum=0.1, maximum=2.0, step=0.1, value=1.0, label="Temperature", + info="Controls randomness of generated audio / 控制生成音频的随机性"), + gr.Slider(minimum=1.0, maximum=3.0, step=0.1, value=1.0, label="Repetition Penalty", + info="Penalizes repetition in generated audio / 惩罚生成音频中的重复"), + gr.Checkbox(label="convert style", value=False), + gr.Checkbox(label="anonymization only", value=False), + ] + + examples = [ + ["examples/source/yae_0.wav", "examples/reference/dingzhen_0.wav", 50, 1.0, 0.5, 0.5, 0.9, 1.0, 1.0, False, False], + ["examples/source/jay_0.wav", "examples/reference/azuma_0.wav", 50, 1.0, 0.5, 0.5, 0.9, 1.0, 1.0, False, False], + ] + + outputs = [ + gr.Audio(label="Stream Output Audio / 流式输出", streaming=True, format='mp3'), + gr.Audio(label="Full Output Audio / 完整输出", streaming=False, format='wav') + ] + + # Launch the Gradio interface + gr.Interface( + fn=vc_wrapper.convert_voice_with_streaming, + description=description, + inputs=inputs, + outputs=outputs, + title="Seed Voice Conversion V2", + examples=examples, + cache_examples=False, + ).launch() + +if __name__ == "__main__": + import argparse + parser = argparse.ArgumentParser() + parser.add_argument("--compile", action="store_true", help="Compile the model using torch.compile") + # V2 custom checkpoints + parser.add_argument("--ar-checkpoint-path", type=str, default=None, + help="Path to custom checkpoint file") + parser.add_argument("--cfm-checkpoint-path", type=str, default=None, + help="Path to custom checkpoint file") + args = parser.parse_args() + main(args) \ No newline at end of file diff --git a/seed-vc/conda-nix-vc-py310.yaml b/seed-vc/conda-nix-vc-py310.yaml new file mode 100644 index 0000000000000000000000000000000000000000..bd59a8ce7f84966e9fd47360fb7af36cc63e3b27 --- /dev/null +++ b/seed-vc/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/seed-vc/configs/astral_quantization/default_2048.yml b/seed-vc/configs/astral_quantization/default_2048.yml new file mode 100644 index 0000000000000000000000000000000000000000..8045f2dfdbfa7e31223cf2e086921e0063944410 --- /dev/null +++ b/seed-vc/configs/astral_quantization/default_2048.yml @@ -0,0 +1,40 @@ +_target_: modules.astral_quantization.default_model.AstralQuantizer +tokenizer_name: "openai/whisper-small" +ssl_model_name: "facebook/hubert-large-ll60k" +ssl_output_layer: 18 +encoder: + _target_: modules.astral_quantization.convnext.ConvNeXtV2Stage + dim: 512 + num_blocks: 12 + intermediate_dim: 1536 + dilation: 1 + input_dim: 1024 +quantizer: + _target_: modules.astral_quantization.bsq.BinarySphericalQuantize + codebook_size: 2048 # codebook size, must be a power of 2 + dim: 512 + entropy_loss_weight: 0.1 + diversity_gamma: 1.0 + spherical: True + enable_entropy_loss: True + soft_entropy_loss: True +decoder: + _target_: modules.astral_quantization.convnext.ConvNeXtV2Stage + dim: 512 + num_blocks: 12 + intermediate_dim: 1536 + dilation: 1 + output_dim: 1024 + gin_channels: 192 +asr_decoder: + _target_: modules.astral_quantization.asr_decoder.ASRDecoder + hidden_dim: 768 + num_heads: 12 + depth: 12 + block_size: 4096 + in_channels: 512 + n_vocab: 51866 + bos_id: 50528 + eos_id: 50527 + dropout_rate: 0.0 + attn_dropout_rate: 0.0 \ No newline at end of file diff --git a/seed-vc/configs/astral_quantization/default_32.yml b/seed-vc/configs/astral_quantization/default_32.yml new file mode 100644 index 0000000000000000000000000000000000000000..7da34ff4a34c6ae5e84b866fdff2d8d4c0a5865c --- /dev/null +++ b/seed-vc/configs/astral_quantization/default_32.yml @@ -0,0 +1,40 @@ +_target_: default_model.AstralQuantizer +tokenizer_name: "openai/whisper-small" +ssl_model_name: "facebook/hubert-large-ll60k" +ssl_output_layer: 18 +encoder: + _target_: modules.convnext.ConvNeXtV2Stage + dim: 512 + num_blocks: 12 + intermediate_dim: 1536 + dilation: 1 + input_dim: 1024 +quantizer: + _target_: modules.bsq.BinarySphericalQuantize + codebook_size: 32 # codebook size, must be a power of 2 + dim: 512 + entropy_loss_weight: 0.1 + diversity_gamma: 1.0 + spherical: True + enable_entropy_loss: True + soft_entropy_loss: True +decoder: + _target_: modules.convnext.ConvNeXtV2Stage + dim: 512 + num_blocks: 12 + intermediate_dim: 1536 + dilation: 1 + output_dim: 1024 + gin_channels: 192 +asr_decoder: + _target_: modules.asr_decoder.ASRDecoder + hidden_dim: 768 + num_heads: 12 + depth: 12 + block_size: 4096 + in_channels: 512 + n_vocab: 51866 + bos_id: 50528 + eos_id: 50527 + dropout_rate: 0.0 + attn_dropout_rate: 0.0 \ No newline at end of file diff --git a/seed-vc/configs/config.json b/seed-vc/configs/config.json new file mode 100644 index 0000000000000000000000000000000000000000..e74f0b4898f6e47e1d198b62cdba989784ce2bb0 --- /dev/null +++ b/seed-vc/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/seed-vc/configs/hifigan.yml b/seed-vc/configs/hifigan.yml new file mode 100644 index 0000000000000000000000000000000000000000..6faee66d988dc54bcc5f747d048c5f7e2a8b759b --- /dev/null +++ b/seed-vc/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/seed-vc/configs/presets/config_dit_mel_seed_uvit_whisper_base_f0_44k.yml b/seed-vc/configs/presets/config_dit_mel_seed_uvit_whisper_base_f0_44k.yml new file mode 100644 index 0000000000000000000000000000000000000000..150560475fd7dedad8da76e87480dd36c85ee07f --- /dev/null +++ b/seed-vc/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/seed-vc/configs/presets/config_dit_mel_seed_uvit_whisper_small_wavenet.yml b/seed-vc/configs/presets/config_dit_mel_seed_uvit_whisper_small_wavenet.yml new file mode 100644 index 0000000000000000000000000000000000000000..7fc8ec4cb659e9d4856e7b9fa1621ba20796fb08 --- /dev/null +++ b/seed-vc/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/seed-vc/configs/presets/config_dit_mel_seed_uvit_xlsr_tiny.yml b/seed-vc/configs/presets/config_dit_mel_seed_uvit_xlsr_tiny.yml new file mode 100644 index 0000000000000000000000000000000000000000..d1870d89bb8979fca9e89a29a4c5ebb96a9eceb6 --- /dev/null +++ b/seed-vc/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/seed-vc/configs/v2/vc_wrapper.yaml b/seed-vc/configs/v2/vc_wrapper.yaml new file mode 100644 index 0000000000000000000000000000000000000000..9d254d516f969f3fff27f8477fc5026074a2a59e --- /dev/null +++ b/seed-vc/configs/v2/vc_wrapper.yaml @@ -0,0 +1,105 @@ +_target_: modules.v2.vc_wrapper.VoiceConversionWrapper +sr: 22050 +hop_size: 256 +mel_fn: + _target_: modules.audio.mel_spectrogram + _partial_: true + n_fft: 1024 + win_size: 1024 + hop_size: 256 + num_mels: 80 + sampling_rate: 22050 + fmin: 0 + fmax: null + center: False +cfm: + _target_: modules.v2.cfm.CFM + estimator: + _target_: modules.v2.dit_wrapper.DiT + time_as_token: true + style_as_token: true + uvit_skip_connection: false + block_size: 8192 + depth: 13 + num_heads: 8 + hidden_dim: 512 + in_channels: 80 + content_dim: 512 + style_encoder_dim: 192 + class_dropout_prob: 0.1 + dropout_rate: 0.0 + attn_dropout_rate: 0.0 +cfm_length_regulator: + _target_: modules.v2.length_regulator.InterpolateRegulator + channels: 512 + is_discrete: true + codebook_size: 2048 + sampling_ratios: [ 1, 1, 1, 1 ] + f0_condition: false +ar: + _target_: modules.v2.ar.NaiveWrapper + model: + _target_: modules.v2.ar.NaiveTransformer + config: + _target_: modules.v2.ar.NaiveModelArgs + dropout: 0.0 + rope_base: 10000.0 + dim: 768 + head_dim: 64 + n_local_heads: 2 + intermediate_size: 2304 + n_head: 12 + n_layer: 12 + vocab_size: 2049 # 1 + 1 for eos +ar_length_regulator: + _target_: modules.v2.length_regulator.InterpolateRegulator + channels: 768 + is_discrete: true + codebook_size: 32 + sampling_ratios: [ ] + f0_condition: false +style_encoder: + _target_: modules.campplus.DTDNN.CAMPPlus + feat_dim: 80 + embedding_size: 192 +content_extractor_narrow: + _target_: modules.astral_quantization.default_model.AstralQuantizer + tokenizer_name: "openai/whisper-small" + ssl_model_name: "facebook/hubert-large-ll60k" + ssl_output_layer: 18 + skip_ssl: true + encoder: &bottleneck_encoder + _target_: modules.astral_quantization.convnext.ConvNeXtV2Stage + dim: 512 + num_blocks: 12 + intermediate_dim: 1536 + dilation: 1 + input_dim: 1024 + quantizer: + _target_: modules.astral_quantization.bsq.BinarySphericalQuantize + codebook_size: 32 # codebook size, must be a power of 2 + dim: 512 + entropy_loss_weight: 0.1 + diversity_gamma: 1.0 + spherical: True + enable_entropy_loss: True + soft_entropy_loss: True +content_extractor_wide: + _target_: modules.astral_quantization.default_model.AstralQuantizer + tokenizer_name: "openai/whisper-small" + ssl_model_name: "facebook/hubert-large-ll60k" + ssl_output_layer: 18 + encoder: *bottleneck_encoder + quantizer: + _target_: modules.astral_quantization.bsq.BinarySphericalQuantize + codebook_size: 2048 # codebook size, must be a power of 2 + dim: 512 + entropy_loss_weight: 0.1 + diversity_gamma: 1.0 + spherical: True + enable_entropy_loss: True + soft_entropy_loss: True +vocoder: + _target_: modules.bigvgan.bigvgan.BigVGAN.from_pretrained + pretrained_model_name_or_path: "nvidia/bigvgan_v2_22khz_80band_256x" + use_cuda_kernel: false diff --git a/seed-vc/data/ft_dataset.py b/seed-vc/data/ft_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..fac88657935006e79fa7b95b4fde4fff191f5b93 --- /dev/null +++ b/seed-vc/data/ft_dataset.py @@ -0,0 +1,126 @@ +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 +def to_mel_fn(wave, mel_fn_args): + return mel_spectrogram(wave, **mel_fn_args) + +class FT_Dataset(torch.utils.data.Dataset): + def __init__( + self, + data_path, + spect_params, + sr=22050, + batch_size=1, + ): + self.data_path = data_path + self.data = [] + for root, _, files in os.walk(data_path): + for file in files: + if file.endswith((".wav", ".mp3", ".flac", ".ogg", ".m4a", ".opus")): + self.data.append(os.path.join(root, file)) + + self.sr = sr + self.mel_fn_args = { + "n_fft": spect_params['n_fft'], + "win_size": spect_params.get('win_length', spect_params.get('win_size', 1024)), + "hop_size": spect_params.get('hop_length', spect_params.get('hop_size', 256)), + "num_mels": spect_params.get('n_mels', spect_params.get('num_mels', 80)), + "sampling_rate": sr, + "fmin": spect_params['fmin'], + "fmax": None if spect_params['fmax'] == "None" else spect_params['fmax'], + "center": False + } + + assert len(self.data) != 0 + 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))) + if orig_sr != self.sr: + speech = librosa.resample(speech, orig_sr, self.sr) + + wave = torch.from_numpy(speech).float().unsqueeze(0) + mel = to_mel_fn(wave, self.mel_fn_args).squeeze(0) + + return wave.squeeze(0), 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 diff --git a/seed-vc/eval.py b/seed-vc/eval.py new file mode 100644 index 0000000000000000000000000000000000000000..117ca7c77aaa35de4593e344599334ef1ea5a3f6 --- /dev/null +++ b/seed-vc/eval.py @@ -0,0 +1,556 @@ +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 + +if torch.cuda.is_available(): + device = torch.device("cuda") +elif torch.backends.mps.is_available(): + device = torch.device("mps") +else: + device = torch.device("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/seed-vc/hf_utils.py b/seed-vc/hf_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..4bd88fa6be55f3eaf4d2cd434bb8757a61cdab5f --- /dev/null +++ b/seed-vc/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=None): + 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/seed-vc/inference.py b/seed-vc/inference.py new file mode 100644 index 0000000000000000000000000000000000000000..84d673b64db2b5349039ff96702533f4be0a4e25 --- /dev/null +++ b/seed-vc/inference.py @@ -0,0 +1,425 @@ +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") +if torch.cuda.is_available(): + device = torch.device("cuda") +elif torch.backends.mps.is_available(): + device = torch.device("mps") +else: + device = torch.device("cpu") + +fp16 = False +def load_models(args): + global fp16 + fp16 = args.fp16 + if not args.f0_condition: + if args.checkpoint is None: + 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 + dit_config_path = args.config + f0_fn = None + else: + if args.checkpoint 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 + dit_config_path = args.config + # 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", type=str, help="Path to the checkpoint file", default=None) + parser.add_argument("--config", 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/seed-vc/inference_v2.py b/seed-vc/inference_v2.py new file mode 100644 index 0000000000000000000000000000000000000000..51c02ecc3512fc20b41ba850b8ee6dc437071daf --- /dev/null +++ b/seed-vc/inference_v2.py @@ -0,0 +1,144 @@ +import os +import argparse +import torch +import yaml +import soundfile as sf +import time +from modules.commons import str2bool + +# Set up device and torch configurations +if torch.cuda.is_available(): + device = torch.device("cuda") +elif torch.backends.mps.is_available(): + device = torch.device("mps") +else: + device = torch.device("cpu") + +dtype = torch.float16 + +# Global variables to store model instances +vc_wrapper_v2 = None + + +def load_v2_models(args): + """Load V2 models using the wrapper from app.py""" + from hydra.utils import instantiate + from omegaconf import DictConfig + cfg = DictConfig(yaml.safe_load(open("configs/v2/vc_wrapper.yaml", "r"))) + vc_wrapper = instantiate(cfg) + vc_wrapper.load_checkpoints(ar_checkpoint_path=args.ar_checkpoint_path, + cfm_checkpoint_path=args.cfm_checkpoint_path) + vc_wrapper.to(device) + vc_wrapper.eval() + + vc_wrapper.setup_ar_caches(max_batch_size=1, max_seq_len=4096, dtype=dtype, device=device) + + if args.compile: + torch._inductor.config.coordinate_descent_tuning = True + torch._inductor.config.triton.unique_kernel_names = True + + if hasattr(torch._inductor.config, "fx_graph_cache"): + # Experimental feature to reduce compilation times, will be on by default in future + torch._inductor.config.fx_graph_cache = True + vc_wrapper.compile_ar() + # vc_wrapper.compile_cfm() + + return vc_wrapper + + +def convert_voice_v2(source_audio_path, target_audio_path, args): + """Convert voice using V2 model""" + global vc_wrapper_v2 + if vc_wrapper_v2 is None: + vc_wrapper_v2 = load_v2_models(args) + + # Use the generator function but collect all outputs + generator = vc_wrapper_v2.convert_voice_with_streaming( + source_audio_path=source_audio_path, + target_audio_path=target_audio_path, + diffusion_steps=args.diffusion_steps, + length_adjust=args.length_adjust, + intelligebility_cfg_rate=args.intelligibility_cfg_rate, + similarity_cfg_rate=args.similarity_cfg_rate, + top_p=args.top_p, + temperature=args.temperature, + repetition_penalty=args.repetition_penalty, + convert_style=args.convert_style, + anonymization_only=args.anonymization_only, + device=device, + dtype=dtype, + stream_output=True + ) + + # Collect all outputs from the generator + for output in generator: + _, full_audio = output + return full_audio + + +def main(args): + # Create output directory if it doesn't exist + os.makedirs(args.output, exist_ok=True) + + start_time = time.time() + converted_audio = convert_voice_v2(args.source, args.target, args) + end_time = time.time() + + if converted_audio is None: + print("Error: Failed to convert voice") + return + + # Save the converted audio + source_name = os.path.basename(args.source).split(".")[0] + target_name = os.path.basename(args.target).split(".")[0] + + # Create a descriptive filename + filename = f"vc_v2_{source_name}_{target_name}_{args.length_adjust}_{args.diffusion_steps}_{args.similarity_cfg_rate}.wav" + + output_path = os.path.join(args.output, filename) + save_sr, converted_audio = converted_audio + sf.write(output_path, converted_audio, save_sr) + + print(f"Voice conversion completed in {end_time - start_time:.2f} seconds") + print(f"Output saved to: {output_path}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Voice Conversion Inference Script") + parser.add_argument("--source", type=str, required=True, + help="Path to source audio file") + parser.add_argument("--target", type=str, required=True, + help="Path to target/reference audio file") + parser.add_argument("--output", type=str, default="./output", + help="Output directory for converted audio") + parser.add_argument("--diffusion-steps", type=int, default=30, + help="Number of diffusion steps") + parser.add_argument("--length-adjust", type=float, default=1.0, + help="Length adjustment factor (<1.0 for speed-up, >1.0 for slow-down)") + parser.add_argument("--compile", type=bool, default=False, + help="Whether to compile the model for faster inference") + + # V2 specific arguments + parser.add_argument("--intelligibility-cfg-rate", type=float, default=0.7, + help="Intelligibility CFG rate for V2 model") + parser.add_argument("--similarity-cfg-rate", type=float, default=0.7, + help="Similarity CFG rate for V2 model") + parser.add_argument("--top-p", type=float, default=0.9, + help="Top-p sampling parameter for V2 model") + parser.add_argument("--temperature", type=float, default=1.0, + help="Temperature sampling parameter for V2 model") + parser.add_argument("--repetition-penalty", type=float, default=1.0, + help="Repetition penalty for V2 model") + parser.add_argument("--convert-style", type=str2bool, default=False, + help="Convert style/emotion/accent for V2 model") + parser.add_argument("--anonymization-only", type=str2bool, default=False, + help="Anonymization only mode for V2 model") + + # V2 custom checkpoints + parser.add_argument("--ar-checkpoint-path", type=str, default=None, + help="Path to custom checkpoint file") + parser.add_argument("--cfm-checkpoint-path", type=str, default=None, + help="Path to custom checkpoint file") + + args = parser.parse_args() + main(args) \ No newline at end of file diff --git a/seed-vc/modules/astral_quantization/bsq.py b/seed-vc/modules/astral_quantization/bsq.py new file mode 100644 index 0000000000000000000000000000000000000000..239d0caaf370afc77ba62961563903ad675b6d98 --- /dev/null +++ b/seed-vc/modules/astral_quantization/bsq.py @@ -0,0 +1,569 @@ +""" +Lookup Free Quantization +Proposed in https://arxiv.org/abs/2310.05737 + +In the simplest setup, each dimension is quantized into {-1, 1}. +An entropy penalty is used to encourage utilization. +""" + +from math import log2, ceil +from functools import partial, cache +from collections import namedtuple +from contextlib import nullcontext + +import torch.distributed as dist +from torch.distributed import nn as dist_nn + +import torch +from torch import nn, einsum +import torch.nn.functional as F +from torch.nn import Module +from torch.amp import autocast + +from einops import rearrange, reduce, pack, unpack + +# constants + +Return = namedtuple('Return', ['quantized', 'indices', 'entropy_aux_loss']) + +LossBreakdown = namedtuple('LossBreakdown', ['per_sample_entropy', 'batch_entropy', 'commitment']) + +# distributed helpers + +@cache +def is_distributed(): + return dist.is_initialized() and dist.get_world_size() > 1 + +def maybe_distributed_mean(t): + if not is_distributed(): + return t + + dist_nn.all_reduce(t) + t = t / dist.get_world_size() + return t + +# helper functions + +def exists(v): + return v is not None + +def identity(t): + return t + +def default(*args): + for arg in args: + if exists(arg): + return arg() if callable(arg) else arg + return None + +def pack_one(t, pattern): + return pack([t], pattern) + +def unpack_one(t, ps, pattern): + return unpack(t, ps, pattern)[0] + +def l2norm(t): + return F.normalize(t, dim = -1) + +# entropy + +def log(t, eps = 1e-5): + return t.clamp(min = eps).log() + +def entropy(prob): + return (-prob * log(prob)).sum(dim=-1) + +# cosine sim linear + +class CosineSimLinear(Module): + def __init__( + self, + dim_in, + dim_out, + scale = 1. + ): + super().__init__() + self.scale = scale + self.weight = nn.Parameter(torch.randn(dim_in, dim_out)) + + def forward(self, x): + x = F.normalize(x, dim = -1) + w = F.normalize(self.weight, dim = 0) + return (x @ w) * self.scale + +def soft_entropy_loss(u, tau=1.0, gamma=1.0): + """ + Compute the soft entropy loss for Binary Spherical Quantization (BSQ). + + Args: + u (torch.Tensor): Input latent embeddings of shape (batch_size, L). + tau (float): Temperature scaling factor. + gamma (float): Weight for the second entropy term. + + Returns: + torch.Tensor: Soft entropy loss. + """ + # Binary quantization: Generate implicit codebook corners + L = u.size(1) # Dimensionality of codebook + corners = torch.tensor([-1.0, 1.0], device=u.device) / (L**0.5) + + # Compute soft quantization probabilities for all dimensions + # q_hat(c|u) for each dimension + prob_matrix = torch.sigmoid(2 * tau * corners.unsqueeze(1) * u.unsqueeze(2)) # Shape: (batch_size, L, 2) + + # Entropy of q_hat(c|u) (independent along each dimension) + entropy_per_dim = -torch.sum(prob_matrix * prob_matrix.log(), dim=-1) # Shape: (batch_size, L) + entropy_term1 = entropy_per_dim.mean() + + # Expected probabilities for dataset entropy (approximation) + expected_probs = prob_matrix.mean(dim=0) # Mean across batch, shape: (L, 2) + entropy_term2 = -torch.sum(expected_probs * expected_probs.log(), dim=-1).mean() + + # Final entropy loss + loss = entropy_term1 - gamma * entropy_term2 + return loss + +# class + +class BinarySphericalQuantize(Module): + def __init__( + self, + *, + dim = None, + codebook_size = None, + entropy_loss_weight = 0.1, + commitment_loss_weight = 0., + diversity_gamma = 1., + straight_through_activation = nn.Identity(), + num_codebooks = 1, + keep_num_codebooks_dim = None, + codebook_scale = 1., # for residual LFQ, codebook scaled down by 2x at each layer + frac_per_sample_entropy = 0.25, # make less than 1. to only use a random fraction of the probs for per sample entropy + has_projections = None, + projection_has_bias = True, + soft_clamp_input_value = None, + cosine_sim_project_in = False, + cosine_sim_project_in_scale = None, + channel_first = None, + experimental_softplus_entropy_loss = False, + entropy_loss_offset = 5., # how much to shift the loss before softplus + spherical = True, # from https://arxiv.org/abs/2406.07548 + force_quantization_f32 = True, # will force the quantization step to be full precision + enable_entropy_loss = True, + soft_entropy_loss = True, + ): + super().__init__() + + # some assert validations + + assert exists(dim) or exists(codebook_size), 'either dim or codebook_size must be specified for LFQ' + assert not exists(codebook_size) or log2(codebook_size).is_integer(), f'your codebook size must be a power of 2 for lookup free quantization (suggested {2 ** ceil(log2(codebook_size))})' + + codebook_size = default(codebook_size, lambda: 2 ** dim) + self.codebook_size = codebook_size + + codebook_dim = int(log2(codebook_size)) + codebook_dims = codebook_dim * num_codebooks + dim = default(dim, codebook_dims) + + has_projections = default(has_projections, dim != codebook_dims) + + if cosine_sim_project_in: + cosine_sim_project_in = default(cosine_sim_project_in_scale, codebook_scale) + project_in_klass = partial(CosineSimLinear, scale = cosine_sim_project_in) + else: + project_in_klass = partial(nn.Linear, bias = projection_has_bias) + + self.project_in = project_in_klass(dim, codebook_dims) if has_projections else nn.Identity() + self.project_out = nn.Linear(codebook_dims, dim, bias = projection_has_bias) if has_projections else nn.Identity() + self.has_projections = has_projections + + self.dim = dim + self.codebook_dim = codebook_dim + self.num_codebooks = num_codebooks + + keep_num_codebooks_dim = default(keep_num_codebooks_dim, num_codebooks > 1) + assert not (num_codebooks > 1 and not keep_num_codebooks_dim) + self.keep_num_codebooks_dim = keep_num_codebooks_dim + + # channel first + + self.channel_first = channel_first + + # straight through activation + + self.activation = straight_through_activation + + # whether to use BSQ (binary spherical quantization) + + self.spherical = spherical + self.maybe_l2norm = (lambda t: l2norm(t) * self.codebook_scale) if spherical else identity + + # entropy aux loss related weights + + assert 0 < frac_per_sample_entropy <= 1. + self.frac_per_sample_entropy = frac_per_sample_entropy + + self.diversity_gamma = diversity_gamma + self.entropy_loss_weight = entropy_loss_weight + + # codebook scale + + self.codebook_scale = codebook_scale + + # commitment loss + + self.commitment_loss_weight = commitment_loss_weight + + # whether to soft clamp the input value from -value to value + + self.soft_clamp_input_value = soft_clamp_input_value + assert not exists(soft_clamp_input_value) or soft_clamp_input_value >= codebook_scale + + # whether to make the entropy loss positive through a softplus (experimental, please report if this worked or not in discussions) + + self.entropy_loss_offset = entropy_loss_offset + self.experimental_softplus_entropy_loss = experimental_softplus_entropy_loss + + # for no auxiliary loss, during inference + + self.register_buffer('mask', 2 ** torch.arange(codebook_dim - 1, -1, -1)) + self.register_buffer('zero', torch.tensor(0.), persistent = False) + + # whether to force quantization step to be f32 + + self.force_quantization_f32 = force_quantization_f32 + + # codes + self.enable_entropy_loss = enable_entropy_loss + self.soft_entropy_loss = soft_entropy_loss + if codebook_size <= 100000: + all_codes = torch.arange(codebook_size) + bits = ((all_codes[..., None].int() & self.mask) != 0).float() + codebook = self.bits_to_codes(bits) + + self.register_buffer('codebook', codebook.float(), persistent = False) + else: + all_codes = torch.arange(pow(2, 16)) + mask = 2 ** torch.arange(16 - 1, -1, -1) + bits = ((all_codes[..., None].int() & mask) != 0).float() + codebook = self.bits_to_codes(bits) + + self.register_buffer('codebook', codebook.float(), persistent = False) + + def bits_to_codes(self, bits): + return bits * self.codebook_scale * 2 - self.codebook_scale + + @property + def dtype(self): + return self.codebook.dtype + + def indices_to_codes( + self, + indices, + project_out = True + ): + is_img_or_video = indices.ndim >= (3 + int(self.keep_num_codebooks_dim)) + should_transpose = default(self.channel_first, is_img_or_video) + + if not self.keep_num_codebooks_dim: + indices = rearrange(indices, '... -> ... 1') + + # indices to codes, which are bits of either -1 or 1 + + bits = ((indices[..., None].int() & self.mask) != 0).to(self.dtype) + + codes = self.bits_to_codes(bits) + + codes = self.maybe_l2norm(codes) + + codes = rearrange(codes, '... c d -> ... (c d)') + + # whether to project codes out to original dimensions + # if the input feature dimensions were not log2(codebook size) + + if project_out: + codes = self.project_out(codes) + + # rearrange codes back to original shape + + if should_transpose: + codes = rearrange(codes, 'b ... d -> b d ...') + + return codes + + def bits_to_z(self, bits): + # assert bits must contain only -1 and 1 + assert torch.all(bits.abs() == 1) + quantized = bits.float() + quantized = self.maybe_l2norm(quantized) + z = self.project_out(quantized) + return z + + def forward( + self, + x, + inv_temperature = 100., + return_loss_breakdown = False, + mask = None, + return_bits = False + ): + """ + einstein notation + b - batch + n - sequence (or flattened spatial dimensions) + d - feature dimension, which is also log2(codebook size) + c - number of codebook dim + """ + + is_img_or_video = x.ndim >= 4 + should_transpose = default(self.channel_first, is_img_or_video) + + # standardize image or video into (batch, seq, dimension) + + if should_transpose: + x = rearrange(x, 'b d ... -> b ... d') + x, ps = pack_one(x, 'b * d') + + assert x.shape[-1] == self.dim, f'expected dimension of {self.dim} but received {x.shape[-1]}' + + x = self.project_in(x) + + # maybe soft clamp + + if exists(self.soft_clamp_input_value): + clamp_value = self.soft_clamp_input_value + x = (x / clamp_value).tanh() * clamp_value + + # split out number of codebooks + + x = rearrange(x, 'b n (c d) -> b n c d', c = self.num_codebooks) + + # maybe l2norm + + x = self.maybe_l2norm(x) + + # whether to force quantization step to be full precision or not + + force_f32 = self.force_quantization_f32 + + quantization_context = partial(autocast, 'cuda', enabled = False) if force_f32 else nullcontext + + with quantization_context(): + + if force_f32: + orig_dtype = x.dtype + x = x.float() + + # quantize by eq 3. + + original_input = x + + codebook_value = torch.ones_like(x) * self.codebook_scale + quantized = torch.where(x > 0, codebook_value, -codebook_value) + if return_bits: + return quantized + + # calculate indices + + indices = reduce((quantized > 0).int() * self.mask.int(), 'b n c d -> b n c', 'sum') + + # maybe l2norm + + quantized = self.maybe_l2norm(quantized) + + # use straight-through gradients (optionally with custom activation fn) if training + + if self.training: + x = self.activation(x) + x = x + (quantized - x).detach() + else: + x = quantized + + # entropy aux loss + if self.soft_entropy_loss: + entropy_aux_loss = soft_entropy_loss(x, tau=1.0, gamma=1.0) + elif self.training and self.enable_entropy_loss: + + if force_f32: + codebook = self.codebook.float() + + codebook = self.maybe_l2norm(codebook) + + # whether to only use a fraction of probs, for reducing memory + + if self.frac_per_sample_entropy < 1.: + # account for mask + if exists(mask): + original_input = original_input[mask] + original_input = rearrange(original_input, 'b n ... -> (b n) ...') + + rand_mask = torch.randn(self.codebook_dim).argsort(dim = -1) < 16 + + sampled_input = original_input[..., rand_mask] + + sampled_distance = -2 * einsum('... i d, j d -> ... i j', sampled_input, codebook) + + sampled_prob = (-sampled_distance * inv_temperature).softmax(dim = -1) + + per_sample_probs = sampled_prob + else: + if exists(mask): + original_input = original_input[mask] + original_input = rearrange(original_input, 'b n ... -> (b n) ...') + # the same as euclidean distance up to a constant + distance = -2 * einsum('... i d, j d -> ... i j', original_input, codebook) + + prob = (-distance * inv_temperature).softmax(dim = -1) + + per_sample_probs = prob + + # calculate per sample entropy + + per_sample_entropy = entropy(per_sample_probs).mean() + + # distribution over all available tokens in the batch + + avg_prob = reduce(per_sample_probs, '... c d -> c d', 'mean') + + avg_prob = maybe_distributed_mean(avg_prob) + + codebook_entropy = entropy(avg_prob).mean() + + # 1. entropy will be nudged to be low for each code, to encourage the network to output confident predictions + # 2. codebook entropy will be nudged to be high, to encourage all codes to be uniformly used within the batch + + entropy_aux_loss = per_sample_entropy - self.diversity_gamma * codebook_entropy + else: + # if not training, just return dummy 0 + entropy_aux_loss = per_sample_entropy = codebook_entropy = self.zero + + # whether to make the entropy loss positive or not through a (shifted) softplus + + if self.training and self.experimental_softplus_entropy_loss: + entropy_aux_loss = F.softplus(entropy_aux_loss + self.entropy_loss_offset) + + # commit loss + + if self.training and self.commitment_loss_weight > 0.: + + commit_loss = F.mse_loss(original_input, quantized.detach(), reduction = 'none') + + if exists(mask): + commit_loss = commit_loss[mask] + + commit_loss = commit_loss.mean() + else: + commit_loss = self.zero + + # input back to original dtype if needed + + if force_f32: + x = x.type(orig_dtype) + + # merge back codebook dim + + x = rearrange(x, 'b n c d -> b n (c d)') + + # project out to feature dimension if needed + + x = self.project_out(x) + + # reconstitute image or video dimensions + + if should_transpose: + x = unpack_one(x, ps, 'b * d') + x = rearrange(x, 'b ... d -> b d ...') + + indices = unpack_one(indices, ps, 'b * c') + + # whether to remove single codebook dim + + if not self.keep_num_codebooks_dim: + indices = rearrange(indices, '... 1 -> ...') + + # complete aux loss + + aux_loss = entropy_aux_loss * self.entropy_loss_weight + commit_loss * self.commitment_loss_weight + + # returns + + ret = Return(x, indices, aux_loss) + + if not return_loss_breakdown: + return ret + + return ret, LossBreakdown(per_sample_entropy, codebook_entropy, commit_loss) + +class GroupedResidualBSQ(Module): + def __init__( + self, + *, + dim, + groups = 1, + accept_image_fmap = False, + **kwargs + ): + super().__init__() + self.dim = dim + self.groups = groups + assert (dim % groups) == 0 + dim_per_group = dim // groups + + self.accept_image_fmap = accept_image_fmap + + self.rvqs = nn.ModuleList([]) + + for _ in range(groups): + self.rvqs.append(LFQ( + dim = dim_per_group, + **kwargs + )) + + self.codebook_size = self.rvqs[0].codebook_size + + @property + def codebooks(self): + return torch.stack(tuple(rvq.codebooks for rvq in self.rvqs)) + + @property + def split_dim(self): + return 1 if self.accept_image_fmap else -1 + + def get_codes_from_indices(self, indices): + codes = tuple(rvq.get_codes_from_indices(chunk_indices) for rvq, chunk_indices in zip(self.rvqs, indices)) + return torch.stack(codes) + + def get_output_from_indices(self, indices): + outputs = tuple(rvq.get_output_from_indices(chunk_indices) for rvq, chunk_indices in zip(self.rvqs, indices)) + return torch.cat(outputs, dim = self.split_dim) + + def forward( + self, + x, + return_all_codes = False + ): + shape, split_dim = x.shape, self.split_dim + assert shape[split_dim] == self.dim + + # split the feature dimension into groups + + x = x.chunk(self.groups, dim = split_dim) + + forward_kwargs = dict( + ) + + # invoke residual vq on each group + + out = tuple(rvq(chunk, **forward_kwargs) for rvq, chunk in zip(self.rvqs, x)) + out = tuple(zip(*out)) + + # otherwise, get all the zipped outputs and combine them + + quantized, all_indices, *maybe_aux_loss = out + + quantized = torch.cat(quantized, dim = split_dim) + all_indices = torch.stack(all_indices) + + ret = (quantized, all_indices, *maybe_aux_loss) + return ret diff --git a/seed-vc/modules/astral_quantization/convnext.py b/seed-vc/modules/astral_quantization/convnext.py new file mode 100644 index 0000000000000000000000000000000000000000..dc7ac97c993ba9b5811f5f57185f5ed5e84547ed --- /dev/null +++ b/seed-vc/modules/astral_quantization/convnext.py @@ -0,0 +1,209 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from typing import List + + +class ConvNextV2LayerNorm(nn.Module): + r"""LayerNorm that supports two data formats: channels_last (default) or channels_first. + The ordering of the dimensions in the inputs. channels_last corresponds to inputs with shape (batch_size, height, + width, channels) while channels_first corresponds to inputs with shape (batch_size, channels, height, width). + """ + + def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"): + super().__init__() + self.weight = nn.Parameter(torch.ones(normalized_shape)) + self.bias = nn.Parameter(torch.zeros(normalized_shape)) + self.eps = eps + self.data_format = data_format + if self.data_format not in ["channels_last", "channels_first"]: + raise NotImplementedError(f"Unsupported data format: {self.data_format}") + self.normalized_shape = (normalized_shape,) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.data_format == "channels_last": + x = torch.nn.functional.layer_norm( + x, self.normalized_shape, self.weight, self.bias, self.eps + ) + elif self.data_format == "channels_first": + input_dtype = x.dtype + x = x.float() + u = x.mean(1, keepdim=True) + s = (x - u).pow(2).mean(1, keepdim=True) + x = (x - u) / torch.sqrt(s + self.eps) + x = x.to(dtype=input_dtype) + x = self.weight[None, :, None] * x + self.bias[None, :, None] + return x + + +class GRN(nn.Module): + def __init__(self, dim): + super().__init__() + self.gamma = nn.Parameter(torch.zeros(1, 1, dim)) + self.beta = nn.Parameter(torch.zeros(1, 1, dim)) + + def forward(self, x): + Gx = torch.norm(x, p=2, dim=1, keepdim=True) + Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6) + return self.gamma * (x * Nx) + self.beta + x + +class InterpolationLayer(nn.Module): + def __init__(self, ): # this is a default of 1 / 50 * (44100 / 512) / 4 + super().__init__() + pass + + def forward(self, x: torch.Tensor, target_len: torch.Tensor, *args, **kwargs) -> torch.Tensor: + x = F.interpolate(x, size=target_len, mode='linear') + return x + +class ConvNeXtV2Stage(nn.Module): + def __init__( + self, + dim: int = 512, + intermediate_dim: int = 2048, + num_blocks: int = 1, + dilation: int = 1, + downsample_layer_indices: List[int] = None, + downsample_factors: List[int] = None, + upsample_layer_indices: List[int] = None, + upsample_factors: List[int] = None, + interpolation_layer_indices: List[int] = None, + input_dim: int = None, + output_dim: int = None, + gin_channels: int = 0, + ): + super().__init__() + # maybe downsample layers + if downsample_layer_indices is not None: + assert downsample_factors is not None + self.downsample_blocks = nn.ModuleList( + [ + nn.Sequential( + ConvNextV2LayerNorm(dim, data_format="channels_first"), + nn.Conv1d( + dim, dim, kernel_size=downsample_factor, stride=downsample_factor + ), + ) for _, downsample_factor in zip(downsample_layer_indices, downsample_factors) + ] + ) + self.downsample_layer_indices = downsample_layer_indices + else: + self.downsample_blocks = nn.ModuleList() + self.downsample_layer_indices = [] + + # maybe upsample layers + if upsample_layer_indices is not None: + assert upsample_factors is not None + self.upsample_blocks = nn.ModuleList( + [ + nn.Sequential( + ConvNextV2LayerNorm(dim, data_format="channels_first"), + nn.ConvTranspose1d( + dim, dim, kernel_size=upsample_factor, stride=upsample_factor + ), + ) for _, upsample_factor in zip(upsample_layer_indices, upsample_factors) + ] + ) + self.upsample_layer_indices = upsample_layer_indices + else: + self.upsample_blocks = nn.ModuleList() + self.upsample_layer_indices = [] + + # maybe interpolation layers + if interpolation_layer_indices is not None: + self.interpolation_blocks = nn.ModuleList( + [ + InterpolationLayer() + for _ in interpolation_layer_indices + ] + ) + self.interpolation_layer_indices = interpolation_layer_indices + else: + self.interpolation_blocks = nn.ModuleList() + self.interpolation_layer_indices = [] + + # main blocks + self.blocks = nn.ModuleList( + [ + ConvNeXtV2Block( + dim=dim, + intermediate_dim=intermediate_dim, + dilation=dilation, + ) + for _ in range(num_blocks) + ] + ) + # maybe input and output projections + if input_dim is not None and input_dim != dim: + self.input_projection = nn.Conv1d(input_dim, dim, kernel_size=1) + else: + self.input_projection = nn.Identity() + if output_dim is not None and output_dim != dim: + self.output_projection = nn.Conv1d(dim, output_dim, kernel_size=1) + else: + self.output_projection = nn.Identity() + + if gin_channels > 0: + self.gin = nn.Conv1d(gin_channels, dim, kernel_size=1) + + def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor: + x = self.input_projection(x) # B, D, T + if hasattr(self, 'gin'): + g = kwargs['g'] + x = x + self.gin(g) + # pad to a multiple of cumprod(downsample_factors) + if len(self.downsample_blocks) > 0: + downsample_factor = 1 + for factor in self.downsample_blocks: + downsample_factor *= factor[1].stride[0] + pad_len = downsample_factor - x.size(-1) % downsample_factor + if pad_len > 0: + x = torch.cat([x, torch.zeros_like(x[:, :, :pad_len])], dim=-1) + + # main blocks + for layer_idx, block in enumerate(self.blocks): + if layer_idx in self.downsample_layer_indices: + x = self.downsample_blocks[self.downsample_layer_indices.index(layer_idx)](x) + if layer_idx in self.upsample_layer_indices: + x = self.upsample_blocks[self.upsample_layer_indices.index(layer_idx)](x) + if layer_idx in self.interpolation_layer_indices: + x = self.interpolation_blocks[self.interpolation_layer_indices.index(layer_idx)](x, target_len=kwargs['target_len']) + x = block(x) + x = self.output_projection(x) + return x + + def setup_caches(self, *args, **kwargs): + pass + + +class ConvNeXtV2Block(nn.Module): + def __init__( + self, + dim: int, + intermediate_dim: int, + dilation: int = 1, + ): + super().__init__() + padding = (dilation * (7 - 1)) // 2 + self.dwconv = nn.Conv1d( + dim, dim, kernel_size=7, padding=padding, groups=dim, dilation=dilation + ) # depthwise conv + self.norm = ConvNextV2LayerNorm(dim, data_format="channels_first") + self.pwconv1 = nn.Linear( + dim, intermediate_dim + ) # pointwise/1x1 convs, implemented with linear layers + self.act = nn.GELU() + self.grn = GRN(intermediate_dim) + self.pwconv2 = nn.Linear(intermediate_dim, dim) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + residual = x + x = self.dwconv(x) + x = self.norm(x) + x = x.transpose(1, 2) # b d n -> b n d + x = self.pwconv1(x) + x = self.act(x) + x = self.grn(x) + x = self.pwconv2(x) + x = x.transpose(1, 2) # b n d -> b d n + return residual + x \ No newline at end of file diff --git a/seed-vc/modules/astral_quantization/default_model.py b/seed-vc/modules/astral_quantization/default_model.py new file mode 100644 index 0000000000000000000000000000000000000000..259515557fc42c36007099c29e71138c2ddae93a --- /dev/null +++ b/seed-vc/modules/astral_quantization/default_model.py @@ -0,0 +1,73 @@ +import torch +from transformers import AutoTokenizer, AutoModel, Wav2Vec2FeatureExtractor + +class AstralQuantizer(torch.nn.Module): + def __init__( + self, + tokenizer_name: str, + ssl_model_name: str, + ssl_output_layer: int, + encoder: torch.nn.Module, + quantizer: torch.nn.Module, + skip_ssl: bool = False, + ): + super().__init__() + self.encoder = encoder + self.quantizer = quantizer + self.tokenizer_name = tokenizer_name + self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_name) + + # Load SSL model from Huggingface + self.ssl_model_name = ssl_model_name + self.ssl_output_layer = ssl_output_layer + self.ssl_feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(ssl_model_name) + + if skip_ssl: # in case the same SSL model has been loaded somewhere else + self.ssl_model = None + else: + self.ssl_model = AutoModel.from_pretrained(ssl_model_name).eval() + self.ssl_model.encoder.layers = self.ssl_model.encoder.layers[:ssl_output_layer] + self.ssl_model.encoder.layer_norm = torch.nn.Identity() + + def load_separate_checkpoint(self, checkpoint_path): + params = torch.load(checkpoint_path, map_location='cpu')['net'] + for key in params.keys(): + for k in list(params[key].keys()): + if k.startswith("module."): + params[key][k[len("module."):]] = params[key][k] + del params[key][k] + self.encoder.load_state_dict(params['encoder']) + self.quantizer.load_state_dict(params['vq']) + if self.decoder is not None: + self.decoder.load_state_dict(params['decoder']) + if self.asr_decoder is not None: + self.asr_decoder.load_state_dict(params['predictor'], strict=False) + + def forward(self, waves_16k, wave_16k_lens, ssl_model=None): + ssl_fn = self.ssl_model if self.ssl_model else ssl_model + assert ssl_fn is not None, "In case in-class SSL model loading is skipped, external ssl_model must be provided" + waves_16k_input_list = [ + waves_16k[bib, :wave_16k_lens[bib]].cpu().numpy() + for bib in range(len(waves_16k)) + ] + alt_inputs = self.ssl_feature_extractor( + waves_16k_input_list, + return_tensors='pt', + return_attention_mask=True, + padding=True, + sampling_rate=16000 + ).to(waves_16k.device) + feature_lens = alt_inputs.data['attention_mask'].sum(-1) // 320 # frame rate of hubert is 50 Hz + + outputs = ssl_fn( + alt_inputs.input_values, + attention_mask=alt_inputs.attention_mask, + ) + last_hidden_states = outputs.last_hidden_state + last_hidden_states = last_hidden_states[:, :feature_lens.max(), :] + feature_lens = feature_lens.clamp(max=last_hidden_states.size(1)) + last_hidden_states = last_hidden_states.transpose(1, 2) + x_hidden = self.encoder(last_hidden_states, feature_lens) + x_hidden = x_hidden.transpose(1, 2) + x_quantized, indices = self.quantizer(x_hidden)[:2] + return x_quantized, indices, feature_lens \ No newline at end of file diff --git a/seed-vc/modules/astral_quantization/transformer.py b/seed-vc/modules/astral_quantization/transformer.py new file mode 100644 index 0000000000000000000000000000000000000000..015ec417d02e3442d8ea120c3802807c73e89bb4 --- /dev/null +++ b/seed-vc/modules/astral_quantization/transformer.py @@ -0,0 +1,254 @@ +# 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. +from dataclasses import dataclass +from typing import Optional + +import torch +import torch.nn as nn +from torch import Tensor +from torch.nn import functional as F +import time + +def find_multiple(n: int, k: int) -> int: + if n % k == 0: + return n + return n + k - (n % k) + +class AdaptiveLayerNorm(nn.Module): + r"""Adaptive Layer Normalization""" + + def __init__(self, d_model, norm) -> None: + super(AdaptiveLayerNorm, self).__init__() + self.project_layer = nn.Linear(d_model, 2 * d_model) + self.norm = norm + self.d_model = d_model + self.eps = self.norm.eps + + def forward(self, input: Tensor, embedding: Tensor = None) -> Tensor: + if embedding is None: + return self.norm(input) + weight, bias = torch.split( + self.project_layer(embedding), + split_size_or_sections=self.d_model, + dim=-1, + ) + return weight * self.norm(input) + bias + + +@dataclass +class ModelArgs: + block_size: int = 2048 + vocab_size: int = 32000 + n_layer: int = 32 + n_head: int = 32 + dim: int = 4096 + intermediate_size: int = None + n_local_heads: int = -1 + head_dim: int = 64 + rope_base: float = 10000 + norm_eps: float = 1e-5 + has_cross_attention: bool = False + context_dim: int = 0 + is_causal: bool = False + dropout_rate: float = 0.1 + attn_dropout_rate: float = 0.1 + + def __post_init__(self): + if self.n_local_heads == -1: + self.n_local_heads = self.n_head + if self.intermediate_size is None: + hidden_dim = 4 * self.dim + n_hidden = int(2 * hidden_dim / 3) + self.intermediate_size = find_multiple(n_hidden, 256) + # self.head_dim = self.dim // self.n_head + +class Transformer(nn.Module): + def __init__(self, config: ModelArgs) -> None: + super().__init__() + self.config = config + + self.layers = nn.ModuleList(TransformerBlock(config) for _ in range(config.n_layer)) + self.norm = AdaptiveLayerNorm(config.dim, RMSNorm(config.dim, eps=config.norm_eps)) + + self.max_batch_size = -1 + self.max_seq_length = config.block_size + freqs_cis = precompute_freqs_cis(self.config.block_size, self.config.head_dim, + self.config.rope_base) + self.register_buffer("freqs_cis", freqs_cis) + + causal_mask = torch.tril( + torch.ones(self.max_seq_length, self.max_seq_length, dtype=torch.bool) + ) + self.register_buffer("causal_mask", causal_mask) + + def forward(self, + x: Tensor, + c: Tensor, + input_pos: Optional[Tensor] = None, + mask: Optional[Tensor] = None, + context: Optional[Tensor] = None, + context_input_pos: Optional[Tensor] = None, + cross_attention_mask: Optional[Tensor] = None, + ) -> Tensor: + if mask is None: + mask = self.causal_mask[:x.size(1), :x.size(1)] + else: + mask = mask[..., input_pos] + freqs_cis = self.freqs_cis[input_pos] + if context is not None: + context_freqs_cis = self.freqs_cis[context_input_pos] + else: + context_freqs_cis = None + skip_in_x_list = [] + for i, layer in enumerate(self.layers): + x = layer(x, c, freqs_cis, mask, context, context_freqs_cis, cross_attention_mask) + x = self.norm(x, c) + return x + + +class TransformerBlock(nn.Module): + def __init__(self, config: ModelArgs) -> None: + super().__init__() + self.attention = Attention(config) + self.feed_forward = FeedForward(config) + self.ffn_norm = AdaptiveLayerNorm(config.dim, RMSNorm(config.dim, eps=config.norm_eps)) + self.attention_norm = AdaptiveLayerNorm(config.dim, RMSNorm(config.dim, eps=config.norm_eps)) + + if config.has_cross_attention: + self.has_cross_attention = True + self.cross_attention = Attention(config, is_cross_attention=True) + self.cross_attention_norm = AdaptiveLayerNorm(config.dim, RMSNorm(config.dim, eps=config.norm_eps)) + else: + self.has_cross_attention = False + + def forward(self, + x: Tensor, + c: Tensor, + freqs_cis: Tensor, + mask: Tensor, + context: Optional[Tensor] = None, + context_freqs_cis: Optional[Tensor] = None, + cross_attention_mask: Optional[Tensor] = None, + ) -> Tensor: + #time_attn_start = time.time() + h = x + self.attention(self.attention_norm(x, c), freqs_cis, mask) + #print(f"time take for attention of sequence length {x.shape[1]} is {time.time() - time_attn_start}") + if self.has_cross_attention: + h = h + self.cross_attention(self.cross_attention_norm(h, c), freqs_cis, cross_attention_mask, context, context_freqs_cis) + out = h + self.feed_forward(self.ffn_norm(h, c)) + return out + + +class Attention(nn.Module): + def __init__(self, config: ModelArgs, is_cross_attention: bool = False): + super().__init__() + assert config.dim % config.n_head == 0 + + total_head_dim = (config.n_head + 2 * config.n_local_heads) * config.head_dim + # key, query, value projections for all heads, but in a batch + if is_cross_attention: + self.wq = nn.Linear(config.dim, config.n_head * config.head_dim, bias=False) + self.wkv = nn.Linear(config.context_dim, 2 * config.n_local_heads * config.head_dim, bias=False) + else: + self.wqkv = nn.Linear(config.dim, total_head_dim, bias=False) + self.wo = nn.Linear(config.head_dim * config.n_head, config.dim, bias=False) + self.kv_cache = None + + self.n_head = config.n_head + self.head_dim = config.head_dim + self.n_local_heads = config.n_local_heads + self.dim = config.dim + self.attn_dropout_rate = config.attn_dropout_rate + + def forward(self, + x: Tensor, + freqs_cis: Tensor, + mask: Tensor, + context: Optional[Tensor] = None, + context_freqs_cis: Optional[Tensor] = None, + ) -> Tensor: + bsz, seqlen, _ = x.shape + + kv_size = self.n_local_heads * self.head_dim + if context is None: + q, k, v = self.wqkv(x).split([kv_size, kv_size, kv_size], dim=-1) + context_seqlen = seqlen + else: + q = self.wq(x) + k, v = self.wkv(context).split([kv_size, kv_size], dim=-1) + context_seqlen = context.shape[1] + + q = q.view(bsz, seqlen, self.n_head, self.head_dim) + k = k.view(bsz, context_seqlen, self.n_local_heads, self.head_dim) + v = v.view(bsz, context_seqlen, self.n_local_heads, self.head_dim) + + q = apply_rotary_emb(q, freqs_cis) + k = apply_rotary_emb(k, context_freqs_cis if context_freqs_cis is not None else freqs_cis) + + q, k, v = map(lambda x: x.transpose(1, 2), (q, k, v)) + + k = k.repeat_interleave(self.n_head // self.n_local_heads, dim=1) + v = v.repeat_interleave(self.n_head // self.n_local_heads, dim=1) + y = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=self.attn_dropout_rate if self.training else 0.0) + + y = y.transpose(1, 2).contiguous().view(bsz, seqlen, self.head_dim * self.n_head) + + y = self.wo(y) + return y + + +class FeedForward(nn.Module): + def __init__(self, config: ModelArgs) -> None: + super().__init__() + self.w1 = nn.Linear(config.dim, config.intermediate_size, bias=False) + self.w3 = nn.Linear(config.dim, config.intermediate_size, bias=False) + self.w2 = nn.Linear(config.intermediate_size, config.dim, bias=False) + self.dropout = nn.Dropout(config.dropout_rate) + + def forward(self, x: Tensor) -> Tensor: + return self.w2(self.dropout(F.silu(self.w1(x)) * self.w3(x))) + + +class RMSNorm(nn.Module): + def __init__(self, dim: int, eps: float = 1e-5): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + + def _norm(self, x): + return x * torch.rsqrt(torch.mean(x * x, dim=-1, keepdim=True) + self.eps) + + def forward(self, x: Tensor) -> Tensor: + output = self._norm(x.float()).type_as(x) + return output * self.weight + + +def precompute_freqs_cis( + seq_len: int, n_elem: int, base: int = 10000, + dtype: torch.dtype = torch.bfloat16 +) -> Tensor: + freqs = 1.0 / (base ** (torch.arange(0, n_elem, 2)[: (n_elem // 2)].float() / n_elem)) + t = torch.arange(seq_len, device=freqs.device) + freqs = torch.outer(t, freqs) + freqs_cis = torch.polar(torch.ones_like(freqs), freqs) + cache = torch.stack([freqs_cis.real, freqs_cis.imag], dim=-1) + return cache.to(dtype=dtype) + + +def apply_rotary_emb(x: Tensor, freqs_cis: Tensor) -> Tensor: + xshaped = x.float().reshape(*x.shape[:-1], -1, 2) + freqs_cis = freqs_cis.view(1, xshaped.size(1), 1, xshaped.size(3), 2) + x_out2 = torch.stack( + [ + xshaped[..., 0] * freqs_cis[..., 0] - xshaped[..., 1] * freqs_cis[..., 1], + xshaped[..., 1] * freqs_cis[..., 0] + xshaped[..., 0] * freqs_cis[..., 1], + ], + -1, + ) + + x_out2 = x_out2.flatten(3) + return x_out2.type_as(x) + diff --git a/seed-vc/modules/audio.py b/seed-vc/modules/audio.py new file mode 100644 index 0000000000000000000000000000000000000000..abe783b0e0af630319700c931eb51d2ce375282b --- /dev/null +++ b/seed-vc/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/seed-vc/modules/bigvgan/activations.py b/seed-vc/modules/bigvgan/activations.py new file mode 100644 index 0000000000000000000000000000000000000000..61f2808a5466b3cf4d041059700993af5527dd29 --- /dev/null +++ b/seed-vc/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/seed-vc/modules/bigvgan/alias_free_activation/cuda/__init__.py b/seed-vc/modules/bigvgan/alias_free_activation/cuda/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/seed-vc/modules/bigvgan/alias_free_activation/cuda/activation1d.py b/seed-vc/modules/bigvgan/alias_free_activation/cuda/activation1d.py new file mode 100644 index 0000000000000000000000000000000000000000..01a25ffc04edbdfe17859653690e3005d441712f --- /dev/null +++ b/seed-vc/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/seed-vc/modules/bigvgan/alias_free_activation/cuda/anti_alias_activation.cpp b/seed-vc/modules/bigvgan/alias_free_activation/cuda/anti_alias_activation.cpp new file mode 100644 index 0000000000000000000000000000000000000000..c5651f77143bd678169eb11564a7cf7a7969a59e --- /dev/null +++ b/seed-vc/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/seed-vc/modules/openvoice/se_extractor.py b/seed-vc/modules/openvoice/se_extractor.py
new file mode 100644
index 0000000000000000000000000000000000000000..d08717984508c893dd88eeccc02b325ca4ea0a6e
--- /dev/null
+++ b/seed-vc/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/seed-vc/modules/openvoice/transforms.py b/seed-vc/modules/openvoice/transforms.py
new file mode 100644
index 0000000000000000000000000000000000000000..a11f799e023864ff7082c1f49c0cc18351a13b47
--- /dev/null
+++ b/seed-vc/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/seed-vc/modules/openvoice/utils.py b/seed-vc/modules/openvoice/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..4e80909d6a03976400322cc0219d1871e9f84bfa
--- /dev/null
+++ b/seed-vc/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/seed-vc/modules/rmvpe.py b/seed-vc/modules/rmvpe.py
new file mode 100644
index 0000000000000000000000000000000000000000..8d7811b2f98f2bf6e25f71ad3a3444e6999d9d4a
--- /dev/null
+++ b/seed-vc/modules/rmvpe.py
@@ -0,0 +1,637 @@
+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"
+ if torch.cuda.is_available():
+ device = "cuda:0"
+ elif torch.backends.mps.is_available():
+ device = "mps"
+ else:
+ device = "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/seed-vc/modules/v2/ar.py b/seed-vc/modules/v2/ar.py
new file mode 100644
index 0000000000000000000000000000000000000000..0611a459c74d1260442855499b03a6885bb4fb04
--- /dev/null
+++ b/seed-vc/modules/v2/ar.py
@@ -0,0 +1,763 @@
+import dataclasses
+import json
+import math
+from collections import OrderedDict
+from functools import partial, wraps
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Optional, Tuple, List
+from tqdm import tqdm
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+from torch import Tensor
+from torch.nn import functional as F
+from torch.utils.checkpoint import checkpoint
+
+
+def find_multiple(n: int, k: int) -> int:
+ if n % k == 0:
+ return n
+ return n + k - (n % k)
+
+def l2norm(t, groups = 1):
+ t = rearrange(t, '... (g d) -> ... g d', g = groups)
+ t = F.normalize(t, p = 2, dim = -1)
+ return rearrange(t, '... g d -> ... (g d)')
+
+@dataclass
+class BaseModelArgs:
+ model_type: str = "base"
+
+ vocab_size: int = 32000
+ n_layer: int = 32
+ n_head: int = 32
+ dim: int = 4096
+ intermediate_size: int = None
+ n_local_heads: int = -1
+ head_dim: int = 64
+ rope_base: float = 10000
+ norm_eps: float = 1e-5
+ max_seq_len: int = 4096
+ dropout: float = 0.0
+ tie_word_embeddings: bool = True
+ attention_qkv_bias: bool = False
+
+ # Gradient checkpointing
+ use_gradient_checkpointing: bool = False
+
+ # Initialize the model
+ initializer_range: float = 0.02
+
+ qk_norm: bool = False
+ layerscale: bool = False
+
+ def __post_init__(self):
+ if self.n_local_heads == -1:
+ self.n_local_heads = self.n_head
+ if self.intermediate_size is None:
+ hidden_dim = 4 * self.dim
+ n_hidden = int(2 * hidden_dim / 3)
+ self.intermediate_size = find_multiple(n_hidden, 256)
+ self.head_dim = self.dim // self.n_head
+
+ def save(self, path: str):
+ with open(path, "w") as f:
+ json.dump(self.__dict__, f, indent=4, sort_keys=True, ensure_ascii=False)
+
+
+@dataclass
+class NaiveModelArgs(BaseModelArgs):
+ model_type: str = "naive"
+
+
+class KVCache(nn.Module):
+ def __init__(
+ self, max_batch_size, max_seq_len, n_heads, head_dim, dtype=torch.bfloat16
+ ):
+ super().__init__()
+ cache_shape = (max_batch_size, n_heads, max_seq_len, head_dim)
+ self.register_buffer("k_cache", torch.zeros(cache_shape, dtype=dtype))
+ self.register_buffer("v_cache", torch.zeros(cache_shape, dtype=dtype))
+
+ def update(self, input_pos, k_val, v_val):
+ # input_pos: [S], k_val: [B, H, S, D]
+ assert input_pos.shape[0] == k_val.shape[2]
+
+ k_out = self.k_cache
+ v_out = self.v_cache
+ k_out[:, :, input_pos] = k_val
+ v_out[:, :, input_pos] = v_val
+
+ return k_out, v_out
+
+
+@dataclass
+class TransformerForwardResult:
+ token_logits: Tensor
+ token_targets: Tensor
+
+
+@dataclass
+class BaseTransformerForwardResult:
+ logits: Tensor
+ hidden_states: Tensor
+
+
+class BaseTransformer(nn.Module):
+ def __init__(
+ self,
+ config: BaseModelArgs,
+ init_weights: bool = True,
+ ) -> None:
+ super().__init__()
+ self.config = config
+
+ # Slow transformer
+ self.embeddings = nn.Embedding(
+ config.vocab_size,
+ config.dim,
+ )
+ self.layers = nn.ModuleList(
+ TransformerBlock(config, use_sdpa=True) for _ in range(config.n_layer)
+ )
+ self.norm = RMSNorm(config.dim, eps=config.norm_eps)
+
+ if self.config.tie_word_embeddings is False:
+ self.output = nn.Linear(
+ config.dim,
+ config.vocab_size,
+ bias=False,
+ )
+
+ self.register_buffer(
+ "freqs_cis",
+ precompute_freqs_cis(
+ config.max_seq_len,
+ config.dim // config.n_head,
+ config.rope_base,
+ ),
+ persistent=False,
+ )
+ self.register_buffer(
+ "causal_mask",
+ torch.tril(
+ torch.ones(
+ config.max_seq_len,
+ config.max_seq_len,
+ dtype=torch.bool,
+ )
+ ),
+ persistent=False,
+ )
+
+ self.output = nn.Linear(
+ config.dim,
+ config.vocab_size,
+ bias=False,
+ )
+
+ # For kv cache
+ self.max_batch_size = -1
+ self.max_seq_len = -1
+
+ if init_weights:
+ self.apply(self._init_weights)
+
+ def setup_caches(
+ self, max_batch_size: int, max_seq_len: int, dtype: torch.dtype = torch.bfloat16, device: torch.device = "cuda"
+ ):
+ if self.max_seq_len >= max_seq_len and self.max_batch_size >= max_batch_size:
+ return
+
+ head_dim = self.config.dim // self.config.n_head
+ max_seq_len = find_multiple(max_seq_len, 8)
+ self.max_seq_len = max_seq_len
+ self.max_batch_size = max_batch_size
+
+ for b in self.layers:
+ b.attention.kv_cache = KVCache(
+ max_batch_size,
+ max_seq_len,
+ self.config.n_local_heads,
+ head_dim,
+ dtype=dtype,
+ ).to(device)
+
+ def embed_base(self, x: Tensor, x_lens: Tensor) -> Tensor:
+ for bib in range(x.size(0)):
+ x[bib, x_lens[bib]:] = self.config.vocab_size - 1
+
+ x_emb = self.embeddings(x)
+ return x, x_emb
+
+ def forward(
+ self,
+ inp: Tensor,
+ key_padding_mask: Optional[Tensor] = None,
+ input_pos: Optional[Tensor] = None,
+ ) -> BaseTransformerForwardResult:
+ seq_len = inp.size(1)
+
+ # Here we want to merge the embeddings of the codebooks
+ # x = self.embed(inp)
+ x = inp.clone()
+
+ if input_pos is None:
+ freqs_cis = self.freqs_cis[:seq_len].repeat(inp.size(0), 1, 1, 1)
+ else:
+ freqs_cis = self.freqs_cis[input_pos]
+
+ # Not that the causal mask here follows the definition of scaled_dot_product_attention
+ # That is, FALSE means masked out
+ # To maintain consistency, key_padding_mask use TRUE to mask out
+ mask = None
+ if key_padding_mask is not None:
+ mask = self.causal_mask[None, None, :seq_len, :seq_len] # (B, N, Q, K)
+ mask = mask & key_padding_mask[:, None, None, :].logical_not()
+
+ for layer in self.layers:
+ if self.config.use_gradient_checkpointing and self.training:
+ x = checkpoint(layer, x, freqs_cis, mask, use_reentrant=True)
+ else:
+ x = layer(x, freqs_cis, mask)
+
+ # We got slow_out here
+ slow_out = self.norm(x)
+
+ if self.config.tie_word_embeddings:
+ token_logits = F.linear(slow_out, self.embeddings.weight)
+ else:
+ token_logits = self.output(slow_out)
+
+ return BaseTransformerForwardResult(
+ logits=token_logits,
+ hidden_states=x,
+ )
+
+ def forward_generate(
+ self,
+ inp: Tensor,
+ input_pos: Optional[Tensor] = None,
+ kv_pos: Optional[Tensor] = None,
+ return_all: bool = False,
+ ) -> BaseTransformerForwardResult:
+ # This is used for generation, optimized for torch compile
+
+ x = inp
+ max_seq_len = self.max_seq_len
+
+ mask = self.causal_mask[None, None, kv_pos, :max_seq_len] # (B, N, Q, K)
+ freqs_cis = self.freqs_cis[input_pos]
+
+ for layer in self.layers:
+ x = layer(x, freqs_cis, mask, input_pos=kv_pos)
+
+ x = x[:, -1:]
+
+ # We got slow_out here
+ slow_out = self.norm(x)
+
+ token_logits = self.output(slow_out)
+
+ return BaseTransformerForwardResult(
+ logits=token_logits,
+ hidden_states=x,
+ )
+
+ def _init_weights(self, module):
+ std = self.config.initializer_range
+ if isinstance(module, nn.Linear):
+ module.weight.data.normal_(mean=0.0, std=std)
+ if module.bias is not None:
+ module.bias.data.zero_()
+ elif isinstance(module, nn.Embedding):
+ module.weight.data.normal_(mean=0.0, std=std)
+ if module.padding_idx is not None:
+ module.weight.data[module.padding_idx].zero_()
+
+class NaiveTransformer(BaseTransformer):
+ def __init__(self, config: NaiveModelArgs) -> None:
+ super().__init__(config, init_weights=False)
+ self.apply(self._init_weights)
+
+ def forward(
+ self,
+ inp: Tensor,
+ cond_lens: Tensor,
+ target: Tensor,
+ target_lens: Tensor,
+ key_padding_mask: Optional[Tensor] = None,
+ input_pos: Optional[Tensor] = None,
+ ) -> TransformerForwardResult:
+ parent_result = super().forward(
+ inp=inp,
+ key_padding_mask=key_padding_mask,
+ input_pos=input_pos,
+ )
+ token_logits = parent_result.logits
+
+ # construct targets for token_logits
+ token_targets = torch.zeros(token_logits.size(0), token_logits.size(1), dtype=torch.long,
+ device=target.device) - 100
+ for bib in range(token_targets.size(0)):
+ token_targets[bib, cond_lens[bib] + 1:cond_lens[bib] + target_lens[bib] + 1] = target[bib, :target_lens[bib]]
+ token_targets[bib, cond_lens[bib] + target_lens[bib] + 1] = self.config.vocab_size - 1
+ return TransformerForwardResult(
+ token_logits=token_logits,
+ token_targets=token_targets,
+ )
+
+ def infer_slow(self, inp: Tensor, input_pos: Optional[Tensor] = None):
+ # no kv cache used
+ parent_result = super().forward(inp, input_pos=input_pos)
+ latent = parent_result.hidden_states[:, -1]
+ base_logits = parent_result.logits[:, -1]
+ base_sampled, _ = topk_sampling(base_logits, top_k=-1, top_p=1.0)
+ return base_sampled
+
+ def forward_generate(
+ self,
+ x: Tensor,
+ input_pos: Optional[Tensor] = None,
+ kv_pos: Optional[Tensor] = None,
+ vq_masks: Optional[Tensor] = None,
+ ) -> TransformerForwardResult:
+ x = super().forward_generate(x, input_pos, kv_pos, vq_masks)
+ return x
+
+class NaiveWrapper(nn.Module):
+ def __init__(self, model: NaiveTransformer) -> None:
+ super().__init__()
+ self.model = model
+ self.sep_token_emb = nn.Parameter(torch.randn(model.config.dim))
+
+ def setup_caches(self, max_batch_size: int, max_seq_len: int, dtype: torch.dtype = torch.bfloat16, device: torch.device = "cuda"):
+ self.model.setup_caches(max_batch_size, max_seq_len, dtype, device)
+
+ def forward(self, cond: Tensor, cond_lens: Tensor, x: Tensor, x_lens: Tensor) -> torch.Tensor:
+ # style_emb = self.style_in(style).unsqueeze(1) # [B, 1, D]
+ sep_token_emb = self.sep_token_emb.expand(x.size(0), 1, -1)
+ _, x_emb = self.model.embed_base(x, x_lens)
+ emb_seq_list = []
+ for i in range(x.size(0)):
+ emb_seq = torch.cat([
+ sep_token_emb[i:i + 1],
+ cond[i:i+1, :cond_lens[i]],
+ sep_token_emb[i:i+1],
+ x_emb[i:i+1, :x_lens[i]]], dim=1)
+ emb_seq_list.append(emb_seq)
+ max_len = max([emb_seq.size(1) for emb_seq in emb_seq_list])
+ emb_seq = torch.cat([
+ F.pad(emb_seq, (0, 0, 0, max_len - emb_seq.size(1)), value=0)
+ for emb_seq in emb_seq_list
+ ], dim=0)
+ # input_pos = torch.arange(emb_seq.size(1), device=emb_seq.device).repeat(emb_seq.size(0), 1)
+ input_pos = torch.zeros(emb_seq.size(0), emb_seq.size(1), device=emb_seq.device, dtype=torch.long)
+ for i in range(x.size(0)):
+ input_pos[i, :cond_lens[i] + 1] = torch.arange(cond_lens[i] + 1, device=emb_seq.device)
+ input_pos[i, cond_lens[i] + 1: cond_lens[i] + x_lens[i] + 2] = torch.arange(x_lens[i] + 1, device=emb_seq.device)
+ out = self.model(emb_seq, cond_lens, x, x_lens, input_pos=input_pos)
+ loss = F.cross_entropy(out.token_logits.transpose(1, 2), out.token_targets.long(), ignore_index=-100)
+ return loss
+
+ @torch.no_grad()
+ def infer(self, cond: Tensor) -> torch.Tensor:
+ sep_token_emb = self.sep_token_emb.expand(1, 1, -1)
+ emb_seq = torch.cat([sep_token_emb, cond, sep_token_emb], dim=1)
+ pred_codes = []
+ input_pos = torch.arange(cond.size(1) + 1, device=cond.device)
+ for i in tqdm(range(4000)):
+ input_pos = torch.cat([input_pos, torch.LongTensor([i]).to(cond.device)], dim=0)
+ base = self.model.infer_slow(emb_seq, input_pos)
+ if base == self.model.config.vocab_size - 1:
+ break
+ new_emb = self.model.embed_base(base, torch.LongTensor([1]).to(base.device))[1]
+ emb_seq = torch.cat([emb_seq, new_emb], dim=1)
+ pred_codes.append(base)
+ return torch.cat(pred_codes, dim=-1)
+
+ @torch.no_grad()
+ def generate(
+ self,
+ prompt_text,
+ prompt_target,
+ compiled_decode_fn = None,
+ **sampling_kwargs,
+ ):
+ sep_token_emb = self.sep_token_emb.expand(1, 1, -1)
+ emb_seq = torch.cat([sep_token_emb, prompt_text, sep_token_emb], dim=1)
+ input_pos = torch.arange(prompt_text.size(1) + 1, device=emb_seq.device)
+ input_pos = torch.cat([input_pos, torch.LongTensor([0]).to(emb_seq.device)])
+ prompt_target_emb = self.model.embed_base(prompt_target,torch.LongTensor([prompt_target.size(1)]).to(prompt_target.device))[1]
+ emb_seq = torch.cat([emb_seq, prompt_target_emb], dim=1)
+ input_pos = torch.cat([input_pos, torch.arange(prompt_target_emb.size(1)).to(input_pos.device) + 1])
+
+ pred_codes = []
+ kv_pos = torch.arange(emb_seq.size(1), device=emb_seq.device)
+ next_tokens = self.decode_one_token_ar(emb_seq, input_pos, kv_pos, suppress_tokens=[self.model.config.vocab_size - 1], **sampling_kwargs)
+ pred_base = next_tokens[0]
+ pred_codes.append(pred_base)
+ new_emb = self.model.embed_base(pred_base.unsqueeze(0), torch.LongTensor([1]).to(pred_base.device))[1]
+ emb_seq = torch.cat([emb_seq, new_emb], dim=1)
+ for _ in tqdm(range(4000)):
+ suppress_eos = len(pred_codes) < 10
+ input_pos = input_pos[-1:] + 1
+ kv_pos = kv_pos[-1:] + 1
+ next_tokens = self.decode_one_token_ar(
+ emb_seq[:, -1:].reshape(1, 1, -1),
+ input_pos.reshape(1),
+ kv_pos.reshape(1),
+ previous_tokens=torch.cat(pred_codes),
+ suppress_tokens=[self.model.config.vocab_size - 1] if suppress_eos else None,
+ compiled_decode_fn=compiled_decode_fn,
+ **sampling_kwargs)
+ pred_base = next_tokens[0]
+ if pred_base == self.model.config.vocab_size - 1:
+ break
+ pred_codes.append(pred_base.clone())
+ new_emb = self.model.embed_base(pred_base.unsqueeze(0), torch.LongTensor([1]).to(pred_base.device))[1]
+ emb_seq = torch.cat([emb_seq, new_emb], dim=1)
+ return torch.stack(pred_codes, dim=-1)
+
+ def decode_one_token_ar(
+ self,
+ x: torch.Tensor,
+ input_pos: torch.Tensor,
+ kv_pos: torch.Tensor,
+ previous_tokens: torch.Tensor = None,
+ compiled_decode_fn = None,
+ **sampling_kwargs,
+ ) -> torch.Tensor:
+ if compiled_decode_fn is not None:
+ x = compiled_decode_fn(x, input_pos, kv_pos)
+ else:
+ x = self.model.forward_generate(x, input_pos, kv_pos)
+
+ sampling_kwargs_main = sampling_kwargs.copy()
+ codebooks = [
+ sample(
+ x.logits,
+ previous_tokens=(
+ previous_tokens[0] if previous_tokens is not None else None
+ ),
+ **sampling_kwargs_main,
+ )[0]
+ ]
+ codebooks = torch.stack(codebooks, dim=0)
+ return codebooks
+
+class TransformerBlock(nn.Module):
+ def __init__(self, config: BaseModelArgs, use_sdpa: bool = True) -> None:
+ super().__init__()
+ self.attention = Attention(config, use_sdpa=use_sdpa)
+ self.feed_forward = FeedForward(config)
+ self.ffn_norm = RMSNorm(config.dim, config.norm_eps)
+ self.attention_norm = RMSNorm(config.dim, config.norm_eps)
+
+ def forward(
+ self, x: Tensor, freqs_cis: Tensor, mask: Tensor, input_pos: Tensor = None
+ ) -> Tensor:
+ h = x + self.attention(self.attention_norm(x), freqs_cis, mask, input_pos)
+ out = h + self.feed_forward(self.ffn_norm(h))
+ return out
+
+
+class Attention(nn.Module):
+ def __init__(self, config: BaseModelArgs, use_sdpa: bool = True):
+ super().__init__()
+ assert config.dim % config.n_head == 0
+
+ total_head_dim = (config.n_head + 2 * config.n_local_heads) * config.head_dim
+ # key, query, value projections for all heads, but in a batch
+ self.wqkv = nn.Linear(
+ config.dim, total_head_dim, bias=config.attention_qkv_bias
+ )
+ self.wo = nn.Linear(config.dim, config.dim, bias=False)
+ self.kv_cache = None
+
+ self.dropout = config.dropout
+ self.n_head = config.n_head
+ self.head_dim = config.head_dim
+ self.n_local_heads = config.n_local_heads
+ self.dim = config.dim
+ self.use_sdpa = use_sdpa
+ self._register_load_state_dict_pre_hook(self.load_hook)
+ self.qk_norm = config.qk_norm
+ self.qk_norm_groups = 1
+ self.qk_norm_scale = 10
+ self.qk_norm_dim_scale = False
+ self.qk_norm_q_scale = self.qk_norm_k_scale = 1
+
+ if self.qk_norm and self.qk_norm_dim_scale:
+ self.qk_norm_q_scale = nn.Parameter(torch.ones(self.n_head, 1, self.head_dim))
+ self.qk_norm_k_scale = nn.Parameter(torch.ones(self.n_head, 1, self.head_dim))
+ def load_hook(self, state_dict, prefix, *args):
+ if prefix + "wq.weight" in state_dict:
+ wq = state_dict.pop(prefix + "wq.weight")
+ wk = state_dict.pop(prefix + "wk.weight")
+ wv = state_dict.pop(prefix + "wv.weight")
+ state_dict[prefix + "wqkv.weight"] = torch.cat([wq, wk, wv])
+
+ def forward(
+ self,
+ x: Tensor,
+ freqs_cis: Tensor,
+ mask: Tensor,
+ input_pos: Optional[Tensor] = None,
+ ) -> Tensor:
+ bsz, seqlen, _ = x.shape
+
+ kv_size = self.n_local_heads * self.head_dim
+ q, k, v = self.wqkv(x).split([self.dim, kv_size, kv_size], dim=-1)
+
+ q = q.view(bsz, seqlen, self.n_head, self.head_dim)
+ k = k.view(bsz, seqlen, self.n_local_heads, self.head_dim)
+ v = v.view(bsz, seqlen, self.n_local_heads, self.head_dim)
+
+ if self.qk_norm:
+ qk_l2norm = partial(l2norm, groups = self.qk_norm_groups)
+ q, k = map(qk_l2norm, (q, k))
+ scale = self.qk_norm_scale
+
+ q = q * self.qk_norm_q_scale
+ k = k * self.qk_norm_k_scale
+
+ q = apply_rotary_emb(q, freqs_cis)
+ k = apply_rotary_emb(k, freqs_cis)
+
+ q, k, v = map(lambda x: x.transpose(1, 2), (q, k, v))
+
+ if self.kv_cache is not None:
+ k, v = self.kv_cache.update(input_pos, k, v)
+
+ k = k.repeat_interleave(self.n_head // self.n_local_heads, dim=1)
+ v = v.repeat_interleave(self.n_head // self.n_local_heads, dim=1)
+
+ if self.use_sdpa:
+ if mask is None:
+ y = F.scaled_dot_product_attention(
+ q,
+ k,
+ v,
+ dropout_p=self.dropout if self.training else 0.0,
+ is_causal=True,
+ # No third party attn_mask here to use flash_attention
+ )
+ else:
+ y = F.scaled_dot_product_attention(
+ q,
+ k,
+ v,
+ attn_mask=mask,
+ dropout_p=self.dropout if self.training else 0.0,
+ )
+ else:
+ y = self.eq_scaled_dot_product_attention(
+ q,
+ k,
+ v,
+ attn_mask=mask,
+ dropout_p=self.dropout if self.training else 0.0,
+ )
+
+ y = y.transpose(1, 2).contiguous().view(bsz, seqlen, self.dim)
+
+ return self.wo(y)
+
+ def eq_scaled_dot_product_attention(
+ self,
+ query,
+ key,
+ value,
+ attn_mask=None,
+ dropout_p=0.0,
+ ) -> torch.Tensor:
+ # This is a standard scaled dot product attention
+ # It's low efficient, but it doesn't raise cuda error
+
+ L, S = query.size(-2), key.size(-2)
+ scale_factor = 1 / math.sqrt(query.size(-1))
+ attn_bias = torch.zeros(1, 1, L, S, dtype=query.dtype, device=query.device)
+
+ if attn_mask is not None:
+ if attn_mask.dtype == torch.bool:
+ attn_bias.masked_fill_(attn_mask.logical_not(), float("-inf"))
+ else:
+ attn_bias += attn_mask
+
+ attn_weight = query @ key.transpose(-2, -1) * scale_factor
+ attn_weight += attn_bias
+ attn_weight = torch.softmax(attn_weight, dim=-1)
+ attn_weight = torch.dropout(attn_weight, dropout_p, train=True)
+
+ return attn_weight @ value
+
+
+class FeedForward(nn.Module):
+ def __init__(self, config: BaseModelArgs) -> None:
+ super().__init__()
+ self.w1 = nn.Linear(config.dim, config.intermediate_size, bias=False)
+ self.w3 = nn.Linear(config.dim, config.intermediate_size, bias=False)
+ self.w2 = nn.Linear(config.intermediate_size, config.dim, bias=False)
+ self.dropout = nn.Dropout(p=config.dropout)
+
+ def forward(self, x: Tensor) -> Tensor:
+ return self.w2(self.dropout(F.silu(self.w1(x)) * self.w3(x)))
+
+
+class RMSNorm(nn.Module):
+ def __init__(self, dim: int, eps: float = 1e-5):
+ super().__init__()
+ self.eps = eps
+ self.weight = nn.Parameter(torch.ones(dim))
+
+ def _norm(self, x):
+ return x * torch.rsqrt(torch.mean(x * x, dim=-1, keepdim=True) + self.eps)
+
+ def forward(self, x: Tensor) -> Tensor:
+ output = self._norm(x.float()).type_as(x)
+ return output * self.weight
+
+
+def precompute_freqs_cis(seq_len: int, n_elem: int, base: int = 10000) -> Tensor:
+ freqs = 1.0 / (
+ base ** (torch.arange(0, n_elem, 2)[: (n_elem // 2)].float() / n_elem)
+ )
+ t = torch.arange(seq_len, device=freqs.device)
+ freqs = torch.outer(t, freqs)
+ freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
+ cache = torch.stack([freqs_cis.real, freqs_cis.imag], dim=-1)
+ return cache.to(dtype=torch.bfloat16)
+
+
+def apply_rotary_emb(x: Tensor, freqs_cis: Tensor) -> Tensor:
+ xshaped = x.float().reshape(*x.shape[:-1], -1, 2)
+ freqs_cis = freqs_cis.view(x.size(0), xshaped.size(1), 1, xshaped.size(3), 2)
+ x_out2 = torch.stack(
+ [
+ xshaped[..., 0] * freqs_cis[..., 0] - xshaped[..., 1] * freqs_cis[..., 1],
+ xshaped[..., 1] * freqs_cis[..., 0] + xshaped[..., 0] * freqs_cis[..., 1],
+ ],
+ -1,
+ )
+
+ x_out2 = x_out2.flatten(3)
+ return x_out2.type_as(x)
+
+def top_k_top_p_filtering(
+ logits, top_k=0, top_p=1.0, filter_value=-float("Inf"), min_tokens_to_keep=1
+):
+ """Filter a distribution of logits using top-k and/or nucleus (top-p) filtering
+ Args:
+ logits: logits distribution shape (batch size, vocabulary size)
+ if top_k > 0: keep only top k tokens with highest probability (top-k filtering).
+ if top_p < 1.0: keep the top tokens with cumulative probability >= top_p (nucleus filtering).
+ Nucleus filtering is described in Holtzman et al. (http://arxiv.org/abs/1904.09751)
+ Make sure we keep at least min_tokens_to_keep per batch example in the output
+ From: https://gist.github.com/thomwolf/1a5a29f6962089e871b94cbd09daf317
+ """
+ if top_k > 0:
+ top_k = min(
+ max(top_k, min_tokens_to_keep), logits.size(-1)
+ ) # Safety check
+ # Remove all tokens with a probability less than the last token of the top-k
+ indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
+ logits[indices_to_remove] = filter_value
+
+ if top_p < 1.0:
+ sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+ cumulative_probs = torch.cumsum(
+ F.softmax(sorted_logits, dim=-1), dim=-1
+ )
+
+ # Remove tokens with cumulative probability above the threshold (token with 0 are kept)
+ sorted_indices_to_remove = cumulative_probs > top_p
+ if min_tokens_to_keep > 1:
+ # Keep at least min_tokens_to_keep (set to min_tokens_to_keep-1 because we add the first one below)
+ sorted_indices_to_remove[..., :min_tokens_to_keep] = 0
+ # Shift the indices to the right to keep also the first token above the threshold
+ sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[
+ ..., :-1
+ ].clone()
+ sorted_indices_to_remove[..., 0] = 0
+
+ # scatter sorted tensors to original indexing
+ indices_to_remove = sorted_indices_to_remove.scatter(
+ 1, sorted_indices, sorted_indices_to_remove
+ )
+ logits[indices_to_remove] = filter_value
+ return logits
+
+def topk_sampling(logits, top_k=10, top_p=1.0, temperature=1.0):
+ # temperature: (`optional`) float
+ # The value used to module the next token probabilities. Must be strictly positive. Default to 1.0.
+ # top_k: (`optional`) int
+ # The number of highest probability vocabulary tokens to keep for top-k-filtering. Between 1 and infinity. Default to 50.
+ # top_p: (`optional`) float
+ # The cumulative probability of parameter highest probability vocabulary tokens to keep for nucleus sampling. Must be between 0 and 1. Default to 1.
+
+ # Temperature (higher temperature => more likely to sample low probability tokens)
+ if temperature != 1.0:
+ logits = logits / temperature
+ # Top-p/top-k filtering
+ logits = top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p)
+ # Sample
+ token = torch.multinomial(F.softmax(logits, dim=-1), num_samples=1)
+ logprobs = F.log_softmax(logits.float(), dim=-1)
+ current_logprobs = logprobs[torch.arange(logprobs.shape[0]), token.squeeze(1)]
+ return token, current_logprobs
+
+def sample(
+ logits,
+ previous_tokens: Optional[torch.Tensor] = None,
+ **sampling_kwargs,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+ probs = logits_to_probs(
+ logits=logits[0, -1], previous_tokens=previous_tokens, **sampling_kwargs
+ )
+ idx_next = multinomial_sample_one_no_sync(probs)
+ return idx_next, probs
+
+def multinomial_sample_one_no_sync(
+ probs_sort,
+): # Does multinomial sampling without a cuda synchronization
+ q = torch.empty_like(probs_sort).exponential_(1)
+ return torch.argmax(probs_sort / q, dim=-1, keepdim=True).to(dtype=torch.int)
+
+
+def logits_to_probs(
+ logits,
+ previous_tokens: Optional[torch.Tensor] = None,
+ suppress_tokens: Optional[List[int]] = None,
+ temperature: torch.Tensor = 0.7,
+ top_p: torch.Tensor = 0.7,
+ repetition_penalty: torch.Tensor = 1.5,
+) -> torch.Tensor:
+ # Apply repetition penalty
+ if previous_tokens is not None:
+ previous_tokens = previous_tokens.long()
+ score = torch.gather(logits, dim=0, index=previous_tokens)
+ score = torch.where(
+ score < 0, score * repetition_penalty, score / repetition_penalty
+ )
+ logits.scatter_(dim=0, index=previous_tokens, src=score)
+ if suppress_tokens is not None:
+ for token in suppress_tokens:
+ logits[token] = -float("Inf")
+
+ # Apply top-p sampling
+ sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+ cum_probs = torch.cumsum(torch.nn.functional.softmax(sorted_logits, dim=-1), dim=-1)
+ sorted_indices_to_remove = cum_probs > top_p
+ sorted_indices_to_remove[0] = False # keep at least one option
+ indices_to_remove = sorted_indices_to_remove.scatter(
+ dim=0, index=sorted_indices, src=sorted_indices_to_remove
+ )
+ logits = logits.masked_fill(indices_to_remove, -float("Inf"))
+
+ logits = logits / max(temperature, 1e-5)
+
+ probs = torch.nn.functional.softmax(logits, dim=-1)
+ return probs
diff --git a/seed-vc/modules/v2/cfm.py b/seed-vc/modules/v2/cfm.py
new file mode 100644
index 0000000000000000000000000000000000000000..929d557b99aaabec19012c810e833aaa922470d5
--- /dev/null
+++ b/seed-vc/modules/v2/cfm.py
@@ -0,0 +1,173 @@
+import torch
+from tqdm import tqdm
+
+class CFM(torch.nn.Module):
+ def __init__(
+ self,
+ estimator: torch.nn.Module,
+ ):
+ super().__init__()
+ self.sigma_min = 1e-6
+ self.estimator = estimator
+ self.in_channels = estimator.in_channels
+ self.criterion = torch.nn.L1Loss()
+
+ @torch.inference_mode()
+ def inference(self,
+ mu: torch.Tensor,
+ x_lens: torch.Tensor,
+ prompt: torch.Tensor,
+ style: torch.Tensor,
+ n_timesteps=10,
+ temperature=1.0,
+ inference_cfg_rate=[0.5, 0.5],
+ random_voice=False,
+ ):
+ """Forward diffusion
+
+ Args:
+ mu (torch.Tensor): output of encoder
+ shape: (batch_size, n_feats, mel_timesteps)
+ x_lens (torch.Tensor): length of each mel-spectrogram
+ shape: (batch_size,)
+ prompt (torch.Tensor): prompt
+ shape: (batch_size, n_feats, prompt_len)
+ style (torch.Tensor): style
+ shape: (batch_size, style_dim)
+ n_timesteps (int): number of diffusion steps
+ temperature (float, optional): temperature for scaling noise. Defaults to 1.0.
+ inference_cfg_rate (float, optional): Classifier-Free Guidance inference introduced in VoiceBox. Defaults to 0.5.
+
+ Returns:
+ sample: generated mel-spectrogram
+ shape: (batch_size, n_feats, mel_timesteps)
+ """
+ B, T = mu.size(0), mu.size(1)
+ z = torch.randn([B, self.in_channels, T], device=mu.device) * temperature
+ t_span = torch.linspace(0, 1, n_timesteps + 1, device=mu.device)
+ t_span = t_span + (-1) * (torch.cos(torch.pi / 2 * t_span) - 1 + t_span)
+ return self.solve_euler(z, x_lens, prompt, mu, style, t_span, inference_cfg_rate, random_voice)
+ def solve_euler(self, x, x_lens, prompt, mu, style, t_span, inference_cfg_rate=[0.5, 0.5], random_voice=False,):
+ """
+ Fixed euler solver for ODEs.
+ Args:
+ x (torch.Tensor): random noise
+ t_span (torch.Tensor): n_timesteps interpolated
+ shape: (n_timesteps + 1,)
+ mu (torch.Tensor): output of encoder
+ shape: (batch_size, n_feats, mel_timesteps)
+ x_lens (torch.Tensor): length of each mel-spectrogram
+ shape: (batch_size,)
+ prompt (torch.Tensor): prompt
+ shape: (batch_size, n_feats, prompt_len)
+ style (torch.Tensor): style
+ shape: (batch_size, style_dim)
+ inference_cfg_rate (float, optional): Classifier-Free Guidance inference introduced in VoiceBox. Defaults to 0.5.
+ sway_sampling (bool, optional): Sway sampling. Defaults to False.
+ amo_sampling (bool, optional): AMO sampling. Defaults to False.
+ """
+ t, _, dt = t_span[0], t_span[-1], t_span[1] - t_span[0]
+
+ # apply prompt
+ prompt_len = prompt.size(-1)
+ prompt_x = torch.zeros_like(x)
+ prompt_x[..., :prompt_len] = prompt[..., :prompt_len]
+ x[..., :prompt_len] = 0
+ for step in tqdm(range(1, len(t_span))):
+ if random_voice:
+ cfg_dphi_dt = self.estimator(
+ torch.cat([x, x], dim=0),
+ torch.cat([torch.zeros_like(prompt_x), torch.zeros_like(prompt_x)], dim=0),
+ torch.cat([x_lens, x_lens], dim=0),
+ torch.cat([t.unsqueeze(0), t.unsqueeze(0)], dim=0),
+ torch.cat([torch.zeros_like(style), torch.zeros_like(style)], dim=0),
+ torch.cat([mu, torch.zeros_like(mu)], dim=0),
+ )
+ cond_txt, uncond = cfg_dphi_dt[0:1], cfg_dphi_dt[1:2]
+ dphi_dt = ((1.0 + inference_cfg_rate[0]) * cond_txt - inference_cfg_rate[0] * uncond)
+ elif all(i == 0 for i in inference_cfg_rate):
+ dphi_dt = self.estimator(x, prompt_x, x_lens, t.unsqueeze(0), style, mu)
+ elif inference_cfg_rate[0] == 0:
+ # Classifier-Free Guidance inference introduced in VoiceBox
+ cfg_dphi_dt = self.estimator(
+ torch.cat([x, x], dim=0),
+ torch.cat([prompt_x, torch.zeros_like(prompt_x)], dim=0),
+ torch.cat([x_lens, x_lens], dim=0),
+ torch.cat([t.unsqueeze(0), t.unsqueeze(0)], dim=0),
+ torch.cat([style, torch.zeros_like(style)], dim=0),
+ torch.cat([mu, mu], dim=0),
+ )
+ cond_txt_spk, cond_txt = cfg_dphi_dt[0:1], cfg_dphi_dt[1:2]
+ dphi_dt = ((1.0 + inference_cfg_rate[1]) * cond_txt_spk - inference_cfg_rate[1] * cond_txt)
+ elif inference_cfg_rate[1] == 0:
+ cfg_dphi_dt = self.estimator(
+ torch.cat([x, x], dim=0),
+ torch.cat([prompt_x, torch.zeros_like(prompt_x)], dim=0),
+ torch.cat([x_lens, x_lens], dim=0),
+ torch.cat([t.unsqueeze(0), t.unsqueeze(0)], dim=0),
+ torch.cat([style, torch.zeros_like(style)], dim=0),
+ torch.cat([mu, torch.zeros_like(mu)], dim=0),
+ )
+ cond_txt_spk, uncond = cfg_dphi_dt[0:1], cfg_dphi_dt[1:2]
+ dphi_dt = ((1.0 + inference_cfg_rate[0]) * cond_txt_spk - inference_cfg_rate[0] * uncond)
+ else:
+ # Multi-condition Classifier-Free Guidance inference introduced in MegaTTS3
+ cfg_dphi_dt = self.estimator(
+ torch.cat([x, x, x], dim=0),
+ torch.cat([prompt_x, torch.zeros_like(prompt_x), torch.zeros_like(prompt_x)], dim=0),
+ torch.cat([x_lens, x_lens, x_lens], dim=0),
+ torch.cat([t.unsqueeze(0), t.unsqueeze(0), t.unsqueeze(0)], dim=0),
+ torch.cat([style, torch.zeros_like(style), torch.zeros_like(style)], dim=0),
+ torch.cat([mu, mu, torch.zeros_like(mu)], dim=0),
+ )
+ cond_txt_spk, cond_txt, uncond = cfg_dphi_dt[0:1], cfg_dphi_dt[1:2], cfg_dphi_dt[2:3]
+ dphi_dt = (1.0 + inference_cfg_rate[0] + inference_cfg_rate[1]) * cond_txt_spk - \
+ inference_cfg_rate[0] * uncond - inference_cfg_rate[1] * cond_txt
+ x = x + dt * dphi_dt
+ t = t + dt
+ if step < len(t_span) - 1:
+ dt = t_span[step + 1] - t
+ x[:, :, :prompt_len] = 0
+
+ return x
+
+ def forward(self, x1, x_lens, prompt_lens, mu, style):
+ """Computes diffusion loss
+
+ Args:
+ x1 (torch.Tensor): Target
+ shape: (batch_size, n_feats, mel_timesteps)
+ mask (torch.Tensor): target mask
+ shape: (batch_size, 1, mel_timesteps)
+ mu (torch.Tensor): output of encoder
+ shape: (batch_size, n_feats, mel_timesteps)
+ spks (torch.Tensor, optional): speaker embedding. Defaults to None.
+ shape: (batch_size, spk_emb_dim)
+
+ Returns:
+ loss: conditional flow matching loss
+ y: conditional flow
+ shape: (batch_size, n_feats, mel_timesteps)
+ """
+ b, _, t = x1.shape
+
+ # random timestep
+ t = torch.rand([b, 1, 1], device=mu.device, dtype=x1.dtype)
+ # sample noise p(x_0)
+ z = torch.randn_like(x1)
+
+ y = (1 - (1 - self.sigma_min) * t) * z + t * x1
+ u = x1 - (1 - self.sigma_min) * z
+ prompt = torch.zeros_like(x1)
+ for bib in range(b):
+ prompt[bib, :, :prompt_lens[bib]] = x1[bib, :, :prompt_lens[bib]]
+ # range covered by prompt are set to 0
+ y[bib, :, :prompt_lens[bib]] = 0
+
+ estimator_out = self.estimator(y, prompt, x_lens, t.squeeze(), style, mu)
+ loss = 0
+ for bib in range(b):
+ loss += self.criterion(estimator_out[bib, :, prompt_lens[bib]:x_lens[bib]], u[bib, :, prompt_lens[bib]:x_lens[bib]])
+ loss /= b
+
+ return loss
diff --git a/seed-vc/modules/v2/dit_model.py b/seed-vc/modules/v2/dit_model.py
new file mode 100644
index 0000000000000000000000000000000000000000..4374ac86a4d4d0869788cdd16087115c4418ba5f
--- /dev/null
+++ b/seed-vc/modules/v2/dit_model.py
@@ -0,0 +1,250 @@
+# 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.
+from dataclasses import dataclass
+from typing import Optional, Union, Tuple, List
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+from torch.nn import functional as F
+import time
+
+def find_multiple(n: int, k: int) -> int:
+ if n % k == 0:
+ return n
+ return n + k - (n % k)
+
+class AdaptiveLayerNorm(nn.Module):
+ r"""Adaptive Layer Normalization"""
+
+ def __init__(self, d_model, norm) -> None:
+ super(AdaptiveLayerNorm, self).__init__()
+ self.linear = nn.Linear(d_model, 6 * d_model)
+ self.act = nn.SiLU()
+ self.norm = norm
+ self.d_model = d_model
+ self.eps = self.norm.eps
+
+ def forward(self, x: Tensor, emb: Tensor) -> Tuple[Tensor]:
+ emb = self.linear(self.act(emb))
+ shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = torch.chunk(emb, 6, dim=-1)
+
+ x = self.norm(x) * (1 + scale_msa) + shift_msa
+ return x, gate_msa, shift_mlp, scale_mlp, gate_mlp
+
+class AdaptiveLayerNormFinal(nn.Module):
+ r"""Adaptive Layer Normalization"""
+
+ def __init__(self, d_model, norm) -> None:
+ super(AdaptiveLayerNormFinal, self).__init__()
+ self.linear = nn.Linear(d_model, 2 * d_model)
+ self.act = nn.SiLU()
+ self.norm = norm
+ self.d_model = d_model
+ self.eps = self.norm.eps
+
+ def forward(self, x: Tensor, emb: Tensor) -> Tuple[Tensor]:
+ emb = self.linear(self.act(emb))
+ scale, shift = torch.chunk(emb, 2, dim=-1)
+
+ x = self.norm(x) * (1 + scale) + shift
+ return x
+
+@dataclass
+class ModelArgs:
+ block_size: int = 2048
+ vocab_size: int = 32000
+ n_layer: int = 32
+ n_head: int = 32
+ dim: int = 4096
+ intermediate_size: int = None
+ n_local_heads: int = -1
+ head_dim: int = 64
+ rope_base: float = 10000
+ norm_eps: float = 1e-5
+ uvit_skip_connection: bool = False
+ time_as_token: bool = False
+ dropout_rate: float = 0.1
+ attn_dropout_rate: float = 0.1
+
+ def __post_init__(self):
+ if self.n_local_heads == -1:
+ self.n_local_heads = self.n_head
+ if self.intermediate_size is None:
+ hidden_dim = 4 * self.dim
+ n_hidden = int(2 * hidden_dim / 3)
+ self.intermediate_size = find_multiple(n_hidden, 256)
+ # self.head_dim = self.dim // self.n_head
+
+class Transformer(nn.Module):
+ def __init__(self, config: ModelArgs) -> None:
+ super().__init__()
+ self.config = config
+
+ self.layers = nn.ModuleList(TransformerBlock(config) for _ in range(config.n_layer))
+ self.norm = AdaptiveLayerNormFinal(config.dim, RMSNorm(config.dim, eps=config.norm_eps))
+
+ self.max_batch_size = -1
+ self.max_seq_length = config.block_size
+
+ self.uvit_skip_connection = self.config.uvit_skip_connection
+ if self.uvit_skip_connection:
+ self.layers_emit_skip = [i for i in range(self.config.n_layer) if i < self.config.n_layer // 2]
+ self.layers_receive_skip = [i for i in range(self.config.n_layer) if i > self.config.n_layer // 2]
+ else:
+ self.layers_emit_skip = []
+ self.layers_receive_skip = []
+ freqs_cis = precompute_freqs_cis(self.config.block_size, self.config.head_dim,
+ self.config.rope_base)
+ self.register_buffer("freqs_cis", freqs_cis)
+
+ causal_mask = torch.tril(
+ torch.ones(self.max_seq_length, self.max_seq_length, dtype=torch.bool)
+ )
+ self.register_buffer("causal_mask", causal_mask)
+
+ def forward(self,
+ x: Tensor,
+ c: Tensor,
+ input_pos: Optional[Tensor] = None,
+ mask: Optional[Tensor] = None,
+ ) -> Tensor:
+ mask = mask[..., input_pos]
+ freqs_cis = self.freqs_cis[input_pos]
+ for i, layer in enumerate(self.layers):
+ x = layer(x, c, freqs_cis, mask)
+ x = self.norm(x, c)
+ return x
+
+
+class TransformerBlock(nn.Module):
+ def __init__(self, config: ModelArgs) -> None:
+ super().__init__()
+ self.attention = Attention(config)
+ self.feed_forward = FeedForward(config)
+ self.ffn_norm = RMSNorm(config.dim, eps=config.norm_eps)
+ self.attention_norm = AdaptiveLayerNorm(config.dim, RMSNorm(config.dim, eps=config.norm_eps))
+ def forward(self,
+ x: Tensor,
+ c: Tensor,
+ freqs_cis: Tensor,
+ mask: Tensor,
+ ) -> Tensor:
+ normed_x, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attention_norm(x, emb=c)
+ # attention
+ attn_output = self.attention(normed_x, freqs_cis, mask)
+ x = x + gate_msa * attn_output
+ normed_x = self.ffn_norm(x) * (1 + scale_mlp) + shift_mlp
+ ff_output = self.feed_forward(normed_x)
+ x = x + gate_mlp * ff_output
+ return x
+
+
+class Attention(nn.Module):
+ def __init__(self, config: ModelArgs, is_cross_attention: bool = False):
+ super().__init__()
+ assert config.dim % config.n_head == 0
+
+ total_head_dim = (config.n_head + 2 * config.n_local_heads) * config.head_dim
+ # key, query, value projections for all heads, but in a batch
+ if is_cross_attention:
+ self.wq = nn.Linear(config.dim, config.n_head * config.head_dim, bias=False)
+ self.wkv = nn.Linear(config.context_dim, 2 * config.n_local_heads * config.head_dim, bias=False)
+ else:
+ self.wqkv = nn.Linear(config.dim, total_head_dim, bias=False)
+ self.wo = nn.Linear(config.head_dim * config.n_head, config.dim, bias=False)
+ self.kv_cache = None
+
+ self.n_head = config.n_head
+ self.head_dim = config.head_dim
+ self.n_local_heads = config.n_local_heads
+ self.dim = config.dim
+ self.attn_dropout_rate = config.attn_dropout_rate
+
+ def forward(self,
+ x: Tensor,
+ freqs_cis: Tensor,
+ mask: Tensor,
+ context: Optional[Tensor] = None,
+ context_freqs_cis: Optional[Tensor] = None,
+ ) -> Tensor:
+ bsz, seqlen, _ = x.shape
+
+ kv_size = self.n_local_heads * self.head_dim
+ q, k, v = self.wqkv(x).split([kv_size, kv_size, kv_size], dim=-1)
+ context_seqlen = seqlen
+
+ q = q.view(bsz, seqlen, self.n_head, self.head_dim)
+ k = k.view(bsz, context_seqlen, self.n_local_heads, self.head_dim)
+ v = v.view(bsz, context_seqlen, self.n_local_heads, self.head_dim)
+
+ q = apply_rotary_emb(q, freqs_cis)
+ k = apply_rotary_emb(k, freqs_cis)
+
+ q, k, v = map(lambda x: x.transpose(1, 2), (q, k, v))
+
+ k = k.repeat_interleave(self.n_head // self.n_local_heads, dim=1)
+ v = v.repeat_interleave(self.n_head // self.n_local_heads, dim=1)
+ y = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0)
+
+ y = y.transpose(1, 2).contiguous().view(bsz, seqlen, self.head_dim * self.n_head)
+
+ y = self.wo(y)
+ return y
+
+
+class FeedForward(nn.Module):
+ def __init__(self, config: ModelArgs) -> None:
+ super().__init__()
+ self.w1 = nn.Linear(config.dim, config.intermediate_size, bias=False)
+ self.w3 = nn.Linear(config.dim, config.intermediate_size, bias=False)
+ self.w2 = nn.Linear(config.intermediate_size, config.dim, bias=False)
+ self.dropout = nn.Dropout(config.dropout_rate)
+
+ def forward(self, x: Tensor) -> Tensor:
+ return self.w2(self.dropout(F.silu(self.w1(x)) * self.w3(x)))
+
+
+class RMSNorm(nn.Module):
+ def __init__(self, dim: int, eps: float = 1e-5):
+ super().__init__()
+ self.eps = eps
+ self.weight = nn.Parameter(torch.ones(dim))
+
+ def _norm(self, x):
+ return x * torch.rsqrt(torch.mean(x * x, dim=-1, keepdim=True) + self.eps)
+
+ def forward(self, x: Tensor) -> Tensor:
+ output = self._norm(x.float()).type_as(x)
+ return output * self.weight
+
+
+def precompute_freqs_cis(
+ seq_len: int, n_elem: int, base: int = 10000,
+ dtype: torch.dtype = torch.bfloat16
+) -> Tensor:
+ freqs = 1.0 / (base ** (torch.arange(0, n_elem, 2)[: (n_elem // 2)].float() / n_elem))
+ t = torch.arange(seq_len, device=freqs.device)
+ freqs = torch.outer(t, freqs)
+ freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
+ cache = torch.stack([freqs_cis.real, freqs_cis.imag], dim=-1)
+ return cache.to(dtype=dtype)
+
+
+def apply_rotary_emb(x: Tensor, freqs_cis: Tensor) -> Tensor:
+ xshaped = x.float().reshape(*x.shape[:-1], -1, 2)
+ freqs_cis = freqs_cis.view(1, xshaped.size(1), 1, xshaped.size(3), 2)
+ x_out2 = torch.stack(
+ [
+ xshaped[..., 0] * freqs_cis[..., 0] - xshaped[..., 1] * freqs_cis[..., 1],
+ xshaped[..., 1] * freqs_cis[..., 0] + xshaped[..., 0] * freqs_cis[..., 1],
+ ],
+ -1,
+ )
+
+ x_out2 = x_out2.flatten(3)
+ return x_out2.type_as(x)
+
diff --git a/seed-vc/modules/v2/dit_wrapper.py b/seed-vc/modules/v2/dit_wrapper.py
new file mode 100644
index 0000000000000000000000000000000000000000..3ac611ab9e33114df905a80e3f847e7ab198c72a
--- /dev/null
+++ b/seed-vc/modules/v2/dit_wrapper.py
@@ -0,0 +1,152 @@
+import torch
+from torch import nn
+import math
+
+from modules.v2.dit_model import ModelArgs, Transformer
+from modules.commons import sequence_mask
+
+from torch.nn.utils import weight_norm
+
+def modulate(x, shift, scale):
+ return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+
+#################################################################################
+# Embedding Layers for Timesteps and Class Labels #
+#################################################################################
+
+class TimestepEmbedder(nn.Module):
+ """
+ Embeds scalar timesteps into vector representations.
+ """
+ def __init__(self, hidden_size, frequency_embedding_size=256):
+ super().__init__()
+ self.mlp = nn.Sequential(
+ nn.Linear(frequency_embedding_size, hidden_size, bias=True),
+ nn.SiLU(),
+ nn.Linear(hidden_size, hidden_size, bias=True),
+ )
+ self.frequency_embedding_size = frequency_embedding_size
+
+ @staticmethod
+ def timestep_embedding(t, dim, max_period=10000, scale=1000):
+ """
+ Create sinusoidal timestep embeddings.
+ :param t: a 1-D Tensor of N indices, one per batch element.
+ These may be fractional.
+ :param dim: the dimension of the output.
+ :param max_period: controls the minimum frequency of the embeddings.
+ :return: an (N, D) Tensor of positional embeddings.
+ """
+ # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
+ half = dim // 2
+ freqs = torch.exp(
+ -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
+ ).to(device=t.device)
+ args = scale * t[:, None].float() * freqs[None]
+ embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+ if dim % 2:
+ embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+ return embedding
+
+ def forward(self, t):
+ t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
+ t_emb = self.mlp(t_freq)
+ return t_emb
+
+
+class DiT(torch.nn.Module):
+ def __init__(
+ self,
+ time_as_token,
+ style_as_token,
+ uvit_skip_connection,
+ block_size,
+ depth,
+ num_heads,
+ hidden_dim,
+ in_channels,
+ content_dim,
+ style_encoder_dim,
+ class_dropout_prob,
+ dropout_rate,
+ attn_dropout_rate,
+ ):
+ super(DiT, self).__init__()
+ self.time_as_token = time_as_token
+ self.style_as_token = style_as_token
+ self.uvit_skip_connection = uvit_skip_connection
+ model_args = ModelArgs(
+ block_size=block_size,
+ n_layer=depth,
+ n_head=num_heads,
+ dim=hidden_dim,
+ head_dim=hidden_dim // num_heads,
+ vocab_size=1, # we don't use this
+ uvit_skip_connection=self.uvit_skip_connection,
+ time_as_token=self.time_as_token,
+ dropout_rate=dropout_rate,
+ attn_dropout_rate=attn_dropout_rate,
+ )
+ self.transformer = Transformer(model_args)
+ self.in_channels = in_channels
+ self.out_channels = in_channels
+ self.num_heads = num_heads
+
+ self.x_embedder = weight_norm(nn.Linear(in_channels, hidden_dim, bias=True))
+
+ self.content_dim = content_dim # for continuous content
+ self.cond_projection = nn.Linear(content_dim, hidden_dim, bias=True) # continuous content
+
+ self.t_embedder = TimestepEmbedder(hidden_dim)
+
+ self.final_mlp = nn.Sequential(
+ nn.Linear(hidden_dim, hidden_dim),
+ nn.SiLU(),
+ nn.Linear(hidden_dim, in_channels),
+ )
+
+ self.class_dropout_prob = class_dropout_prob
+
+ self.cond_x_merge_linear = nn.Linear(hidden_dim + in_channels + in_channels, hidden_dim)
+ self.style_in = nn.Linear(style_encoder_dim, hidden_dim)
+
+ def forward(self, x, prompt_x, x_lens, t, style, cond):
+ class_dropout = False
+ content_dropout = False
+ if self.training and torch.rand(1) < self.class_dropout_prob:
+ class_dropout = True
+ if self.training and torch.rand(1) < 0.5:
+ content_dropout = True
+ cond_in_module = self.cond_projection
+
+ B, _, T = x.size()
+
+ t1 = self.t_embedder(t) # (N, D)
+ cond = cond_in_module(cond)
+
+ x = x.transpose(1, 2)
+ prompt_x = prompt_x.transpose(1, 2)
+
+ x_in = torch.cat([x, prompt_x, cond], dim=-1)
+ if class_dropout:
+ x_in[..., self.in_channels:self.in_channels*2] = 0
+ if content_dropout:
+ x_in[..., self.in_channels*2:] = 0
+ x_in = self.cond_x_merge_linear(x_in) # (N, T, D)
+
+ style = self.style_in(style)
+ style = torch.zeros_like(style) if class_dropout else style
+ if self.style_as_token:
+ x_in = torch.cat([style.unsqueeze(1), x_in], dim=1)
+ if self.time_as_token:
+ x_in = torch.cat([t1.unsqueeze(1), x_in], dim=1)
+ x_mask = sequence_mask(x_lens + self.style_as_token + self.time_as_token, max_length=x_in.size(1)).to(x.device).unsqueeze(1)
+ input_pos = torch.arange(x_in.size(1)).to(x.device)
+ x_mask_expanded = x_mask[:, None, :].repeat(1, 1, x_in.size(1), 1)
+ x_res = self.transformer(x_in, t1.unsqueeze(1), input_pos, x_mask_expanded)
+ x_res = x_res[:, 1:] if self.time_as_token else x_res
+ x_res = x_res[:, 1:] if self.style_as_token else x_res
+ x = self.final_mlp(x_res)
+ x = x.transpose(1, 2)
+ return x
diff --git a/seed-vc/modules/v2/length_regulator.py b/seed-vc/modules/v2/length_regulator.py
new file mode 100644
index 0000000000000000000000000000000000000000..c8fba22ed62e367ef3abe88b5a04c5ee40658349
--- /dev/null
+++ b/seed-vc/modules/v2/length_regulator.py
@@ -0,0 +1,105 @@
+from typing import Tuple
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+from modules.commons import sequence_mask
+import numpy as np
+
+# f0_bin = 256
+f0_max = 1100.0
+f0_min = 50.0
+f0_mel_min = 1127 * np.log(1 + f0_min / 700)
+f0_mel_max = 1127 * np.log(1 + f0_max / 700)
+
+def f0_to_coarse(f0, f0_bin):
+ f0_mel = 1127 * (1 + f0 / 700).log()
+ a = (f0_bin - 2) / (f0_mel_max - f0_mel_min)
+ b = f0_mel_min * a - 1.
+ f0_mel = torch.where(f0_mel > 0, f0_mel * a - b, f0_mel)
+ # torch.clip_(f0_mel, min=1., max=float(f0_bin - 1))
+ f0_coarse = torch.round(f0_mel).long()
+ f0_coarse = f0_coarse * (f0_coarse > 0)
+ f0_coarse = f0_coarse + ((f0_coarse < 1) * 1)
+ f0_coarse = f0_coarse * (f0_coarse < f0_bin)
+ f0_coarse = f0_coarse + ((f0_coarse >= f0_bin) * (f0_bin - 1))
+ return f0_coarse
+
+class InterpolateRegulator(nn.Module):
+ def __init__(
+ self,
+ channels: int,
+ sampling_ratios: Tuple,
+ is_discrete: bool = False,
+ in_channels: int = None, # only applies to continuous input
+ codebook_size: int = 1024, # for discrete only
+ out_channels: int = None,
+ groups: int = 1,
+ f0_condition: bool = False,
+ n_f0_bins: int = 512,
+ ):
+ super().__init__()
+ self.sampling_ratios = sampling_ratios
+ out_channels = out_channels or channels
+ model = nn.ModuleList([])
+ if len(sampling_ratios) > 0:
+ self.interpolate = True
+ for _ in sampling_ratios:
+ module = nn.Conv1d(channels, channels, 3, 1, 1)
+ norm = nn.GroupNorm(groups, channels)
+ act = nn.Mish()
+ model.extend([module, norm, act])
+ else:
+ self.interpolate = False
+ model.append(
+ nn.Conv1d(channels, out_channels, 1, 1) if channels != out_channels else nn.Identity()
+ )
+ self.model = nn.Sequential(*model)
+ self.embedding = nn.Embedding(codebook_size, channels)
+ self.is_discrete = is_discrete
+
+ self.mask_token = nn.Parameter(torch.zeros(1, channels))
+
+ if f0_condition:
+ self.f0_embedding = nn.Embedding(n_f0_bins, channels)
+ self.f0_condition = f0_condition
+ self.n_f0_bins = n_f0_bins
+ self.f0_bins = torch.arange(2, 1024, 1024 // n_f0_bins)
+ self.f0_mask = nn.Parameter(torch.zeros(1, channels))
+ else:
+ self.f0_condition = False
+
+ if not is_discrete:
+ self.content_in_proj = nn.Linear(in_channels, channels)
+
+ def forward(self, x, ylens=None, f0=None):
+ if self.is_discrete:
+ if len(x.size()) == 2:
+ x = self.embedding(x)
+ else:
+ x = self.embedding(x[:, 0])
+ else:
+ x = self.content_in_proj(x)
+ # x in (B, T, D)
+
+ if self.interpolate:
+ mask = sequence_mask(ylens).unsqueeze(-1)
+ x = F.interpolate(x.transpose(1, 2).contiguous(), size=ylens.max(), mode='nearest')
+ else:
+ x = x.transpose(1, 2).contiguous()
+ mask = None
+ # mask = mask[:, :x.size(2), :]
+ # ylens = ylens.clamp(max=x.size(2)).long()
+ if self.f0_condition:
+ if f0 is None:
+ x = x + self.f0_mask.unsqueeze(-1)
+ else:
+ # quantized_f0 = torch.bucketize(f0, self.f0_bins.to(f0.device)) # (N, T)
+ quantized_f0 = f0_to_coarse(f0, self.n_f0_bins)
+ quantized_f0 = quantized_f0.clamp(0, self.n_f0_bins - 1).long()
+ f0_emb = self.f0_embedding(quantized_f0)
+ f0_emb = F.interpolate(f0_emb.transpose(1, 2).contiguous(), size=ylens.max(), mode='nearest')
+ x = x + f0_emb
+ out = self.model(x).transpose(1, 2).contiguous()
+ out = out * mask if mask is not None else out
+ olens = ylens
+ return out, olens
diff --git a/seed-vc/modules/v2/model.py b/seed-vc/modules/v2/model.py
new file mode 100644
index 0000000000000000000000000000000000000000..c1453ef438c8ef3a22b8ab151db36aa75c93f0ff
--- /dev/null
+++ b/seed-vc/modules/v2/model.py
@@ -0,0 +1,302 @@
+# 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.
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+from torch.nn import functional as F
+
+
+def find_multiple(n: int, k: int) -> int:
+ if n % k == 0:
+ return n
+ return n + k - (n % k)
+
+class AdaptiveLayerNorm(nn.Module):
+ r"""Adaptive Layer Normalization"""
+
+ def __init__(self, d_model, norm) -> None:
+ super(AdaptiveLayerNorm, self).__init__()
+ self.project_layer = nn.Linear(d_model, 2 * d_model)
+ self.norm = norm
+ self.d_model = d_model
+ self.eps = self.norm.eps
+
+ def forward(self, input: Tensor, embedding: Tensor = None) -> Tensor:
+ if embedding is None:
+ return self.norm(input)
+ weight, bias = torch.split(
+ self.project_layer(embedding),
+ split_size_or_sections=self.d_model,
+ dim=-1,
+ )
+ return weight * self.norm(input) + bias
+
+
+@dataclass
+class ModelArgs:
+ block_size: int = 2048
+ vocab_size: int = 32000
+ n_layer: int = 32
+ n_head: int = 32
+ dim: int = 4096
+ intermediate_size: int = None
+ n_local_heads: int = -1
+ head_dim: int = 64
+ rope_base: float = 10000
+ norm_eps: float = 1e-5
+ has_cross_attention: bool = False
+ context_dim: int = 0
+ uvit_skip_connection: bool = False
+ time_as_token: bool = False
+
+ def __post_init__(self):
+ if self.n_local_heads == -1:
+ self.n_local_heads = self.n_head
+ if self.intermediate_size is None:
+ hidden_dim = 4 * self.dim
+ n_hidden = int(2 * hidden_dim / 3)
+ self.intermediate_size = find_multiple(n_hidden, 256)
+ # self.head_dim = self.dim // self.n_head
+
+class Transformer(nn.Module):
+ def __init__(self, config: ModelArgs) -> None:
+ super().__init__()
+ self.config = config
+
+ self.layers = nn.ModuleList(TransformerBlock(config) for _ in range(config.n_layer))
+ self.norm = AdaptiveLayerNorm(config.dim, RMSNorm(config.dim, eps=config.norm_eps))
+
+ self.freqs_cis: Optional[Tensor] = None
+ self.mask_cache: Optional[Tensor] = None
+ self.max_batch_size = -1
+ self.max_seq_length = -1
+
+ def setup_caches(self, max_batch_size, max_seq_length, use_kv_cache=False):
+ if self.max_seq_length >= max_seq_length and self.max_batch_size >= max_batch_size:
+ return
+ head_dim = self.config.dim // self.config.n_head
+ max_seq_length = find_multiple(max_seq_length, 8)
+ self.max_seq_length = max_seq_length
+ self.max_batch_size = max_batch_size
+ dtype = self.norm.project_layer.weight.dtype
+ device = self.norm.project_layer.weight.device
+
+ self.freqs_cis = precompute_freqs_cis(self.config.block_size, self.config.head_dim,
+ self.config.rope_base, dtype).to(device)
+ self.causal_mask = torch.tril(torch.ones(self.max_seq_length, self.max_seq_length, dtype=torch.bool)).to(device)
+ self.use_kv_cache = use_kv_cache
+ self.uvit_skip_connection = self.config.uvit_skip_connection
+ if self.uvit_skip_connection:
+ self.layers_emit_skip = [i for i in range(self.config.n_layer) if i < self.config.n_layer // 2]
+ self.layers_receive_skip = [i for i in range(self.config.n_layer) if i > self.config.n_layer // 2]
+ else:
+ self.layers_emit_skip = []
+ self.layers_receive_skip = []
+
+ def forward(self,
+ x: Tensor,
+ c: Tensor,
+ input_pos: Optional[Tensor] = None,
+ mask: Optional[Tensor] = None,
+ context: Optional[Tensor] = None,
+ context_input_pos: Optional[Tensor] = None,
+ cross_attention_mask: Optional[Tensor] = None,
+ ) -> Tensor:
+ assert self.freqs_cis is not None, "Caches must be initialized first"
+ if mask is None: # in case of non-causal model
+ if not self.training and self.use_kv_cache:
+ mask = self.causal_mask[None, None, input_pos]
+ else:
+ mask = self.causal_mask[None, None, input_pos]
+ mask = mask[..., input_pos]
+ freqs_cis = self.freqs_cis[input_pos]
+ if context is not None:
+ context_freqs_cis = self.freqs_cis[context_input_pos]
+ else:
+ context_freqs_cis = None
+ skip_in_x_list = []
+ for i, layer in enumerate(self.layers):
+ if self.uvit_skip_connection and i in self.layers_receive_skip:
+ skip_in_x = skip_in_x_list.pop(-1)
+ else:
+ skip_in_x = None
+ x = layer(x, c, input_pos, freqs_cis, mask, context, context_freqs_cis, cross_attention_mask, skip_in_x)
+ if self.uvit_skip_connection and i in self.layers_emit_skip:
+ skip_in_x_list.append(x)
+ x = self.norm(x, c)
+ return x
+
+ @classmethod
+ def from_name(cls, name: str):
+ return cls(ModelArgs.from_name(name))
+
+
+class TransformerBlock(nn.Module):
+ def __init__(self, config: ModelArgs) -> None:
+ super().__init__()
+ self.attention = Attention(config)
+ self.feed_forward = FeedForward(config)
+ self.ffn_norm = AdaptiveLayerNorm(config.dim, RMSNorm(config.dim, eps=config.norm_eps))
+ self.attention_norm = AdaptiveLayerNorm(config.dim, RMSNorm(config.dim, eps=config.norm_eps))
+
+ if config.has_cross_attention:
+ self.has_cross_attention = True
+ self.cross_attention = Attention(config, is_cross_attention=True)
+ self.cross_attention_norm = AdaptiveLayerNorm(config.dim, RMSNorm(config.dim, eps=config.norm_eps))
+ else:
+ self.has_cross_attention = False
+
+ if config.uvit_skip_connection:
+ self.skip_in_linear = nn.Linear(config.dim * 2, config.dim)
+ self.uvit_skip_connection = True
+ else:
+ self.uvit_skip_connection = False
+
+ self.time_as_token = config.time_as_token
+
+ def forward(self,
+ x: Tensor,
+ c: Tensor,
+ input_pos: Tensor,
+ freqs_cis: Tensor,
+ mask: Tensor,
+ context: Optional[Tensor] = None,
+ context_freqs_cis: Optional[Tensor] = None,
+ cross_attention_mask: Optional[Tensor] = None,
+ skip_in_x: Optional[Tensor] = None,
+ ) -> Tensor:
+ c = None if self.time_as_token else c
+ if self.uvit_skip_connection and skip_in_x is not None:
+ x = self.skip_in_linear(torch.cat([x, skip_in_x], dim=-1))
+ h = x + self.attention(self.attention_norm(x, c), freqs_cis, mask, input_pos)
+ if self.has_cross_attention:
+ h = h + self.cross_attention(self.cross_attention_norm(h, c), freqs_cis, cross_attention_mask, input_pos, context, context_freqs_cis)
+ out = h + self.feed_forward(self.ffn_norm(h, c))
+ return out
+
+
+class Attention(nn.Module):
+ def __init__(self, config: ModelArgs, is_cross_attention: bool = False):
+ super().__init__()
+ assert config.dim % config.n_head == 0
+
+ total_head_dim = (config.n_head + 2 * config.n_local_heads) * config.head_dim
+ # key, query, value projections for all heads, but in a batch
+ if is_cross_attention:
+ self.wq = nn.Linear(config.dim, config.n_head * config.head_dim, bias=False)
+ self.wkv = nn.Linear(config.context_dim, 2 * config.n_local_heads * config.head_dim, bias=False)
+ else:
+ self.wqkv = nn.Linear(config.dim, total_head_dim, bias=False)
+ self.wo = nn.Linear(config.head_dim * config.n_head, config.dim, bias=False)
+ self.kv_cache = None
+
+ self.n_head = config.n_head
+ self.head_dim = config.head_dim
+ self.n_local_heads = config.n_local_heads
+ self.dim = config.dim
+ # self._register_load_state_dict_pre_hook(self.load_hook)
+
+ # def load_hook(self, state_dict, prefix, *args):
+ # if prefix + "wq.weight" in state_dict:
+ # wq = state_dict.pop(prefix + "wq.weight")
+ # wk = state_dict.pop(prefix + "wk.weight")
+ # wv = state_dict.pop(prefix + "wv.weight")
+ # state_dict[prefix + "wqkv.weight"] = torch.cat([wq, wk, wv])
+
+ def forward(self,
+ x: Tensor,
+ freqs_cis: Tensor,
+ mask: Tensor,
+ input_pos: Optional[Tensor] = None,
+ context: Optional[Tensor] = None,
+ context_freqs_cis: Optional[Tensor] = None,
+ ) -> Tensor:
+ bsz, seqlen, _ = x.shape
+
+ kv_size = self.n_local_heads * self.head_dim
+ if context is None:
+ q, k, v = self.wqkv(x).split([kv_size, kv_size, kv_size], dim=-1)
+ context_seqlen = seqlen
+ else:
+ q = self.wq(x)
+ k, v = self.wkv(context).split([kv_size, kv_size], dim=-1)
+ context_seqlen = context.shape[1]
+
+ q = q.view(bsz, seqlen, self.n_head, self.head_dim)
+ k = k.view(bsz, context_seqlen, self.n_local_heads, self.head_dim)
+ v = v.view(bsz, context_seqlen, self.n_local_heads, self.head_dim)
+
+ q = apply_rotary_emb(q, freqs_cis)
+ k = apply_rotary_emb(k, context_freqs_cis if context_freqs_cis is not None else freqs_cis)
+
+ q, k, v = map(lambda x: x.transpose(1, 2), (q, k, v))
+
+ if self.kv_cache is not None:
+ k, v = self.kv_cache.update(input_pos, k, v)
+
+ k = k.repeat_interleave(self.n_head // self.n_local_heads, dim=1)
+ v = v.repeat_interleave(self.n_head // self.n_local_heads, dim=1)
+ y = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0)
+
+ y = y.transpose(1, 2).contiguous().view(bsz, seqlen, self.head_dim * self.n_head)
+
+ y = self.wo(y)
+ return y
+
+
+class FeedForward(nn.Module):
+ def __init__(self, config: ModelArgs) -> None:
+ super().__init__()
+ self.w1 = nn.Linear(config.dim, config.intermediate_size, bias=False)
+ self.w3 = nn.Linear(config.dim, config.intermediate_size, bias=False)
+ self.w2 = nn.Linear(config.intermediate_size, config.dim, bias=False)
+
+ def forward(self, x: Tensor) -> Tensor:
+ return self.w2(F.silu(self.w1(x)) * self.w3(x))
+
+
+class RMSNorm(nn.Module):
+ def __init__(self, dim: int, eps: float = 1e-5):
+ super().__init__()
+ self.eps = eps
+ self.weight = nn.Parameter(torch.ones(dim))
+
+ def _norm(self, x):
+ return x * torch.rsqrt(torch.mean(x * x, dim=-1, keepdim=True) + self.eps)
+
+ def forward(self, x: Tensor) -> Tensor:
+ output = self._norm(x.float()).type_as(x)
+ return output * self.weight
+
+
+def precompute_freqs_cis(
+ seq_len: int, n_elem: int, base: int = 10000,
+ dtype: torch.dtype = torch.bfloat16
+) -> Tensor:
+ freqs = 1.0 / (base ** (torch.arange(0, n_elem, 2)[: (n_elem // 2)].float() / n_elem))
+ t = torch.arange(seq_len, device=freqs.device)
+ freqs = torch.outer(t, freqs)
+ freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
+ cache = torch.stack([freqs_cis.real, freqs_cis.imag], dim=-1)
+ return cache.to(dtype=dtype)
+
+
+def apply_rotary_emb(x: Tensor, freqs_cis: Tensor) -> Tensor:
+ xshaped = x.float().reshape(*x.shape[:-1], -1, 2)
+ freqs_cis = freqs_cis.view(1, xshaped.size(1), 1, xshaped.size(3), 2)
+ x_out2 = torch.stack(
+ [
+ xshaped[..., 0] * freqs_cis[..., 0] - xshaped[..., 1] * freqs_cis[..., 1],
+ xshaped[..., 1] * freqs_cis[..., 0] + xshaped[..., 0] * freqs_cis[..., 1],
+ ],
+ -1,
+ )
+
+ x_out2 = x_out2.flatten(3)
+ return x_out2.type_as(x)
diff --git a/seed-vc/modules/v2/vc_wrapper.py b/seed-vc/modules/v2/vc_wrapper.py
new file mode 100644
index 0000000000000000000000000000000000000000..26e2dd0c5ecbbb0e8c19cea2c8472450148eb97e
--- /dev/null
+++ b/seed-vc/modules/v2/vc_wrapper.py
@@ -0,0 +1,664 @@
+import torch
+import librosa
+import torchaudio
+import numpy as np
+from pydub import AudioSegment
+from hf_utils import load_custom_model_from_hf
+
+DEFAULT_REPO_ID = "Plachta/Seed-VC"
+DEFAULT_CFM_CHECKPOINT = "v2/cfm_small.pth"
+DEFAULT_AR_CHECKPOINT = "v2/ar_base.pth"
+
+DEFAULT_CE_REPO_ID = "Plachta/ASTRAL-quantization"
+DEFAULT_CE_NARROW_CHECKPOINT = "bsq32/bsq32_light.pth"
+DEFAULT_CE_WIDE_CHECKPOINT = "bsq2048/bsq2048_light.pth"
+
+DEFAULT_SE_REPO_ID = "funasr/campplus"
+DEFAULT_SE_CHECKPOINT = "campplus_cn_common.bin"
+
+class VoiceConversionWrapper(torch.nn.Module):
+ def __init__(
+ self,
+ sr: int,
+ hop_size: int,
+ mel_fn: callable,
+ cfm: torch.nn.Module,
+ cfm_length_regulator: torch.nn.Module,
+ content_extractor_narrow: torch.nn.Module,
+ content_extractor_wide: torch.nn.Module,
+ ar_length_regulator: torch.nn.Module,
+ ar: torch.nn.Module,
+ style_encoder: torch.nn.Module,
+ vocoder: torch.nn.Module,
+ ):
+ super(VoiceConversionWrapper, self).__init__()
+ self.sr = sr
+ self.hop_size = hop_size
+ self.mel_fn = mel_fn
+ self.cfm = cfm
+ self.cfm_length_regulator = cfm_length_regulator
+ self.content_extractor_narrow = content_extractor_narrow
+ self.content_extractor_wide = content_extractor_wide
+ self.vocoder = vocoder
+ self.ar_length_regulator = ar_length_regulator
+ self.ar = ar
+ self.style_encoder = style_encoder
+ # Set streaming parameters
+ self.overlap_frame_len = 16
+ self.bitrate = "320k"
+ self.compiled_decode_fn = None
+ self.dit_compiled = False
+ self.dit_max_context_len = 30 # in seconds
+ self.ar_max_content_len = 1500 # in num of narrow tokens
+ self.compile_len = 87 * self.dit_max_context_len
+
+ def forward_cfm(self, content_indices_wide, content_lens, mels, mel_lens, style_vectors):
+ device = content_indices_wide.device
+ B = content_indices_wide.size(0)
+ cond, _ = self.cfm_length_regulator(content_indices_wide, ylens=mel_lens)
+
+ # randomly set a length as prompt
+ prompt_len_max = mel_lens - 1
+ prompt_len = (torch.rand([B], device=device) * prompt_len_max).floor().to(dtype=torch.long)
+ prompt_len[torch.rand([B], device=device) < 0.1] = 0
+
+ loss = self.cfm(mels, mel_lens, prompt_len, cond, style_vectors)
+ return loss
+
+ def forward_ar(self, content_indices_narrow, content_indices_wide, content_lens):
+ device = content_indices_narrow.device
+ duration_reduced_narrow_tokens = []
+ duration_reduced_narrow_lens = []
+ for bib in range(content_indices_narrow.size(0)):
+ reduced, reduced_len = self.duration_reduction_func(content_indices_narrow[bib])
+ duration_reduced_narrow_tokens.append(reduced)
+ duration_reduced_narrow_lens.append(reduced_len)
+ duration_reduced_narrow_tokens = torch.nn.utils.rnn.pad_sequence(duration_reduced_narrow_tokens,
+ batch_first=True, padding_value=0).to(device)
+ duration_reduced_narrow_lens = torch.LongTensor(duration_reduced_narrow_lens).to(device)
+
+ # interpolate speech token to match acoustic feature length
+ cond, _ = self.ar_length_regulator(duration_reduced_narrow_tokens)
+ loss = self.ar(cond, duration_reduced_narrow_lens, content_indices_wide, content_lens)
+ return loss
+
+ def forward(self, waves_16k, mels, wave_lens_16k, mel_lens, forward_ar=False, forward_cfm=True):
+ """
+ Forward pass for the model.
+ """
+ # extract wide content features as both AR and CFM models use them
+ with torch.no_grad():
+ _, content_indices_wide, content_lens = self.content_extractor_wide(waves_16k, wave_lens_16k)
+ if forward_ar:
+ # extract narrow content features for AR model
+ _, content_indices_narrow, _ = self.content_extractor_narrow(waves_16k, wave_lens_16k, ssl_model=self.content_extractor_wide.ssl_model)
+ loss_ar = self.forward_ar(content_indices_narrow.clone(), content_indices_wide.clone(), content_lens)
+ else:
+ loss_ar = torch.tensor(0.0, device=waves_16k.device, dtype=waves_16k.dtype)
+ if forward_cfm:
+ style_vectors = self.compute_style(waves_16k, wave_lens_16k)
+ loss_cfm = self.forward_cfm(content_indices_wide, content_lens, mels, mel_lens, style_vectors)
+ else:
+ loss_cfm = torch.tensor(0.0, device=waves_16k.device, dtype=waves_16k.dtype)
+ return loss_ar, loss_cfm
+
+ def compile_ar(self):
+ """
+ Compile the AR model for inference.
+ """
+ self.compiled_decode_fn = torch.compile(
+ self.ar.model.forward_generate,
+ fullgraph=True,
+ backend="inductor" if torch.cuda.is_available() else "aot_eager",
+ mode="reduce-overhead" if torch.cuda.is_available() else None,
+ )
+
+ def compile_cfm(self):
+ self.cfm.estimator.transformer = torch.compile(
+ self.cfm.estimator.transformer,
+ fullgraph=True,
+ backend="inductor" if torch.cuda.is_available() else "aot_eager",
+ mode="reduce-overhead" if torch.cuda.is_available() else None,
+ )
+ self.dit_compiled = True
+
+ @staticmethod
+ def strip_prefix(state_dict: dict, prefix: str = "module.") -> dict:
+ """
+ Strip the prefix from the state_dict keys.
+ """
+ new_state_dict = {}
+ for k, v in state_dict.items():
+ if k.startswith(prefix):
+ new_key = k[len(prefix):]
+ else:
+ new_key = k
+ new_state_dict[new_key] = v
+ return new_state_dict
+
+ @staticmethod
+ def duration_reduction_func(token_seq, n_gram=1):
+ """
+ Args:
+ token_seq: (T,)
+ Returns:
+ reduced_token_seq: (T')
+ reduced_token_seq_len: T'
+ """
+ n_gram_seq = token_seq.unfold(0, n_gram, 1)
+ mask = torch.all(n_gram_seq[1:] != n_gram_seq[:-1], dim=1)
+ reduced_token_seq = torch.cat(
+ (n_gram_seq[0, :n_gram], n_gram_seq[1:, -1][mask])
+ )
+ return reduced_token_seq, len(reduced_token_seq)
+
+ @staticmethod
+ def crossfade(chunk1, chunk2, overlap):
+ """Apply crossfade between two audio chunks."""
+ 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
+
+ def _stream_wave_chunks(self, vc_wave, processed_frames, vc_mel, overlap_wave_len,
+ generated_wave_chunks, previous_chunk, is_last_chunk, stream_output):
+ """
+ Helper method to handle streaming wave chunks.
+
+ Args:
+ vc_wave: The current wave chunk
+ processed_frames: Number of frames processed so far
+ vc_mel: The mel spectrogram
+ overlap_wave_len: Length of overlap between chunks
+ generated_wave_chunks: List of generated wave chunks
+ previous_chunk: Previous wave chunk for crossfading
+ is_last_chunk: Whether this is the last chunk
+ stream_output: Whether to stream the output
+
+ Returns:
+ Tuple of (processed_frames, previous_chunk, should_break, mp3_bytes, full_audio)
+ where should_break indicates if processing should stop
+ mp3_bytes is the MP3 bytes if streaming, None otherwise
+ full_audio is the full audio if this is the last chunk, None otherwise
+ """
+ mp3_bytes = None
+ full_audio = None
+
+ if processed_frames == 0:
+ if is_last_chunk:
+ output_wave = vc_wave[0].cpu().numpy()
+ generated_wave_chunks.append(output_wave)
+
+ if stream_output:
+ output_wave_int16 = (output_wave * 32768.0).astype(np.int16)
+ mp3_bytes = AudioSegment(
+ output_wave_int16.tobytes(), frame_rate=self.sr,
+ sample_width=output_wave_int16.dtype.itemsize, channels=1
+ ).export(format="mp3", bitrate=self.bitrate).read()
+ full_audio = (self.sr, np.concatenate(generated_wave_chunks))
+ else:
+ return processed_frames, previous_chunk, True, None, np.concatenate(generated_wave_chunks)
+
+ return processed_frames, previous_chunk, True, mp3_bytes, full_audio
+
+ 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_mel.size(2) - self.overlap_frame_len
+
+ if stream_output:
+ output_wave_int16 = (output_wave * 32768.0).astype(np.int16)
+ mp3_bytes = AudioSegment(
+ output_wave_int16.tobytes(), frame_rate=self.sr,
+ sample_width=output_wave_int16.dtype.itemsize, channels=1
+ ).export(format="mp3", bitrate=self.bitrate).read()
+
+ elif is_last_chunk:
+ output_wave = self.crossfade(previous_chunk.cpu().numpy(), vc_wave[0].cpu().numpy(), overlap_wave_len)
+ generated_wave_chunks.append(output_wave)
+ processed_frames += vc_mel.size(2) - self.overlap_frame_len
+
+ if stream_output:
+ output_wave_int16 = (output_wave * 32768.0).astype(np.int16)
+ mp3_bytes = AudioSegment(
+ output_wave_int16.tobytes(), frame_rate=self.sr,
+ sample_width=output_wave_int16.dtype.itemsize, channels=1
+ ).export(format="mp3", bitrate=self.bitrate).read()
+ full_audio = (self.sr, np.concatenate(generated_wave_chunks))
+ else:
+ return processed_frames, previous_chunk, True, None, np.concatenate(generated_wave_chunks)
+
+ return processed_frames, previous_chunk, True, mp3_bytes, full_audio
+
+ else:
+ output_wave = self.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_mel.size(2) - self.overlap_frame_len
+
+ if stream_output:
+ output_wave_int16 = (output_wave * 32768.0).astype(np.int16)
+ mp3_bytes = AudioSegment(
+ output_wave_int16.tobytes(), frame_rate=self.sr,
+ sample_width=output_wave_int16.dtype.itemsize, channels=1
+ ).export(format="mp3", bitrate=self.bitrate).read()
+
+ return processed_frames, previous_chunk, False, mp3_bytes, full_audio
+
+ def load_checkpoints(
+ self,
+ cfm_checkpoint_path = None,
+ ar_checkpoint_path = None,
+ ):
+ if cfm_checkpoint_path is None:
+ cfm_checkpoint_path = load_custom_model_from_hf(
+ repo_id=DEFAULT_REPO_ID,
+ model_filename=DEFAULT_CFM_CHECKPOINT,
+ )
+ else:
+ print(f"Loading CFM checkpoint from {cfm_checkpoint_path}...")
+ if ar_checkpoint_path is None:
+ ar_checkpoint_path = load_custom_model_from_hf(
+ repo_id=DEFAULT_REPO_ID,
+ model_filename=DEFAULT_AR_CHECKPOINT,
+ )
+ else:
+ print(f"Loading AR checkpoint from {ar_checkpoint_path}...")
+ # cfm
+ cfm_checkpoint = torch.load(cfm_checkpoint_path, map_location="cpu")
+ cfm_length_regulator_state_dict = self.strip_prefix(cfm_checkpoint["net"]['length_regulator'], "module.")
+ cfm_state_dict = self.strip_prefix(cfm_checkpoint["net"]['cfm'], "module.")
+ missing_keys, unexpected_keys = self.cfm.load_state_dict(cfm_state_dict, strict=False)
+ missing_keys, unexpected_keys = self.cfm_length_regulator.load_state_dict(cfm_length_regulator_state_dict, strict=False)
+
+ # ar
+ ar_checkpoint = torch.load(ar_checkpoint_path, map_location="cpu")
+ ar_length_regulator_state_dict = self.strip_prefix(ar_checkpoint["net"]['length_regulator'], "module.")
+ ar_state_dict = self.strip_prefix(ar_checkpoint["net"]['ar'], "module.")
+ missing_keys, unexpected_keys = self.ar.load_state_dict(ar_state_dict, strict=False)
+ missing_keys, unexpected_keys = self.ar_length_regulator.load_state_dict(ar_length_regulator_state_dict, strict=False)
+
+ # content extractor
+ content_extractor_narrow_checkpoint_path = load_custom_model_from_hf(
+ repo_id=DEFAULT_CE_REPO_ID,
+ model_filename=DEFAULT_CE_NARROW_CHECKPOINT,
+ )
+ content_extractor_narrow_checkpoint = torch.load(content_extractor_narrow_checkpoint_path, map_location="cpu")
+ self.content_extractor_narrow.load_state_dict(
+ content_extractor_narrow_checkpoint, strict=False
+ )
+
+ content_extractor_wide_checkpoint_path = load_custom_model_from_hf(
+ repo_id=DEFAULT_CE_REPO_ID,
+ model_filename=DEFAULT_CE_WIDE_CHECKPOINT,
+ )
+ content_extractor_wide_checkpoint = torch.load(content_extractor_wide_checkpoint_path, map_location="cpu")
+ self.content_extractor_wide.load_state_dict(
+ content_extractor_wide_checkpoint, strict=False
+ )
+
+ # style encoder
+ style_encoder_checkpoint_path = load_custom_model_from_hf(DEFAULT_SE_REPO_ID, DEFAULT_SE_CHECKPOINT, config_filename=None)
+ style_encoder_checkpoint = torch.load(style_encoder_checkpoint_path, map_location="cpu")
+ self.style_encoder.load_state_dict(style_encoder_checkpoint, strict=False)
+
+ def setup_ar_caches(self, max_batch_size=1, max_seq_len=4096, dtype=torch.float32, device=torch.device("cpu")):
+ self.ar.setup_caches(max_batch_size=max_batch_size, max_seq_len=max_seq_len, dtype=dtype, device=device)
+
+ @torch.no_grad()
+ def compute_style(self, waves_16k: torch.Tensor, wave_lens_16k: torch.Tensor = None):
+ if wave_lens_16k is None:
+ wave_lens_16k = torch.tensor([waves_16k.size(-1)], dtype=torch.int32).to(waves_16k.device)
+ feat_list = []
+ for bib in range(waves_16k.size(0)):
+ feat = torchaudio.compliance.kaldi.fbank(waves_16k[bib:bib + 1, :wave_lens_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(waves_16k.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
+ ]
+ feat = torch.stack(feat_list, dim=0)
+ style = self.style_encoder(feat, feat_lens)
+ return style
+
+ @torch.no_grad()
+ @torch.inference_mode()
+ def convert_timbre(
+ self,
+ source_audio_path: str,
+ target_audio_path: str,
+ diffusion_steps: int = 30,
+ length_adjust: float = 1.0,
+ inference_cfg_rate: float = 0.5,
+ use_sway_sampling: bool = False,
+ use_amo_sampling: bool = False,
+ device: torch.device = torch.device("cpu"),
+ dtype: torch.dtype = torch.float32,
+ ):
+ source_wave = librosa.load(source_audio_path, sr=self.sr)[0]
+ target_wave = librosa.load(target_audio_path, sr=self.sr)[0]
+ source_wave_tensor = torch.tensor(source_wave).unsqueeze(0).to(device)
+ target_wave_tensor = torch.tensor(target_wave).unsqueeze(0).to(device)
+
+ # get 16khz audio
+ source_wave_16k = librosa.resample(source_wave, orig_sr=self.sr, target_sr=16000)
+ target_wave_16k = librosa.resample(target_wave, orig_sr=self.sr, target_sr=16000)
+ source_wave_16k_tensor = torch.tensor(source_wave_16k).unsqueeze(0).to(device)
+ target_wave_16k_tensor = torch.tensor(target_wave_16k).unsqueeze(0).to(device)
+
+ # compute mel spectrogram
+ source_mel = self.mel_fn(source_wave_tensor)
+ target_mel = self.mel_fn(target_wave_tensor)
+ source_mel_len = source_mel.size(2)
+ target_mel_len = target_mel.size(2)
+
+ with torch.autocast(device_type=device.type, dtype=dtype):
+ # compute content features
+ _, source_content_indices, _ = self.content_extractor_wide(source_wave_16k_tensor, [source_wave_16k.size])
+ _, target_content_indices, _ = self.content_extractor_wide(target_wave_16k_tensor, [target_wave_16k.size])
+
+ # compute style features
+ target_style = self.compute_style(target_wave_16k_tensor)
+
+ # Length regulation
+ cond, _ = self.cfm_length_regulator(source_content_indices, ylens=torch.LongTensor([source_mel_len]).to(device))
+ prompt_condition, _, = self.cfm_length_regulator(target_content_indices, ylens=torch.LongTensor([target_mel_len]).to(device))
+
+ cat_condition = torch.cat([prompt_condition, cond], dim=1)
+ # generate mel spectrogram
+ vc_mel = self.cfm.inference(
+ cat_condition,
+ torch.LongTensor([cat_condition.size(1)]).to(device),
+ target_mel, target_style, diffusion_steps,
+ inference_cfg_rate=inference_cfg_rate,
+ sway_sampling=use_sway_sampling,
+ amo_sampling=use_amo_sampling,
+ )
+ vc_mel = vc_mel[:, :, target_mel_len:]
+ vc_wave = self.vocoder(vc_mel.float()).squeeze()[None]
+ return vc_wave.cpu().numpy()
+
+ @torch.no_grad()
+ @torch.inference_mode()
+ def convert_voice(
+ self,
+ source_audio_path: str,
+ target_audio_path: str,
+ diffusion_steps: int = 30,
+ length_adjust: float = 1.0,
+ inference_cfg_rate: float = 0.5,
+ top_p: float = 0.7,
+ temperature: float = 0.7,
+ repetition_penalty: float = 1.5,
+ use_sway_sampling: bool = False,
+ use_amo_sampling: bool = False,
+ device: torch.device = torch.device("cpu"),
+ dtype: torch.dtype = torch.float32,
+ ):
+ source_wave = librosa.load(source_audio_path, sr=self.sr)[0]
+ target_wave = librosa.load(target_audio_path, sr=self.sr)[0]
+ source_wave_tensor = torch.tensor(source_wave).unsqueeze(0).to(device)
+ target_wave_tensor = torch.tensor(target_wave).unsqueeze(0).to(device)
+
+ # get 16khz audio
+ source_wave_16k = librosa.resample(source_wave, orig_sr=self.sr, target_sr=16000)
+ target_wave_16k = librosa.resample(target_wave, orig_sr=self.sr, target_sr=16000)
+ source_wave_16k_tensor = torch.tensor(source_wave_16k).unsqueeze(0).to(device)
+ target_wave_16k_tensor = torch.tensor(target_wave_16k).unsqueeze(0).to(device)
+
+ # compute mel spectrogram
+ source_mel = self.mel_fn(source_wave_tensor)
+ target_mel = self.mel_fn(target_wave_tensor)
+ source_mel_len = source_mel.size(2)
+ target_mel_len = target_mel.size(2)
+
+ with torch.autocast(device_type=device.type, dtype=dtype):
+ # compute content features
+ _, source_content_indices, _ = self.content_extractor_wide(source_wave_16k_tensor, [source_wave_16k.size])
+ _, target_content_indices, _ = self.content_extractor_wide(target_wave_16k_tensor, [target_wave_16k.size])
+
+ _, source_narrow_indices, _ = self.content_extractor_narrow(source_wave_16k_tensor,
+ [source_wave_16k.size], ssl_model=self.content_extractor_wide.ssl_model)
+ _, target_narrow_indices, _ = self.content_extractor_narrow(target_wave_16k_tensor,
+ [target_wave_16k.size], ssl_model=self.content_extractor_wide.ssl_model)
+
+ src_narrow_reduced, src_narrow_len = self.duration_reduction_func(source_narrow_indices[0], 1)
+ tgt_narrow_reduced, tgt_narrow_len = self.duration_reduction_func(target_narrow_indices[0], 1)
+
+ ar_cond = self.ar_length_regulator(torch.cat([tgt_narrow_reduced, src_narrow_reduced], dim=0)[None])[0]
+
+ ar_out = self.ar.generate(ar_cond, target_content_indices, top_p=top_p, temperature=temperature, repetition_penalty=repetition_penalty)
+ ar_out_mel_len = torch.LongTensor([int(source_mel_len / source_content_indices.size(-1) * ar_out.size(-1) * length_adjust)]).to(device)
+ # compute style features
+ target_style = self.compute_style(target_wave_16k_tensor)
+
+ # Length regulation
+ cond, _ = self.cfm_length_regulator(ar_out, ylens=torch.LongTensor([ar_out_mel_len]).to(device))
+ prompt_condition, _, = self.cfm_length_regulator(target_content_indices, ylens=torch.LongTensor([target_mel_len]).to(device))
+
+ cat_condition = torch.cat([prompt_condition, cond], dim=1)
+ # generate mel spectrogram
+ vc_mel = self.cfm.inference(
+ cat_condition,
+ torch.LongTensor([cat_condition.size(1)]).to(device),
+ target_mel, target_style, diffusion_steps,
+ inference_cfg_rate=inference_cfg_rate,
+ sway_sampling=use_sway_sampling,
+ amo_sampling=use_amo_sampling,
+ )
+ vc_mel = vc_mel[:, :, target_mel_len:]
+ vc_wave = self.vocoder(vc_mel.float()).squeeze()[None]
+ return vc_wave.cpu().numpy()
+
+ def _process_content_features(self, audio_16k_tensor, is_narrow=False):
+ """Process audio through Whisper model to extract features."""
+ content_extractor_fn = self.content_extractor_narrow if is_narrow else self.content_extractor_wide
+ if audio_16k_tensor.size(-1) <= 16000 * 30:
+ # Compute content features
+ _, content_indices, _ = content_extractor_fn(audio_16k_tensor, [audio_16k_tensor.size(-1)], ssl_model=self.content_extractor_wide.ssl_model)
+ else:
+ # Process long audio in chunks
+ overlapping_time = 5 # 5 seconds
+ features_list = []
+ buffer = None
+ traversed_time = 0
+ while traversed_time < audio_16k_tensor.size(-1):
+ if buffer is None: # first chunk
+ chunk = audio_16k_tensor[:, traversed_time:traversed_time + 16000 * 30]
+ else:
+ chunk = torch.cat([
+ buffer,
+ audio_16k_tensor[:, traversed_time:traversed_time + 16000 * (30 - overlapping_time)]
+ ], dim=-1)
+ _, chunk_content_indices, _ = content_extractor_fn(chunk, [chunk.size(-1)], ssl_model=self.content_extractor_wide.ssl_model)
+ if traversed_time == 0:
+ features_list.append(chunk_content_indices)
+ else:
+ features_list.append(chunk_content_indices[:, 50 * overlapping_time:])
+ buffer = chunk[:, -16000 * overlapping_time:]
+ traversed_time += 30 * 16000 if traversed_time == 0 else chunk.size(-1) - 16000 * overlapping_time
+ content_indices = torch.cat(features_list, dim=1)
+
+ return content_indices
+
+ @torch.no_grad()
+ @torch.inference_mode()
+ def convert_voice_with_streaming(
+ self,
+ source_audio_path: str,
+ target_audio_path: str,
+ diffusion_steps: int = 30,
+ length_adjust: float = 1.0,
+ intelligebility_cfg_rate: float = 0.7,
+ similarity_cfg_rate: float = 0.7,
+ top_p: float = 0.7,
+ temperature: float = 0.7,
+ repetition_penalty: float = 1.5,
+ convert_style: bool = False,
+ anonymization_only: bool = False,
+ device: torch.device = torch.device("cuda"),
+ dtype: torch.dtype = torch.float16,
+ stream_output: bool = True,
+ ):
+ """
+ Convert voice with streaming support for long audio files.
+
+ Args:
+ source_audio_path: Path to source audio file
+ target_audio_path: Path to target audio file
+ diffusion_steps: Number of diffusion steps (default: 30)
+ length_adjust: Length adjustment factor (default: 1.0)
+ intelligebility_cfg_rate: CFG rate for intelligibility (default: 0.7)
+ similarity_cfg_rate: CFG rate for similarity (default: 0.7)
+ top_p: Top-p sampling parameter (default: 0.7)
+ temperature: Temperature for sampling (default: 0.7)
+ repetition_penalty: Repetition penalty (default: 1.5)
+ device: Device to use (default: cpu)
+ dtype: Data type to use (default: float32)
+ stream_output: Whether to stream the output (default: True)
+
+ Returns:
+ If stream_output is True, yields (mp3_bytes, full_audio) tuples
+ If stream_output is False, returns the full audio as a numpy array
+ """
+ # Load audio
+ source_wave = librosa.load(source_audio_path, sr=self.sr)[0]
+ target_wave = librosa.load(target_audio_path, sr=self.sr)[0]
+
+ # Limit target audio to 25 seconds
+ target_wave = target_wave[:self.sr * (self.dit_max_context_len - 5)]
+
+ source_wave_tensor = torch.tensor(source_wave).unsqueeze(0).float().to(device)
+ target_wave_tensor = torch.tensor(target_wave).unsqueeze(0).float().to(device)
+
+ # Resample to 16kHz for feature extraction
+ source_wave_16k = librosa.resample(source_wave, orig_sr=self.sr, target_sr=16000)
+ target_wave_16k = librosa.resample(target_wave, orig_sr=self.sr, target_sr=16000)
+ source_wave_16k_tensor = torch.tensor(source_wave_16k).unsqueeze(0).to(device)
+ target_wave_16k_tensor = torch.tensor(target_wave_16k).unsqueeze(0).to(device)
+
+ # Compute mel spectrograms
+ source_mel = self.mel_fn(source_wave_tensor)
+ target_mel = self.mel_fn(target_wave_tensor)
+ source_mel_len = source_mel.size(2)
+ target_mel_len = target_mel.size(2)
+
+ # Set up chunk processing parameters
+ max_context_window = self.sr // self.hop_size * self.dit_max_context_len
+ overlap_wave_len = self.overlap_frame_len * self.hop_size
+
+ with torch.autocast(device_type=device.type, dtype=dtype):
+ # Compute content features
+ source_content_indices = self._process_content_features(source_wave_16k_tensor, is_narrow=False)
+ target_content_indices = self._process_content_features(target_wave_16k_tensor, is_narrow=False)
+ # Compute style features
+ target_style = self.compute_style(target_wave_16k_tensor)
+ prompt_condition, _, = self.cfm_length_regulator(target_content_indices,
+ ylens=torch.LongTensor([target_mel_len]).to(device))
+
+ # prepare for streaming
+ generated_wave_chunks = []
+ processed_frames = 0
+ previous_chunk = None
+ if convert_style:
+ with torch.autocast(device_type=device.type, dtype=dtype):
+ source_narrow_indices = self._process_content_features(source_wave_16k_tensor, is_narrow=True)
+ target_narrow_indices = self._process_content_features(target_wave_16k_tensor, is_narrow=True)
+ src_narrow_reduced, src_narrow_len = self.duration_reduction_func(source_narrow_indices[0], 1)
+ tgt_narrow_reduced, tgt_narrow_len = self.duration_reduction_func(target_narrow_indices[0], 1)
+ # Process src_narrow_reduced in chunks of max 1000 tokens
+ max_chunk_size = self.ar_max_content_len - tgt_narrow_len
+
+ # Process src_narrow_reduced in chunks
+ for i in range(0, len(src_narrow_reduced), max_chunk_size):
+ is_last_chunk = i + max_chunk_size >= len(src_narrow_reduced)
+ with torch.autocast(device_type=device.type, dtype=dtype):
+ chunk = src_narrow_reduced[i:i + max_chunk_size]
+ if anonymization_only:
+ chunk_ar_cond = self.ar_length_regulator(chunk[None])[0]
+ chunk_ar_out = self.ar.generate(chunk_ar_cond, torch.zeros([1, 0]).long().to(device),
+ compiled_decode_fn=self.compiled_decode_fn,
+ top_p=top_p, temperature=temperature,
+ repetition_penalty=repetition_penalty)
+ else:
+ # For each chunk, we need to include tgt_narrow_reduced as context
+ chunk_ar_cond = self.ar_length_regulator(torch.cat([tgt_narrow_reduced, chunk], dim=0)[None])[0]
+ chunk_ar_out = self.ar.generate(chunk_ar_cond, target_content_indices, compiled_decode_fn=self.compiled_decode_fn,
+ top_p=top_p, temperature=temperature,
+ repetition_penalty=repetition_penalty)
+ chunkar_out_mel_len = torch.LongTensor([int(source_mel_len / source_content_indices.size(
+ -1) * chunk_ar_out.size(-1) * length_adjust)]).to(device)
+ # Length regulation
+ chunk_cond, _ = self.cfm_length_regulator(chunk_ar_out, ylens=torch.LongTensor([chunkar_out_mel_len]).to(device))
+ cat_condition = torch.cat([prompt_condition, chunk_cond], dim=1)
+ original_len = cat_condition.size(1)
+ # pad cat_condition to compile_len
+ if self.dit_compiled:
+ cat_condition = torch.nn.functional.pad(cat_condition,
+ (0, 0, 0, self.compile_len - cat_condition.size(1),),
+ value=0)
+ # Voice Conversion
+ vc_mel = self.cfm.inference(
+ cat_condition,
+ torch.LongTensor([original_len]).to(device),
+ target_mel, target_style, diffusion_steps,
+ inference_cfg_rate=[intelligebility_cfg_rate, similarity_cfg_rate],
+ random_voice=anonymization_only,
+ )
+ vc_mel = vc_mel[:, :, target_mel_len:original_len]
+ vc_wave = self.vocoder(vc_mel).squeeze()[None]
+ processed_frames, previous_chunk, should_break, mp3_bytes, full_audio = self._stream_wave_chunks(
+ vc_wave, processed_frames, vc_mel, overlap_wave_len,
+ generated_wave_chunks, previous_chunk, is_last_chunk, stream_output
+ )
+
+ if stream_output and mp3_bytes is not None:
+ yield mp3_bytes, full_audio
+ if should_break:
+ break
+ else:
+ cond, _ = self.cfm_length_regulator(source_content_indices, ylens=torch.LongTensor([source_mel_len]).to(device))
+
+ # Process in chunks for streaming
+ max_source_window = max_context_window - target_mel.size(2)
+
+ # 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)
+ original_len = cat_condition.size(1)
+ # pad cat_condition to compile_len
+ if self.dit_compiled:
+ cat_condition = torch.nn.functional.pad(cat_condition,
+ (0, 0, 0, self.compile_len - cat_condition.size(1),), value=0)
+ with torch.autocast(device_type=device.type, dtype=torch.float32): # force CFM to use float32
+ # Voice Conversion
+ vc_mel = self.cfm.inference(
+ cat_condition,
+ torch.LongTensor([original_len]).to(device),
+ target_mel, target_style, diffusion_steps,
+ inference_cfg_rate=[intelligebility_cfg_rate, similarity_cfg_rate],
+ random_voice=anonymization_only,
+ )
+ vc_mel = vc_mel[:, :, target_mel_len:original_len]
+ vc_wave = self.vocoder(vc_mel).squeeze()[None]
+
+ processed_frames, previous_chunk, should_break, mp3_bytes, full_audio = self._stream_wave_chunks(
+ vc_wave, processed_frames, vc_mel, overlap_wave_len,
+ generated_wave_chunks, previous_chunk, is_last_chunk, stream_output
+ )
+
+ if stream_output and mp3_bytes is not None:
+ yield mp3_bytes, full_audio
+ if should_break:
+ break
diff --git a/seed-vc/modules/wavenet.py b/seed-vc/modules/wavenet.py
new file mode 100644
index 0000000000000000000000000000000000000000..2d2c3b862593e2c0f78b2f8cb2d38c4cbbe573eb
--- /dev/null
+++ b/seed-vc/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/seed-vc/optimizers.py b/seed-vc/optimizers.py
new file mode 100644
index 0000000000000000000000000000000000000000..d97b8a7bfbb6f6f226db8bae2e474b44e24e0d1a
--- /dev/null
+++ b/seed-vc/optimizers.py
@@ -0,0 +1,120 @@
+#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-6,
+ weight_decay=0.01,
+ )
+ 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
+
+class MinLRExponentialLR(torch.optim.lr_scheduler.ExponentialLR):
+ def __init__(self, optimizer, gamma, min_lr=1e-5):
+ self.min_lr = min_lr
+ super().__init__(optimizer, gamma)
+
+ def get_lr(self):
+ lrs = super().get_lr()
+ return [max(lr, self.min_lr) for lr in lrs]
+
+def build_single_optimizer(model, lr,):
+ model_parameters = model.parameters()
+ parameters_require_grad = filter(lambda p: p.requires_grad, model_parameters)
+ optim = AdamW(
+ parameters_require_grad,
+ lr=lr,
+ betas=(0.9, 0.98),
+ eps=1e-6,
+ weight_decay=0.01,
+ )
+
+ scheduler = MinLRExponentialLR(optim, gamma=0.999996, min_lr=1e-5)
+
+ return optim, scheduler
\ No newline at end of file
diff --git a/seed-vc/real-time-gui.py b/seed-vc/real-time-gui.py
new file mode 100644
index 0000000000000000000000000000000000000000..38ad345f08d4aabe174c2239c71302c3dca4d45f
--- /dev/null
+++ b/seed-vc/real-time-gui.py
@@ -0,0 +1,1190 @@
+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 = None
+
+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
+ if device.type == "mps":
+ start_event = torch.mps.event.Event(enable_timing=True)
+ end_event = torch.mps.event.Event(enable_timing=True)
+ torch.mps.synchronize()
+ else:
+ 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()
+ if device.type == "mps":
+ torch.mps.synchronize() # MPS - Wait for the events to be recorded!
+ else:
+ 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 = device
+
+
+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 Files", "*.wav"),
+ ("MP3 Files", "*.mp3"),
+ ("FLAC Files", "*.flac"),
+ ("M4A Files", "*.m4a"),
+ ("OGG Files", "*.ogg"),
+ ("Opus Files", "*.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):
+ if device.type == "mps":
+ torch.mps.empty_cache()
+ else:
+ 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 = min(500, 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
+ if device.type == "mps":
+ start_event = torch.mps.event.Event(enable_timing=True)
+ end_event = torch.mps.event.Event(enable_timing=True)
+ torch.mps.synchronize()
+ else:
+ 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()
+ if device.type == "mps":
+ torch.mps.synchronize() # MPS - Wait for the events to be recorded!
+ else:
+ 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:
+ # ]
+ torch.from_numpy(librosa.resample(self.input_wav[-indata.shape[0] - 2 * self.zc :].cpu().numpy(), orig_sr=self.gui_config.samplerate, target_sr=16000)[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!")
+ if device.type == "mps":
+ start_event = torch.mps.event.Event(enable_timing=True)
+ end_event = torch.mps.event.Event(enable_timing=True)
+ torch.mps.synchronize()
+ else:
+ 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()
+ if device.type == "mps":
+ torch.mps.synchronize() # MPS - Wait for the events to be recorded!
+ else:
+ 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
+ )
+
+ tensor = cor_nom[0, 0] / cor_den[0, 0]
+ if tensor.numel() > 1: # If tensor has multiple elements
+ if sys.platform == "darwin":
+ _, sola_offset = torch.max(tensor, dim=0)
+ sola_offset = sola_offset.item()
+ else:
+ sola_offset = torch.argmax(tensor, dim=0).item()
+ else:
+ sola_offset = tensor.item()
+
+ 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)
+ parser.add_argument("--gpu", type=int, help="Which GPU id to use", default=0)
+ args = parser.parse_args()
+ cuda_target = f"cuda:{args.gpu}" if args.gpu else "cuda"
+
+ if torch.cuda.is_available():
+ device = torch.device(cuda_target)
+ elif torch.backends.mps.is_available():
+ device = torch.device("mps")
+ else:
+ device = torch.device("cpu")
+ gui = GUI(args)
\ No newline at end of file
diff --git a/seed-vc/requirements-mac.txt b/seed-vc/requirements-mac.txt
new file mode 100644
index 0000000000000000000000000000000000000000..1d841d79c0b9eacf0386d75cbe9cc8a843db1f35
--- /dev/null
+++ b/seed-vc/requirements-mac.txt
@@ -0,0 +1,25 @@
+--extra-index-url https://download.pytorch.org/whl/cu121
+torch --pre --extra-index-url https://download.pytorch.org/whl/nightly/cpu
+torchvision --pre --extra-index-url https://download.pytorch.org/whl/nightly/cpu
+torchaudio --pre --extra-index-url https://download.pytorch.org/whl/nightly/cpu
+accelerate
+scipy==1.13.1
+librosa==0.10.2
+huggingface-hub>=0.28.1
+munch==4.0.0
+einops==0.8.0
+descript-audio-codec==1.0.0
+gradio==5.23.0
+pydub==0.25.1
+resemblyzer
+jiwer==3.0.3
+transformers==4.46.3
+FreeSimpleGUI==5.1.1
+soundfile==0.12.1
+sounddevice==0.5.0
+modelscope==1.18.1
+funasr==1.1.5
+numpy==1.26.4
+pyyaml
+python-dotenv
+hydra-core==1.3.2
\ No newline at end of file
diff --git a/seed-vc/requirements.txt b/seed-vc/requirements.txt
new file mode 100644
index 0000000000000000000000000000000000000000..399c6649b58b757c3096351ef79a0cfa218390e8
--- /dev/null
+++ b/seed-vc/requirements.txt
@@ -0,0 +1,27 @@
+torch --pre --index-url https://download.pytorch.org/whl/nightly/cu126
+torchvision --pre --index-url https://download.pytorch.org/whl/nightly/cu126
+torchaudio --pre --index-url https://download.pytorch.org/whl/nightly/cu126
+accelerate
+torch==2.4.0
+torchvision==0.19.0
+torchaudio==2.4.0
+scipy==1.13.1
+librosa==0.10.2
+huggingface-hub>=0.28.1
+munch==4.0.0
+einops==0.8.0
+descript-audio-codec==1.0.0
+gradio==5.23.0
+pydub==0.25.1
+resemblyzer
+jiwer==3.0.3
+transformers==4.46.3
+FreeSimpleGUI==5.1.1
+soundfile==0.12.1
+sounddevice==0.5.0
+modelscope==1.18.1
+funasr==1.1.5
+numpy==1.26.4
+hydra-core==1.3.2
+pyyaml
+python-dotenv
\ No newline at end of file
diff --git a/seed-vc/seed_vc_wrapper.py b/seed-vc/seed_vc_wrapper.py
new file mode 100644
index 0000000000000000000000000000000000000000..c40d12094dacb210d8ce4875d8eb1affeb069b04
--- /dev/null
+++ b/seed-vc/seed_vc_wrapper.py
@@ -0,0 +1,461 @@
+import torch
+import torchaudio
+import librosa
+import numpy as np
+from pydub import AudioSegment
+import yaml
+from modules.commons import build_model, load_checkpoint, recursive_munch
+from hf_utils import load_custom_model_from_hf
+from modules.campplus.DTDNN import CAMPPlus
+from modules.bigvgan import bigvgan
+from modules.audio import mel_spectrogram
+from modules.rmvpe import RMVPE
+from transformers import AutoFeatureExtractor, WhisperModel
+
+class SeedVCWrapper:
+ def __init__(self, device=None):
+ """
+ Initialize the Seed-VC wrapper with all necessary models and configurations.
+
+ Args:
+ device: torch device to use. If None, will be automatically determined.
+ """
+ # Set device
+ if device is None:
+ if torch.cuda.is_available():
+ self.device = torch.device("cuda")
+ elif torch.backends.mps.is_available():
+ self.device = torch.device("mps")
+ else:
+ self.device = torch.device("cpu")
+ else:
+ self.device = device
+
+ # Load base model and configuration
+ self._load_base_model()
+
+ # Load F0 conditioned model
+ self._load_f0_model()
+
+ # Load additional modules
+ self._load_additional_modules()
+
+ # Set streaming parameters
+ self.overlap_frame_len = 16
+ self.bitrate = "320k"
+
+ def _load_base_model(self):
+ """Load the base DiT model for voice conversion."""
+ 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'])
+ self.model = build_model(model_params, stage='DiT')
+ self.hop_length = config['preprocess_params']['spect_params']['hop_length']
+ self.sr = config['preprocess_params']['sr']
+
+ # Load checkpoints
+ self.model, _, _, _ = load_checkpoint(
+ self.model, None, dit_checkpoint_path,
+ load_only_params=True, ignore_modules=[], is_distributed=False
+ )
+ for key in self.model:
+ self.model[key].eval()
+ self.model[key].to(self.device)
+ self.model.cfm.estimator.setup_caches(max_batch_size=1, max_seq_length=8192)
+
+ # Set up mel spectrogram function
+ 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": self.sr,
+ "fmin": 0,
+ "fmax": None,
+ "center": False
+ }
+ self.to_mel = lambda x: mel_spectrogram(x, **mel_fn_args)
+
+ # Load whisper model
+ whisper_name = model_params.speech_tokenizer.whisper_name if hasattr(model_params.speech_tokenizer, 'whisper_name') else "openai/whisper-small"
+ self.whisper_model = WhisperModel.from_pretrained(whisper_name, torch_dtype=torch.float16).to(self.device)
+ del self.whisper_model.decoder
+ self.whisper_feature_extractor = AutoFeatureExtractor.from_pretrained(whisper_name)
+
+ def _load_f0_model(self):
+ """Load the F0 conditioned model for voice conversion."""
+ 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.pth",
+ "config_dit_mel_seed_uvit_whisper_base_f0_44k.yml"
+ )
+ config = yaml.safe_load(open(dit_config_path, 'r'))
+ model_params = recursive_munch(config['model_params'])
+ self.model_f0 = build_model(model_params, stage='DiT')
+ self.hop_length_f0 = config['preprocess_params']['spect_params']['hop_length']
+ self.sr_f0 = config['preprocess_params']['sr']
+
+ # Load checkpoints
+ self.model_f0, _, _, _ = load_checkpoint(
+ self.model_f0, None, dit_checkpoint_path,
+ load_only_params=True, ignore_modules=[], is_distributed=False
+ )
+ for key in self.model_f0:
+ self.model_f0[key].eval()
+ self.model_f0[key].to(self.device)
+ self.model_f0.cfm.estimator.setup_caches(max_batch_size=1, max_seq_length=8192)
+
+ # Set up mel spectrogram function for F0 model
+ mel_fn_args_f0 = {
+ "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": self.sr_f0,
+ "fmin": 0,
+ "fmax": None,
+ "center": False
+ }
+ self.to_mel_f0 = lambda x: mel_spectrogram(x, **mel_fn_args_f0)
+
+ def _load_additional_modules(self):
+ """Load additional modules like CAMPPlus, BigVGAN, and RMVPE."""
+ # Load CAMPPlus
+ campplus_ckpt_path = load_custom_model_from_hf("funasr/campplus", "campplus_cn_common.bin", config_filename=None)
+ self.campplus_model = CAMPPlus(feat_dim=80, embedding_size=192)
+ self.campplus_model.load_state_dict(torch.load(campplus_ckpt_path, map_location="cpu"))
+ self.campplus_model.eval()
+ self.campplus_model.to(self.device)
+
+ # Load BigVGAN models
+ self.bigvgan_model = bigvgan.BigVGAN.from_pretrained('nvidia/bigvgan_v2_22khz_80band_256x', use_cuda_kernel=False)
+ self.bigvgan_model.remove_weight_norm()
+ self.bigvgan_model = self.bigvgan_model.eval().to(self.device)
+
+ self.bigvgan_44k_model = bigvgan.BigVGAN.from_pretrained('nvidia/bigvgan_v2_44khz_128band_512x', use_cuda_kernel=False)
+ self.bigvgan_44k_model.remove_weight_norm()
+ self.bigvgan_44k_model = self.bigvgan_44k_model.eval().to(self.device)
+
+ # Load RMVPE for F0 extraction
+ model_path = load_custom_model_from_hf("lj1995/VoiceConversionWebUI", "rmvpe.pt", None)
+ self.rmvpe = RMVPE(model_path, is_half=False, device=self.device)
+
+ @staticmethod
+ def adjust_f0_semitones(f0_sequence, n_semitones):
+ """Adjust F0 values by a number of semitones."""
+ factor = 2 ** (n_semitones / 12)
+ return f0_sequence * factor
+
+ @staticmethod
+ def crossfade(chunk1, chunk2, overlap):
+ """Apply crossfade between two audio chunks."""
+ 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
+
+ def _stream_wave_chunks(self, vc_wave, processed_frames, vc_target, overlap_wave_len,
+ generated_wave_chunks, previous_chunk, is_last_chunk, stream_output, sr):
+ """
+ Helper method to handle streaming wave chunks.
+
+ Args:
+ vc_wave: The current wave chunk
+ processed_frames: Number of frames processed so far
+ vc_target: The target mel spectrogram
+ overlap_wave_len: Length of overlap between chunks
+ generated_wave_chunks: List of generated wave chunks
+ previous_chunk: Previous wave chunk for crossfading
+ is_last_chunk: Whether this is the last chunk
+ stream_output: Whether to stream the output
+ sr: Sample rate
+
+ Returns:
+ Tuple of (processed_frames, previous_chunk, should_break, mp3_bytes, full_audio)
+ where should_break indicates if processing should stop
+ mp3_bytes is the MP3 bytes if streaming, None otherwise
+ full_audio is the full audio if this is the last chunk, None otherwise
+ """
+ mp3_bytes = None
+ full_audio = None
+
+ if processed_frames == 0:
+ if is_last_chunk:
+ output_wave = vc_wave[0].cpu().numpy()
+ generated_wave_chunks.append(output_wave)
+
+ if stream_output:
+ output_wave_int16 = (output_wave * 32768.0).astype(np.int16)
+ mp3_bytes = AudioSegment(
+ output_wave_int16.tobytes(), frame_rate=sr,
+ sample_width=output_wave_int16.dtype.itemsize, channels=1
+ ).export(format="mp3", bitrate=self.bitrate).read()
+ full_audio = (sr, np.concatenate(generated_wave_chunks))
+ else:
+ return processed_frames, previous_chunk, True, None, np.concatenate(generated_wave_chunks)
+
+ return processed_frames, previous_chunk, True, mp3_bytes, full_audio
+
+ 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) - self.overlap_frame_len
+
+ if stream_output:
+ output_wave_int16 = (output_wave * 32768.0).astype(np.int16)
+ mp3_bytes = AudioSegment(
+ output_wave_int16.tobytes(), frame_rate=sr,
+ sample_width=output_wave_int16.dtype.itemsize, channels=1
+ ).export(format="mp3", bitrate=self.bitrate).read()
+
+ elif is_last_chunk:
+ output_wave = self.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) - self.overlap_frame_len
+
+ if stream_output:
+ output_wave_int16 = (output_wave * 32768.0).astype(np.int16)
+ mp3_bytes = AudioSegment(
+ output_wave_int16.tobytes(), frame_rate=sr,
+ sample_width=output_wave_int16.dtype.itemsize, channels=1
+ ).export(format="mp3", bitrate=self.bitrate).read()
+ full_audio = (sr, np.concatenate(generated_wave_chunks))
+ else:
+ return processed_frames, previous_chunk, True, None, np.concatenate(generated_wave_chunks)
+
+ return processed_frames, previous_chunk, True, mp3_bytes, full_audio
+
+ else:
+ output_wave = self.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) - self.overlap_frame_len
+
+ if stream_output:
+ output_wave_int16 = (output_wave * 32768.0).astype(np.int16)
+ mp3_bytes = AudioSegment(
+ output_wave_int16.tobytes(), frame_rate=sr,
+ sample_width=output_wave_int16.dtype.itemsize, channels=1
+ ).export(format="mp3", bitrate=self.bitrate).read()
+
+ return processed_frames, previous_chunk, False, mp3_bytes, full_audio
+
+ def _process_whisper_features(self, audio_16k, is_source=True):
+ """Process audio through Whisper model to extract features."""
+ if audio_16k.size(-1) <= 16000 * 30:
+ # If audio is short enough, process in one go
+ inputs = self.whisper_feature_extractor(
+ [audio_16k.squeeze(0).cpu().numpy()],
+ return_tensors="pt",
+ return_attention_mask=True,
+ sampling_rate=16000
+ )
+ input_features = self.whisper_model._mask_input_features(
+ inputs.input_features, attention_mask=inputs.attention_mask
+ ).to(self.device)
+ outputs = self.whisper_model.encoder(
+ input_features.to(self.whisper_model.encoder.dtype),
+ head_mask=None,
+ output_attentions=False,
+ output_hidden_states=False,
+ return_dict=True,
+ )
+ features = outputs.last_hidden_state.to(torch.float32)
+ features = features[:, :audio_16k.size(-1) // 320 + 1]
+ else:
+ # Process long audio in chunks
+ overlapping_time = 5 # 5 seconds
+ features_list = []
+ buffer = None
+ traversed_time = 0
+ while traversed_time < audio_16k.size(-1):
+ if buffer is None: # first chunk
+ chunk = audio_16k[:, traversed_time:traversed_time + 16000 * 30]
+ else:
+ chunk = torch.cat([
+ buffer,
+ audio_16k[:, traversed_time:traversed_time + 16000 * (30 - overlapping_time)]
+ ], dim=-1)
+ inputs = self.whisper_feature_extractor(
+ [chunk.squeeze(0).cpu().numpy()],
+ return_tensors="pt",
+ return_attention_mask=True,
+ sampling_rate=16000
+ )
+ input_features = self.whisper_model._mask_input_features(
+ inputs.input_features, attention_mask=inputs.attention_mask
+ ).to(self.device)
+ outputs = self.whisper_model.encoder(
+ input_features.to(self.whisper_model.encoder.dtype),
+ head_mask=None,
+ output_attentions=False,
+ output_hidden_states=False,
+ return_dict=True,
+ )
+ chunk_features = outputs.last_hidden_state.to(torch.float32)
+ chunk_features = chunk_features[:, :chunk.size(-1) // 320 + 1]
+ if traversed_time == 0:
+ features_list.append(chunk_features)
+ else:
+ features_list.append(chunk_features[:, 50 * overlapping_time:])
+ buffer = chunk[:, -16000 * overlapping_time:]
+ traversed_time += 30 * 16000 if traversed_time == 0 else chunk.size(-1) - 16000 * overlapping_time
+ features = torch.cat(features_list, dim=1)
+
+ return features
+
+ @torch.no_grad()
+ @torch.inference_mode()
+ def convert_voice(self, source, target, diffusion_steps=10, length_adjust=1.0,
+ inference_cfg_rate=0.7, f0_condition=False, auto_f0_adjust=True,
+ pitch_shift=0, stream_output=True):
+ """
+ Convert both timbre and voice from source to target.
+
+ Args:
+ source: Path to source audio file
+ target: Path to target audio file
+ diffusion_steps: Number of diffusion steps (default: 10)
+ length_adjust: Length adjustment factor (default: 1.0)
+ inference_cfg_rate: Inference CFG rate (default: 0.7)
+ f0_condition: Whether to use F0 conditioning (default: False)
+ auto_f0_adjust: Whether to automatically adjust F0 (default: True)
+ pitch_shift: Pitch shift in semitones (default: 0)
+ stream_output: Whether to stream the output (default: True)
+
+ Returns:
+ If stream_output is True, yields (mp3_bytes, full_audio) tuples
+ If stream_output is False, returns the full audio as a numpy array
+ """
+ # Select appropriate models based on F0 condition
+ inference_module = self.model if not f0_condition else self.model_f0
+ mel_fn = self.to_mel if not f0_condition else self.to_mel_f0
+ bigvgan_fn = self.bigvgan_model if not f0_condition else self.bigvgan_44k_model
+ 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_wave_len = self.overlap_frame_len * hop_length
+
+ # 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(self.device)
+ ref_audio = torch.tensor(ref_audio[:sr * 25]).unsqueeze(0).float().to(self.device)
+
+ # Resample to 16kHz for feature extraction
+ ref_waves_16k = torchaudio.functional.resample(ref_audio, sr, 16000)
+ converted_waves_16k = torchaudio.functional.resample(source_audio, sr, 16000)
+
+ # Extract Whisper features
+ S_alt = self._process_whisper_features(converted_waves_16k, is_source=True)
+ S_ori = self._process_whisper_features(ref_waves_16k, is_source=False)
+
+ # Compute mel spectrograms
+ mel = mel_fn(source_audio.to(self.device).float())
+ mel2 = mel_fn(ref_audio.to(self.device).float())
+
+ # Set target lengths
+ target_lengths = torch.LongTensor([int(mel.size(2) * length_adjust)]).to(mel.device)
+ target2_lengths = torch.LongTensor([mel2.size(2)]).to(mel2.device)
+
+ # Compute style features
+ 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 = self.campplus_model(feat2.unsqueeze(0))
+
+ # Process F0 if needed
+ if f0_condition:
+ F0_ori = self.rmvpe.infer_from_audio(ref_waves_16k[0], thred=0.03)
+ F0_alt = self.rmvpe.infer_from_audio(converted_waves_16k[0], thred=0.03)
+
+ if self.device == "mps":
+ F0_ori = torch.from_numpy(F0_ori).float().to(self.device)[None]
+ F0_alt = torch.from_numpy(F0_alt).float().to(self.device)[None]
+ else:
+ F0_ori = torch.from_numpy(F0_ori).to(self.device)[None]
+ F0_alt = torch.from_numpy(F0_alt).to(self.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] = self.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 = 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
+ )
+
+ # Process in chunks for streaming
+ max_source_window = max_context_window - mel2.size(2)
+ processed_frames = 0
+ generated_wave_chunks = []
+ previous_chunk = None
+
+ # 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=self.device.type, dtype=torch.float16):
+ # 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 = bigvgan_fn(vc_target.float())[0]
+
+ processed_frames, previous_chunk, should_break, mp3_bytes, full_audio = self._stream_wave_chunks(
+ vc_wave, processed_frames, vc_target, overlap_wave_len,
+ generated_wave_chunks, previous_chunk, is_last_chunk, stream_output, sr
+ )
+
+ if stream_output and mp3_bytes is not None:
+ yield mp3_bytes, full_audio
+
+ if should_break:
+ if not stream_output:
+ return full_audio
+ break
+
+ if not stream_output:
+ return np.concatenate(generated_wave_chunks)
+
+ return None, None
\ No newline at end of file
diff --git a/seed-vc/train.py b/seed-vc/train.py
new file mode 100644
index 0000000000000000000000000000000000000000..cd5f365c51f5ca124f803896c324d97a3db246cd
--- /dev/null
+++ b/seed-vc/train.py
@@ -0,0 +1,437 @@
+import os
+import sys
+os.environ['HF_HUB_CACHE'] = './checkpoints/hf_cache'
+import torch
+import torch.multiprocessing as mp
+import random
+import librosa
+import yaml
+import argparse
+import torchaudio
+import torchaudio.compliance.kaldi as kaldi
+import glob
+from tqdm import tqdm
+import shutil
+
+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
+ shutil.copyfile(config_path, os.path.join(self.log_dir, os.path.basename(config_path)))
+ 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
+ available_checkpoints = glob.glob(os.path.join(self.log_dir, "DiT_epoch_*_step_*.pth"))
+ if len(available_checkpoints) > 0:
+ 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 we have more than 2
+ 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, training from scratch.")
+
+ def build_sv_model(self, device, config):
+ 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):
+ 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
+ self.bigvgan_model = bigvgan.BigVGAN.from_pretrained(vocoder_name, use_cuda_kernel=False)
+ 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_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)
+ # remove decoder to save memory
+ 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].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().long()
+ prompt_len[torch.rand([B], device=alt_cond.device) < 0.1] = 0
+
+ # for prompt cond token, use 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 the prompt only
+ 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)
+ 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()
+ torch.nn.utils.clip_grad_norm_(self.model.cfm.parameters(), 10.0)
+ 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,
+ f'DiT_epoch_{self.epoch:05d}_step_{self.iters:05d}.pth'
+ )
+ 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:
+ 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 final model..')
+ 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)
+ print(f"Final model saved at {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,
+ device=args.device
+ )
+ trainer.train()
+
+if __name__ == '__main__':
+ if sys.platform == 'win32':
+ mp.freeze_support()
+ mp.set_start_method('spawn', force=True)
+
+ 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)
+ parser.add_argument("--gpu", type=int, help="Which GPU id to use", default=0)
+ args = parser.parse_args()
+ if torch.backends.mps.is_available():
+ args.device = "mps"
+ else:
+ args.device = f"cuda:{args.gpu}" if args.gpu else "cuda:0"
+ main(args)
diff --git a/seed-vc/train_v2.py b/seed-vc/train_v2.py
new file mode 100644
index 0000000000000000000000000000000000000000..a1bb437b3652ed05f71cb3a190469047eec383ba
--- /dev/null
+++ b/seed-vc/train_v2.py
@@ -0,0 +1,345 @@
+import os
+import sys
+os.environ['HF_HUB_CACHE'] = './checkpoints/hf_cache'
+import torch
+import torch.multiprocessing as mp
+import random
+import librosa
+import yaml
+import argparse
+import torchaudio
+import torchaudio.compliance.kaldi as kaldi
+import glob
+import time
+from tqdm import tqdm
+import shutil
+import accelerate
+from optimizers import build_optimizer
+from data.ft_dataset import build_ft_dataloader
+import hydra
+from omegaconf import DictConfig
+
+from accelerate import Accelerator
+from accelerate import DistributedDataParallelKwargs
+from accelerate.logging import get_logger
+
+class Trainer:
+ def __init__(
+ self,
+ config_path,
+ pretrained_cfm_ckpt_path,
+ pretrained_ar_ckpt_path,
+ data_dir,
+ run_name,
+ batch_size=0,
+ num_workers=0,
+ steps=1000,
+ save_interval=500,
+ max_epochs=1000,
+ train_cfm=True,
+ train_ar=False,
+ mixed_precision=None,
+ ):
+ self.config_path = config_path
+ self.mixed_precision = mixed_precision
+
+ # Load configuration
+ self.config = yaml.safe_load(open(config_path))
+
+ # Setup logging directory
+ self.log_dir = os.path.join("runs", run_name)
+ if not os.path.exists(self.log_dir):
+ os.makedirs(self.log_dir, exist_ok=True)
+ shutil.copy(config_path, os.path.join(self.log_dir, os.path.basename(config_path)))
+
+ # Setup accelerator
+ ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True, broadcast_buffers=False)
+ self.accelerator = Accelerator(
+ project_dir=self.log_dir,
+ split_batches=True,
+ kwargs_handlers=[ddp_kwargs],
+ mixed_precision=mixed_precision
+ )
+ self.device = self.accelerator.device
+
+ # Initialize training parameters
+ self._init_dataloader(
+ data_dir=data_dir,
+ batch_size=batch_size,
+ num_workers=num_workers,
+ spect_params=self.config['mel_fn'],
+ sr=self.config['sr'],
+ )
+
+ # Initialize models and optimizers
+ self._init_models(train_cfm=train_cfm, train_ar=train_ar)
+
+ # Load checkpoint if available
+ self._load_checkpoint(pretrained_cfm_ckpt_path, pretrained_ar_ckpt_path)
+
+ # Initialize training parameters
+ self.iters = 0
+ self.start_epoch = 0
+ self.log_interval = 10
+ self.max_steps = steps
+ self.save_interval = save_interval
+ self.max_epochs = max_epochs
+
+ def _init_dataloader(self, data_dir, batch_size, num_workers, spect_params, sr):
+ self.spect_params = spect_params
+ self.sr = sr
+ # Initialize dataloader
+ self.train_dataloader = build_ft_dataloader(
+ data_dir,
+ spect_params,
+ self.sr,
+ batch_size=batch_size,
+ num_workers=num_workers,
+ )
+
+ def _init_models(self, train_cfm=True, train_ar=False):
+ """Initialize models and optimizers"""
+ assert train_cfm or train_ar, "At least one model should be trained"
+ self.train_cfm = train_cfm
+ self.train_ar = train_ar
+ # Initialize main model
+ self._init_main_model(train_cfm=train_cfm, train_ar=train_ar)
+
+ # Initialize optimizers
+ self._init_optimizers()
+
+
+ def _init_main_model(self, train_cfm=True, train_ar=False):
+ """Initialize the main model"""
+ with self.accelerator.main_process_first():
+ cfg = DictConfig(self.config)
+ self.model = hydra.utils.instantiate(cfg).to(self.device)
+ for p in self.model.parameters():
+ p.requires_grad = False
+ if train_cfm:
+ for p in self.model.cfm.parameters():
+ p.requires_grad = True
+ for p in self.model.cfm_length_regulator.parameters():
+ p.requires_grad = True
+ if train_ar:
+ for p in self.model.ar.parameters():
+ p.requires_grad = True
+ for p in self.model.ar_length_regulator.parameters():
+ p.requires_grad = True
+
+
+ def _init_optimizers(self):
+ """Initialize optimizers and schedulers"""
+ from optimizers import build_single_optimizer
+ self.optimizer, self.scheduler = build_single_optimizer(
+ self.model,
+ lr=2e-5,
+ )
+ self.optimizer = self.accelerator.prepare(self.optimizer)
+ self.scheduler = self.accelerator.prepare(self.scheduler)
+
+ def _find_checkpoint(self, name_pattern, max_keep=1):
+ """Find checkpoint files in the specified directory"""
+ available_checkpoints = glob.glob(os.path.join(self.log_dir, name_pattern))
+ if len(available_checkpoints) > max_keep - 1:
+ # 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 self.accelerator.is_main_process
+ and len(available_checkpoints) > max_keep
+ ):
+ os.remove(earliest_checkpoint)
+ print(f"Removed {earliest_checkpoint}")
+ return latest_checkpoint
+ else:
+ return None
+
+ def _load_checkpoint(self, pretrained_cfm_ckpt_path, pretrained_ar_ckpt_path):
+ """Load checkpoint if available"""
+ cfm_checkpoint_path = pretrained_cfm_ckpt_path or self._find_checkpoint("CFM_epoch_*_step_*.pth", max_keep=1)
+ ar_checkpoint_path = pretrained_ar_ckpt_path or self._find_checkpoint("AR_epoch_*_step_*.pth", max_keep=1)
+
+ with self.accelerator.main_process_first():
+ if cfm_checkpoint_path:
+ print(f"Loading CFM checkpoint from {cfm_checkpoint_path}")
+ if ar_checkpoint_path:
+ print(f"Loading AR checkpoint from {ar_checkpoint_path}")
+ self.model.load_checkpoints(cfm_checkpoint_path=cfm_checkpoint_path, ar_checkpoint_path=ar_checkpoint_path)
+ self.model = self.accelerator.prepare(self.model)
+
+ def filter_state_dict_shapes(self, params, model):
+ model_state_dict = model.state_dict()
+ filtered_state_dict = {
+ k: v
+ for k, v in params.items()
+ if k in model_state_dict and v.shape == model_state_dict[k].shape
+ }
+ skipped_keys = set(params.keys()) - set(filtered_state_dict.keys())
+ if skipped_keys:
+ print(
+ f"Warning: Skipped loading some keys due to shape mismatch: {skipped_keys}"
+ )
+ return filtered_state_dict, skipped_keys
+
+ def train(self):
+ """Main training loop"""
+ for epoch in range(self.start_epoch, self.start_epoch + 1000):
+ epoch_start_time = time.time()
+
+ try:
+ self.train_dataloader.sampler.set_epoch(epoch)
+ except AttributeError:
+ pass
+
+ self.model.train()
+
+ for i, batch in enumerate(tqdm(self.train_dataloader)):
+ # Process batch
+ self._process_batch(epoch, i, batch)
+ if self.iters >= self.max_steps and self.accelerator.is_main_process:
+ print("Reached max steps, stopping training")
+ self._save_checkpoint(epoch)
+ exit()
+
+ # Log epoch completion
+ if self.accelerator.is_main_process:
+ print(f"Epoch {epoch} completed in {time.time() - epoch_start_time:.2f} seconds")
+
+ if epoch + 1 >= self.max_epochs and self.accelerator.is_main_process:
+ print("Reached max epochs, stopping training")
+ self._save_checkpoint(epoch)
+ exit()
+
+ def _process_batch(self, epoch, i, batch):
+ """Process a single batch"""
+ # Move batch to device
+ waves, mels, wave_lens, mel_lens = batch
+ # Resample to 16kHz for ASR models
+ waves_16k = torchaudio.functional.resample(waves, self.sr, 16000)
+ wave_lengths_16k = (wave_lens.float() * 16000 / self.sr).long()
+
+ # Forward pass and loss calculation
+ with self.accelerator.autocast():
+ loss_ar, loss_cfm = self.model(
+ waves_16k.to(self.device),
+ mels.to(self.device),
+ wave_lengths_16k.to(self.device),
+ mel_lens.to(self.device),
+ forward_ar=self.train_ar,
+ forward_cfm=self.train_cfm,
+ )
+
+ loss = loss_ar + loss_cfm
+
+ self.accelerator.backward(loss)
+
+ grad_norm_g = torch.nn.utils.clip_grad_norm_(
+ self.model.parameters(), 1000.0
+ )
+ self.optimizer.step()
+ self.scheduler.step(self.iters)
+ self.optimizer.zero_grad()
+
+ # Log training progress
+ self._log_training_progress(epoch, i, loss, loss_ar, loss_cfm, grad_norm_g)
+
+ # Save checkpoint
+ if self.iters != 0 and self.iters % self.save_interval == 0 and self.accelerator.is_main_process:
+ self._save_checkpoint(epoch)
+
+ # Increment iteration counter
+ self.iters += 1
+
+ def _log_training_progress(self, epoch, i, loss, loss_ar, loss_cfm, grad_norm_g):
+ """Log training progress to tensorboard and wandb"""
+ if self.iters % self.log_interval == 0 and self.accelerator.is_main_process:
+ with torch.no_grad():
+ cur_lr = self.scheduler.get_last_lr()[0] if i != 0 else 0
+
+ # Log to console
+ print("Epoch %d, Iteration %d, Loss: %.4f, Loss AR: %.4f, Loss CFM: %.4f, Grad Norm: %.4f, LR: %.6f"
+ % (epoch, i, loss.item(), loss_ar.item(), loss_cfm.item(), grad_norm_g, cur_lr))
+
+ def _save_checkpoint(self, epoch):
+ """Save model checkpoint"""
+ print('Saving checkpoint...')
+ if self.train_ar:
+ state = {
+ 'net': {
+ 'ar': self.accelerator.unwrap_model(self.model).ar.state_dict(),
+ 'length_regulator': self.accelerator.unwrap_model(self.model).ar_length_regulator.state_dict(),
+ },
+ 'iters': self.iters,
+ 'epoch': epoch,
+ }
+ save_path = os.path.join(self.log_dir, 'AR_epoch_%05d_step_%05d.pth' % (epoch, self.iters))
+ torch.save(state, save_path)
+ print(f"Saved AR checkpoint to {save_path}")
+
+ # Find all checkpoints and remove old ones
+ self._remove_old_checkpoints("AR_epoch_*_step_*.pth", max_keep=1)
+ if self.train_cfm:
+ state = {
+ 'net': {
+ 'cfm': self.accelerator.unwrap_model(self.model).cfm.state_dict(),
+ 'length_regulator': self.accelerator.unwrap_model(self.model).cfm_length_regulator.state_dict(),
+ },
+ 'iters': self.iters,
+ 'epoch': epoch,
+ }
+ save_path = os.path.join(self.log_dir, 'CFM_epoch_%05d_step_%05d.pth' % (epoch, self.iters))
+ torch.save(state, save_path)
+ print(f"Saved CFM checkpoint to {save_path}")
+
+ # Find all checkpoints and remove old ones
+ self._remove_old_checkpoints("CFM_epoch_*_step_*.pth", max_keep=1)
+ def _remove_old_checkpoints(self, name_pattern, max_keep=1):
+ """Remove old checkpoints"""
+ checkpoints = glob.glob(os.path.join(self.log_dir, name_pattern))
+ if len(checkpoints) > max_keep:
+ # Sort by step
+ checkpoints.sort(key=lambda x: int(x.split('_')[-1].split('.')[0]))
+ # Remove all except last 1
+ for cp in checkpoints[:-max_keep]:
+ os.remove(cp)
+
+def main(args):
+ trainer = Trainer(
+ config_path=args.config,
+ pretrained_cfm_ckpt_path=args.pretrained_cfm_ckpt,
+ pretrained_ar_ckpt_path=args.pretrained_ar_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,
+ train_cfm=args.train_cfm,
+ train_ar=args.train_ar,
+ )
+ trainer.train()
+
+if __name__ == '__main__':
+ parser = argparse.ArgumentParser()
+ parser.add_argument('--config', type=str, default='configs/v2/vc_wrapper.yaml')
+ parser.add_argument('--pretrained-cfm-ckpt', type=str, default=None)
+ parser.add_argument('--pretrained-ar-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)
+ parser.add_argument('--train-cfm', action='store_true', help='Train CFM model')
+ parser.add_argument('--train-ar', action='store_true', help='Train AR model')
+ args = parser.parse_args()
+ main(args)