Bug fixes for offline mode

Author: SortAnon

controllable_talknet.py

@@ -36,9 +36,14 @@
 from denoiser import Denoiser
 
 app = JupyterDash(__name__)
-UPLOAD_DIRECTORY = "/content"
+RUN_PATH = os.path.dirname(os.path.realpath(__file__))
+if RUN_PATH == "/content":
+    UI_MODE = "colab"
+else:
+    UI_MODE = "offline"
 torch.set_grad_enabled(False)
 
+app.title = "Controllable TalkNet"
 app.layout = html.Div(
     children=[
         html.H1(
@@ -89,7 +94,7 @@
             },
         ),
         html.Label(
-            "Upload reference audio to " + UPLOAD_DIRECTORY,
+            "Upload reference audio to " + RUN_PATH,
             htmlFor="reference-dropdown",
         ),
         dcc.Store(id="current-f0s"),
@@ -453,27 +458,23 @@ def preprocess_tokens(tokens, blank):
 
 
 def get_duration(wav_name, transcript):
-    if not os.path.exists(os.path.join(UPLOAD_DIRECTORY, "output")):
-        os.mkdir(os.path.join(UPLOAD_DIRECTORY, "output"))
+    if not os.path.exists(os.path.join(RUN_PATH, "temp")):
+        os.mkdir(os.path.join(RUN_PATH, "temp"))
     if "_" not in transcript:
         generate_json(
-            os.path.join(UPLOAD_DIRECTORY, "output", wav_name + "_conv.wav")
+            os.path.join(RUN_PATH, "temp", wav_name + "_conv.wav")
             + "|"
             + transcript.strip(),
-            os.path.join(UPLOAD_DIRECTORY, "output", wav_name + ".json"),
+            os.path.join(RUN_PATH, "temp", wav_name + ".json"),
         )
     else:
         generate_json(
-            os.path.join(UPLOAD_DIRECTORY, "output", wav_name + "_conv.wav")
-            + "|"
-            + "dummy",
-            os.path.join(UPLOAD_DIRECTORY, "output", wav_name + ".json"),
+            os.path.join(RUN_PATH, "temp", wav_name + "_conv.wav") + "|" + "dummy",
+            os.path.join(RUN_PATH, "temp", wav_name + ".json"),
         )
 
     data_config = {
-        "manifest_filepath": os.path.join(
-            UPLOAD_DIRECTORY, "output", wav_name + ".json"
-        ),
+        "manifest_filepath": os.path.join(RUN_PATH, "temp", wav_name + ".json"),
         "sample_rate": 22050,
         "batch_size": 1,
     }
@@ -645,7 +646,7 @@ def update_pitch_options(value):
 def update_filelist(n_clicks):
     filelist = []
     supported_formats = [".wav", ".ogg", ".mp3", "flac", ".aac"]
-    for x in sorted(os.listdir(UPLOAD_DIRECTORY)):
+    for x in sorted(os.listdir(RUN_PATH)):
         if x[-4:].lower() in supported_formats:
             filelist.append({"label": x, "value": x})
     return filelist
@@ -664,18 +665,18 @@ def update_filelist(n_clicks):
 )
 def select_file(dropdown_value):
     if dropdown_value is not None:
