seed-vc: change output name, put examples in the container, use cuda 12.1

2026-10-09 00:09:55 +00:00 · 2024-10-23 16:56:13 +05:30
parent 7d5fd4f029
commit 7e12c7c3a5
3 changed files with 138 additions and 68 deletions
@@ -1,3 +1,2 @@
 outputs
-examples
 Dockerfile
@@ -1,4 +1,4 @@
-FROM pytorch/pytorch:2.4.0-cuda12.4-cudnn9-devel
+FROM pytorch/pytorch:2.5.0-cuda12.1-cudnn9-runtime

 USER root
 ENV SHELL=/bin/bash
@@ -6,7 +6,7 @@ import os
 import os.path as osp
 import yaml

-warnings.simplefilter('ignore')
+warnings.simplefilter("ignore")

 # load packages
 import random
@@ -24,31 +24,44 @@ from hf_utils import load_custom_model_from_hf
 # Load model and configuration
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

+
 def load_models(args):
    if not args.f0_condition:
-        dit_checkpoint_path, dit_config_path = load_custom_model_from_hf("Plachta/Seed-VC",
-                                                                         "DiT_seed_v2_uvit_whisper_small_wavenet_bigvgan_pruned.pth",
-                                                                         "config_dit_mel_seed_uvit_whisper_small_wavenet.yml")
+        dit_checkpoint_path, dit_config_path = load_custom_model_from_hf(
+            "Plachta/Seed-VC",
+            "DiT_seed_v2_uvit_whisper_small_wavenet_bigvgan_pruned.pth",
+            "config_dit_mel_seed_uvit_whisper_small_wavenet.yml",
+        )
        f0_extractor = None
    else:
-        dit_checkpoint_path, dit_config_path = load_custom_model_from_hf("Plachta/Seed-VC",
-                                                                         "DiT_seed_v2_uvit_facodec_small_wavenet_f0_bigvgan_pruned.pth",
-                                                                         "config_dit_mel_seed_facodec_small_wavenet_f0.yml")
+        dit_checkpoint_path, dit_config_path = load_custom_model_from_hf(
+            "Plachta/Seed-VC",
+            "DiT_seed_v2_uvit_facodec_small_wavenet_f0_bigvgan_pruned.pth",
+            "config_dit_mel_seed_facodec_small_wavenet_f0.yml",
+        )
        # f0 extractor
        from modules.rmvpe import RMVPE

-        model_path = load_custom_model_from_hf("lj1995/VoiceConversionWebUI", "rmvpe.pt", None)
+        model_path = load_custom_model_from_hf(
+            "lj1995/VoiceConversionWebUI", "rmvpe.pt", None
+        )
        f0_extractor = RMVPE(model_path, is_half=False, device=device)

-    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']
+    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)
+    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)
@@ -57,24 +70,31 @@ def load_models(args):
    # 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_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)

    from modules.bigvgan import bigvgan
-    bigvgan_model = bigvgan.BigVGAN.from_pretrained('nvidia/bigvgan_v2_22khz_80band_256x', use_cuda_kernel=False)
+
+    bigvgan_model = bigvgan.BigVGAN.from_pretrained(
+        "nvidia/bigvgan_v2_22khz_80band_256x", 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)

    if model_params.speech_tokenizer.type == "facodec":
-        ckpt_path, config_path = load_custom_model_from_hf("Plachta/FAcodec", 'pytorch_model.bin', 'config.yml')
+        ckpt_path, config_path = load_custom_model_from_hf(
+            "Plachta/FAcodec", "pytorch_model.bin", "config.yml"
+        )

        codec_config = yaml.safe_load(open(config_path))
-        codec_model_params = recursive_munch(codec_config['model_params'])
+        codec_model_params = recursive_munch(codec_config["model_params"])
        codec_encoder = build_model(codec_model_params, stage="codec")

        ckpt_params = torch.load(ckpt_path, map_location="cpu")
@@ -83,43 +103,71 @@ def load_models(args):
            codec_encoder[key].load_state_dict(ckpt_params[key], strict=False)
        _ = [codec_encoder[key].eval() for key in codec_encoder]
        _ = [codec_encoder[key].to(device) for key in codec_encoder]
