Update app.py

app.py

@@ -1,8 +1,9 @@
 import torch
 import os
+import torchaudio
 import gradio as gr
 import matplotlib.pyplot as plt
-from scipy.io.wavfile import write
+
 
 device="cpu"
 
@@ -23,21 +24,26 @@ tacotron2 = tacotron2.to(device)
 tacotron2.eval()
 
 # Load Nvidia Waveglow from Hub
-waveglow = torch.hub.load(
-    "NVIDIA/DeepLearningExamples:torchhub",
-    "nvidia_waveglow",
-    model_math="fp32",
-    pretrained=False,
-)
-checkpoint = torch.hub.load_state_dict_from_url(
-    "https://api.ngc.nvidia.com/v2/models/nvidia/waveglowpyt_fp32/versions/1/files/nvidia_waveglowpyt_fp32_20190306.pth",  # noqa: E501
-    progress=False,
-    map_location=device,
-)
-state_dict = {key.replace("module.", ""): value for key, value in checkpoint["state_dict"].items()}
+# waveglow = torch.hub.load(
+#     "NVIDIA/DeepLearningExamples:torchhub",
+#     "nvidia_waveglow",
+#     model_math="fp32",
+#     pretrained=False,
+# )
+# checkpoint = torch.hub.load_state_dict_from_url(
+#     "https://api.ngc.nvidia.com/v2/models/nvidia/waveglowpyt_fp32/versions/1/files/nvidia_waveglowpyt_fp32_20190306.pth",  # noqa: E501
+#     progress=False,
+#     map_location=device,
+# )
+# state_dict = {key.replace("module.", ""): value for key, value in checkpoint["state_dict"].items()}
+
+# waveglow.load_state_dict(state_dict)
+# waveglow = waveglow.remove_weightnorm(waveglow)
+# waveglow = waveglow.to(device)
+# waveglow.eval()
 
-waveglow.load_state_dict(state_dict)
-waveglow = waveglow.remove_weightnorm(waveglow)
+waveglow_pretrained_model = os.path.join(os.getcwd(), 'waveglow_256channels_ljs_v3.pt')
+waveglow = torch.load(waveglow_pretrained_model, map_location=device)['model']
 waveglow = waveglow.to(device)
 waveglow.eval()
 
@@ -61,9 +67,10 @@ def inference(text):
     #Save Audio
     audio_numpy = audio[0].data.cpu().numpy()
     rate = 22050
-    write("output.wav", rate, audio_numpy)
+    write("output1.wav", rate, audio_numpy)
+    torchaudio.save("output2.wav", audio[0:1].cpu(), sample_rate=22050)
     
-    return "output.wav","test.png"
+    return "output1.wav", "output2.wav", "test.png"
   
 title="TACOTRON 2"
 description="Nepali Speech TACOTRON 2: The Tacotron 2 model for generating mel spectrograms from text. To use it, simply add you text or click on one of the examples to load them. Read more at the links below."