-        if not os.path.exists(os.path.join(UPLOAD_DIRECTORY, "output")):
-            os.mkdir(os.path.join(UPLOAD_DIRECTORY, "output"))
-        ffmpeg.input(os.path.join(UPLOAD_DIRECTORY, dropdown_value)).output(
-            os.path.join(UPLOAD_DIRECTORY, "output", dropdown_value + "_conv.wav"),
+        if not os.path.exists(os.path.join(RUN_PATH, "temp")):
+            os.mkdir(os.path.join(RUN_PATH, "temp"))
+        ffmpeg.input(os.path.join(RUN_PATH, dropdown_value)).output(
+            os.path.join(RUN_PATH, "temp", dropdown_value + "_conv.wav"),
             ar="22050",
             ac="1",
             acodec="pcm_s16le",
             map_metadata="-1",
             fflags="+bitexact",
         ).overwrite_output().run(quiet=True)
         fo_with_silence, f0_wo_silence = crepe_f0(
-            os.path.join(UPLOAD_DIRECTORY, "output", dropdown_value + "_conv.wav")
+            os.path.join(RUN_PATH, "temp", dropdown_value + "_conv.wav")
         )
         return [
             "Analyzed " + dropdown_value,
@@ -724,25 +725,25 @@ def download_model(model, custom_model):
             drive_id = model
         if drive_id == "" or drive_id is None:
             return ("Missing Drive ID", None, None)
-        if not os.path.exists(os.path.join(UPLOAD_DIRECTORY, "models")):
-            os.mkdir(os.path.join(UPLOAD_DIRECTORY, "models"))
-        if not os.path.exists(os.path.join(UPLOAD_DIRECTORY, "models", drive_id)):
-            os.mkdir(os.path.join(UPLOAD_DIRECTORY, "models", drive_id))
-            zip_path = os.path.join(UPLOAD_DIRECTORY, "models", drive_id, "model.zip")
+        if not os.path.exists(os.path.join(RUN_PATH, "models")):
+            os.mkdir(os.path.join(RUN_PATH, "models"))
+        if not os.path.exists(os.path.join(RUN_PATH, "models", drive_id)):
+            os.mkdir(os.path.join(RUN_PATH, "models", drive_id))
+            zip_path = os.path.join(RUN_PATH, "models", drive_id, "model.zip")
             gdown.download(
                 d + drive_id,
                 zip_path,
                 quiet=False,
             )
             if not os.path.exists(zip_path):
-                os.rmdir(os.path.join(UPLOAD_DIRECTORY, "models", drive_id))
+                os.rmdir(os.path.join(RUN_PATH, "models", drive_id))
                 return ("Model download failed", None, None)
             if os.stat(zip_path).st_size < 16:
                 os.remove(zip_path)
-                os.rmdir(os.path.join(UPLOAD_DIRECTORY, "models", drive_id))
+                os.rmdir(os.path.join(RUN_PATH, "models", drive_id))
                 return ("Model zip is empty", None, None)
             with zipfile.ZipFile(zip_path, "r") as zip_ref:
-                zip_ref.extractall(os.path.join(UPLOAD_DIRECTORY, "models", drive_id))
+                zip_ref.extractall(os.path.join(RUN_PATH, "models", drive_id))
             os.remove(zip_path)
 
         # Download super-resolution HiFi-GAN
@@ -757,8 +758,8 @@ def download_model(model, custom_model):
 
         return (
             None,
-            os.path.join(UPLOAD_DIRECTORY, "models", drive_id, "TalkNetSpect.nemo"),
-            os.path.join(UPLOAD_DIRECTORY, "models", drive_id, "hifiganmodel"),
+            os.path.join(RUN_PATH, "models", drive_id, "TalkNetSpect.nemo"),
+            os.path.join(RUN_PATH, "models", drive_id, "hifiganmodel"),
         )
     except Exception as e:
         return (str(e), None, None)
@@ -879,7 +880,10 @@ def generate_audio(
                 spect = tnmodel.force_spectrogram(
                     tokens=tokens,
                     durs=torch.from_numpy(durs).view(1, -1).to("cuda:0"),
-                    f0=torch.FloatTensor(f0s).view(1, -1).to("cuda:0"),
+                    f0=torch.FloatTensor(f0s)
+                    .view(1, -1)
+                    .type(torch.LongTensor)
+                    .to("cuda:0"),
                 )
 
             if hifipath != hifigan_path:

requirements.txt

@@ -0,0 +1,26 @@
+tensorflow==2.4.1
+tensorboard==2.4.1
+dash==1.21.0
+jupyter-dash==0.4.0
+psola==0.0.1
+wget==3.2
+unidecode==1.2.0
+pysptk==0.1.18
+frozendict==2.0.3
+torch==1.8.1+cu111
+torchvision==0.9.1+cu111
+torchaudio==0.8.1
+torchtext==0.9.1
+torch_stft==0.1.4
+kaldiio==2.17.2
+pydub==0.25.1
+pyannote.audio==1.1.2
+g2p_en==2.1.0
+pesq==0.0.2
+pystoi==0.3.3
+crepe==0.0.12
+resampy==0.2.2
+ffmpeg-python==0.2.0
+tqdm
+gdown==3.13.0
+editdistance==0.5.3

talknet_offline.py

@@ -0,0 +1,7 @@
+if __name__ == '__main__':
+    from controllable_talknet import *
+    app.run_server(
+        mode="external",
+        debug=False,
+        threaded=True,
+    )