Update README.md

README.md

@@ -34,13 +34,57 @@ Please notice that we encourage you to read our tutorials and learn more about
 ### Using the Vocoder
 
 ```python
-import torch
-from speechbrain.pretrained import UnitHIFIGAN
-
-hifi_gan_unit = UnitHIFIGAN.from_hparams(source="chaanks/hifigan-unit-hubert-l7-k128-ljspeech-libritts", savedir="tmpdir_vocoder")
-codes = torch.randint(0, 99, (100,))
-waveform = hifi_gan_unit.decode_unit(codes)
-
+import torchaudio
+from speechbrain.inference.encoders import MelSpectrogramEncoder
+from speechbrain.inference.vocoders import UnitHIFIGAN
+from speechbrain.lobes.models.huggingface_transformers.discrete_hubert import (
+    DiscreteHuBERT,
+)
+
+speaker_encoder_source = "speechbrain/spkrec-ecapa-voxceleb-mel-spec"
+speech_encoder_source = "facebook/hubert-base-ls960"
+kmeans_folder = "speechbrain/SSL_Quantization"
+kmeans_filename = "LJSpeech_hubert_k128_L7.pt"
+layer = 7
+vocoder_source = "chaanks/hifigan-unit-hubert-l7-k128-ljspeech-libritts"
+save_path = "tmpdir"
+device = "cuda"
+sample_rate = 16000
+
+wav = "chaanks/hifigan-unit-hubert-l7-k128-ljspeech-libritts/test.wav"
+
+speaker_encoder = MelSpectrogramEncoder.from_hparams(
+    source=speaker_encoder_source,
+    run_opts={"device": str(device)},
+    savedir=save_path + "/spk_encoder",
+)
+
+speech_encoder = DiscreteHuBERT(
+    source=speech_encoder_source,
+    save_path=save_path + "/speech_encoder",
+    kmeans_filename=kmeans_filename,
+    kmeans_cache_dir=save_path + "/kmeans",
+    kmeans_repo_id=kmeans_folder,
+    output_norm=False,
+    freeze=True,
+    freeze_feature_extractor=True,
+    apply_spec_augment=False,
+    output_all_hiddens=True,
+    ssl_layer_num=layer,
+).to(device)
+
+vocoder = UnitHIFIGAN.from_hparams(
+    source=vocoder_source,
+    run_opts={"device": str(device)},
+    savedir=save_path + "/vocoder",
+)
+
+audio = speaker_encoder.load_audio(wav)
+audio = audio.to(device)
+spk = speaker_encoder.encode_waveform(audio)
+_, codes = speech_encoder(audio.unsqueeze(0))
+waveform = vocoder.decode_unit(codes.squeeze(0), spk=spk.reshape(-1))
+torchaudio.save("test.wav", waveform.cpu(), sample_rate=sample_rate)
 ```