-        speechtokenizer_set = ('facodec', codec_encoder, None)
+        speechtokenizer_set = ("facodec", codec_encoder, None)
    elif model_params.speech_tokenizer.type == "whisper":
        from transformers import AutoFeatureExtractor, WhisperModel
-        whisper_name = model_params.speech_tokenizer.whisper_name if hasattr(model_params.speech_tokenizer, 'whisper_name') else "whisper-large-v3"
-        whisper_model = WhisperModel.from_pretrained(whisper_name, torch_dtype=torch.float16).to(device)
+
+        whisper_name = (
+            model_params.speech_tokenizer.whisper_name
+            if hasattr(model_params.speech_tokenizer, "whisper_name")
+            else "whisper-large-v3"
+        )
+        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)
-        speechtokenizer_set = ('whisper', whisper_model, whisper_feature_extractor)
+        speechtokenizer_set = ("whisper", whisper_model, whisper_feature_extractor)
    else:
-        raise ValueError(f"Unsupported speech tokenizer type: {model_params.speech_tokenizer.type}")
-
+        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'],
+        "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'].get('fmax', "None") == "None" else 8000,
-        "center": False
+        "fmin": config["preprocess_params"].get("fmin", 0),
+        "fmax": (
+            None if config["preprocess_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, speechtokenizer_set, f0_extractor, bigvgan_model, campplus_model, to_mel, mel_fn_args
+    return (
+        model,
+        speechtokenizer_set,
+        f0_extractor,
+        bigvgan_model,
+        campplus_model,
+        to_mel,
+        mel_fn_args,
+    )
+

 def adjust_f0_semitones(f0_sequence, n_semitones):
    factor = 2 ** (n_semitones / 12)
    return f0_sequence * factor

+
@torch.no_grad()
 def main(args):
-    model, speechtokenizer_set, f0_extractor, bigvgan_model, campplus_model, to_mel, mel_fn_args = load_models(args)
-    sr = mel_fn_args['sampling_rate']
+    (
+        model,
+        speechtokenizer_set,
+        f0_extractor,
+        bigvgan_model,
+        campplus_model,
+        to_mel,
+        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
@@ -132,19 +180,21 @@ def main(args):
    source_audio = librosa.load(source, sr=sr)[0]
    ref_audio = librosa.load(target_name, sr=sr)[0]

-    source_audio = source_audio[:sr * 30]
+    source_audio = source_audio[: sr * 30]
    source_audio = torch.tensor(source_audio).unsqueeze(0).float().to(device)

-    ref_audio = ref_audio[:(sr * 30 - source_audio.size(-1))]
+    ref_audio = ref_audio[: (sr * 30 - source_audio.size(-1))]
    ref_audio = torch.tensor(ref_audio).unsqueeze(0).float().to(device)

    source_waves_16k = torchaudio.functional.resample(source_audio, sr, 16000)
    ref_waves_16k = torchaudio.functional.resample(ref_audio, sr, 16000)

    converted_waves_24k = torchaudio.functional.resample(source_audio, sr, 24000)
-    wave_lengths_24k = torch.LongTensor([converted_waves_24k.size(1)]).to(converted_waves_24k.device)
+    wave_lengths_24k = torch.LongTensor([converted_waves_24k.size(1)]).to(
+        converted_waves_24k.device
+    )
    waves_input = converted_waves_24k.unsqueeze(1)
-    if speechtokenizer_set[0] == 'facodec':
+    if speechtokenizer_set[0] == "facodec":
        codec_encoder = speechtokenizer_set[1]
        z = codec_encoder.encoder(waves_input)
        (quantized, codes) = codec_encoder.quantizer(z, waves_input)
@@ -156,15 +206,18 @@ def main(args):
        z = codec_encoder.encoder(waves_input)
        (quantized, codes) = codec_encoder.quantizer(z, waves_input)
        S_ori = torch.cat([codes[1], codes[0]], dim=1)
-    elif speechtokenizer_set[0] == 'whisper':
+    elif speechtokenizer_set[0] == "whisper":
        whisper_model = speechtokenizer_set[1]
        whisper_feature_extractor = speechtokenizer_set[2]
        converted_waves_16k = torchaudio.functional.resample(source_audio, sr, 16000)
-        alt_inputs = whisper_feature_extractor([converted_waves_16k.squeeze(0).cpu().numpy()],
-                                               return_tensors="pt",
-                                               return_attention_mask=True,)
+        alt_inputs = whisper_feature_extractor(
+            [converted_waves_16k.squeeze(0).cpu().numpy()],
+            return_tensors="pt",
+            return_attention_mask=True,
+        )
        alt_input_features = whisper_model._mask_input_features(
-            alt_inputs.input_features, attention_mask=alt_inputs.attention_mask).to(device)
+            alt_inputs.input_features, attention_mask=alt_inputs.attention_mask
+        ).to(device)
        with torch.no_grad():
            alt_outputs = whisper_model.encoder(
                alt_input_features.to(whisper_model.encoder.dtype),
@@ -174,14 +227,17 @@ def main(args):
                return_dict=True,
            )
        S_alt = alt_outputs.last_hidden_state.to(torch.float32)
-        S_alt = S_alt[:, :converted_waves_16k.size(-1)//320 + 1]
+        S_alt = S_alt[:, : converted_waves_16k.size(-1) // 320 + 1]

        ori_waves_16k = torchaudio.functional.resample(ref_audio, sr, 16000)
-        ori_inputs = whisper_feature_extractor([ori_waves_16k.squeeze(0).cpu().numpy()],
-                                               return_tensors="pt",
-                                               return_attention_mask=True)
+        ori_inputs = whisper_feature_extractor(
+            [ori_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)
+            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),
@@ -191,7 +247,7 @@ def main(args):
                return_dict=True,
            )
        S_ori = ori_outputs.last_hidden_state.to(torch.float32)
-        S_ori = S_ori[:, :ori_waves_16k.size(-1) // 320 + 1]
+        S_ori = S_ori[:, : ori_waves_16k.size(-1) // 320 + 1]
    else:
        raise ValueError(f"Unsupported speech tokenizer type: {speechtokenizer_set[0]}")

@@ -201,16 +257,17 @@ def main(args):
    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 = 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))

    if f0_condition:
        waves_16k = torchaudio.functional.resample(waves_24k, sr, 16000)
-        converted_waves_16k = torchaudio.functional.resample(converted_waves_24k, sr, 16000)
+        converted_waves_16k = torchaudio.functional.resample(
+            converted_waves_24k, sr, 16000
+        )
        F0_ori = f0_extractor.infer_from_audio(waves_16k[0], thred=0.03)
        F0_alt = f0_extractor.infer_from_audio(converted_waves_16k[0], thred=0.03)

@@ -228,31 +285,45 @@ def main(args):
        # 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_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)
+            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, _, prompt_codes, commitment_loss, codebook_loss = model.length_regulator(S_ori, ylens=target2_lengths, n_quantizers=3, f0=F0_ori)
+    cond, _, codes, commitment_loss, codebook_loss = model.length_regulator(
+        S_alt, ylens=target_lengths, n_quantizers=3, f0=shifted_f0_alt
+    )
+    prompt_condition, _, prompt_codes, commitment_loss, codebook_loss = (
+        model.length_regulator(S_ori, ylens=target2_lengths, n_quantizers=3, f0=F0_ori)
+    )
    cat_condition = torch.cat([prompt_condition, cond], dim=1)

    time_vc_start = time.time()
    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):]
+        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
    # if f0_condition:
-    vc_wave = bigvgan_model(vc_target).squeeze(1)  # wav_gen is FloatTensor with shape [B(1), 1, T_time] and values in [-1, 1]
+    vc_wave = bigvgan_model(vc_target).squeeze(
+        1
+    )  # wav_gen is FloatTensor with shape [B(1), 1, T_time] and values in [-1, 1]

    time_vc_end = time.time()
    print(f"RTF: {(time_vc_end - time_vc_start) / vc_wave.size(-1) * sr}")
@@ -260,7 +331,7 @@ def main(args):
    source_name = source.split("/")[-1].split(".")[0]
    target_name = target_name.split("/")[-1].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)
+    torchaudio.save(os.path.join(args.output, f"vc_out.wav"), vc_wave.cpu(), sr)


 if __name__ == "__main__":