Update app.py

app.py

@@ -4,8 +4,13 @@ import soundfile as sf
 import gradio as gr
 import torchaudio
 import os
+from huggingface_hub import hf_hub_download
 
-from Amphion.models.ns3_codec import FACodecEncoder, FACodecDecoder
+from Amphion.models.ns3_codec import (
+    FACodecEncoder,
+    FACodecDecoder,
+    FACodecRedecoder,
+)
 
 fa_encoder = FACodecEncoder(
     ngf=32,
@@ -31,15 +36,27 @@ fa_decoder = FACodecDecoder(
     use_gr_residual_phone=True,
 )
 
-fa_encoder.load_state_dict(torch.load("ns3_facodec_encoder.bin"))
-fa_decoder.load_state_dict(torch.load("ns3_facodec_decoder.bin"))
+fa_redecoder = FACodecRedecoder()
+
+# encoder_ckpt = hf_hub_download(repo_id="amphion/naturalspeech3_facodec", filename="ns3_facodec_encoder.bin")
+# decoder_ckpt = hf_hub_download(repo_id="amphion/naturalspeech3_facodec", filename="ns3_facodec_decoder.bin")
+# redecoder_ckpt = hf_hub_download(repo_id="amphion/naturalspeech3_facodec", filename="ns3_facodec_redecoder.bin")
+
+encoder_ckpt = "ns3_facodec_encoder.bin"
+decoder_ckpt = "ns3_facodec_decoder.bin"
+redecoder_ckpt = "ns3_facodec_redecoder.bin"
+
+fa_encoder.load_state_dict(torch.load(encoder_ckpt))
+fa_decoder.load_state_dict(torch.load(decoder_ckpt))
+fa_redecoder.load_state_dict(torch.load(redecoder_ckpt))
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 fa_encoder = fa_encoder.to(device)
 fa_decoder = fa_decoder.to(device)
+fa_redecoder = fa_redecoder.to(device)
 fa_encoder.eval()
 fa_decoder.eval()
-
+fa_redecoder.eval()
 
 def codec_inference(speech_path):
 
@@ -61,23 +78,69 @@ def codec_inference(speech_path):
     return result_path
 
 
+def codec_voice_conversion(speech_path_a, speech_path_b):
+
+    with torch.no_grad():
+
+        wav_a, sr = librosa.load(speech_path_a, sr=16000)
+        wav_a = torch.tensor(wav_a).to(device).unsqueeze(0).unsqueeze(0)
+        wav_b, sr = librosa.load(speech_path_b, sr=16000)
+        wav_b = torch.tensor(wav_b).to(device).unsqueeze(0).unsqueeze(0)
+
+        enc_out_a = fa_encoder(wav_a)
+        enc_out_b = fa_encoder(wav_b)
+
+        vq_post_emb_a, vq_id_a, _, quantized, spk_embs_a = fa_decoder(
+            enc_out_a, eval_vq=False, vq=True
+        )
+        vq_post_emb_b, vq_id_b, _, quantized, spk_embs_b = fa_decoder(
+            enc_out_b, eval_vq=False, vq=True
+        )
+
+        recon_wav_a = fa_decoder.inference(vq_post_emb_a, spk_embs_a)
+        recon_wav_b = fa_decoder.inference(vq_post_emb_b, spk_embs_b)
+
+        vq_post_emb_a_to_b = fa_redecoder.vq2emb(
+            vq_id_a, spk_embs_b, use_residual=False
+        )
+        recon_wav_a_to_b = fa_redecoder.inference(vq_post_emb_a_to_b, spk_embs_b)
+
+    os.makedirs("temp", exist_ok=True)
+    recon_a_result_path = "temp/result_a.wav"
+    recon_b_result_path = "temp/result_b.wav"
+    vc_result_path = "temp/result_vc.wav"
+    sf.write(vc_result_path, recon_wav_a_to_b[0, 0].cpu().numpy(), 16000)
+    sf.write(recon_a_result_path, recon_wav_a[0, 0].cpu().numpy(), 16000)
+    sf.write(recon_b_result_path, recon_wav_b[0, 0].cpu().numpy(), 16000)
+
+    return recon_a_result_path, recon_b_result_path, vc_result_path
+
+
 demo_inputs = [
     gr.Audio(
         sources=["upload", "microphone"],
-        label="Upload the speech file",
+        label="Upload the source speech file",
+        type="filepath",
+    ),
+    gr.Audio(
+        sources=["upload", "microphone"],
+        label="Upload the reference speech file",
         type="filepath",
     ),
 ]
 
-demo_outputs = gr.Audio(label="")
+demo_outputs = [
+    gr.Audio(label="Source speech reconstructed"),
+    gr.Audio(label="Reference speech reconstructed"),
+    gr.Audio(label="Voice conversion result"),
+]
 
 demo = gr.Interface(
-    fn=codec_inference,
+    fn=codec_voice_conversion,
     inputs=demo_inputs,
     outputs=demo_outputs,
     title="NaturalSpeech3 FACodec",
-    description=
-    """
+    description="""
     ## FACodec: Speech Codec with Attribute Factorization used for NaturalSpeech 3
 
     [![arXiv](https://img.shields.io/badge/arXiv-Paper-<COLOR>.svg)](https://arxiv.org/pdf/2403.03100.pdf)
@@ -96,3 +159,4 @@ demo = gr.Interface(
 
 if __name__ == "__main__":
     demo.launch()
+