Split download into two functions

Author: SortAnon

controllable_talknet.py

@@ -21,7 +21,7 @@
 import time
 import uuid
 from core import extract, vocoder, reconstruct
-from core.download import download_from_drive
+from core.download import download_from_drive, download_reconst
 
 app = JupyterDash(__name__)
 DEVICE = "cuda:0"
@@ -781,6 +781,14 @@ def generate_audio(
 
             # Reconstruction
             if "srec" in pitch_options:
+                load_error = download_reconst(RUN_PATH)
+                if load_error is not None:
+                    return [
+                        None,
+                        load_error,
+                        playback_hide,
+                        None,
+                    ]
                 new_spect = reconstruct_inst.reconstruct(spect)
                 if rec_voc is None:
                     rec_voc = vocoder.HiFiGAN(

core/download.py

@@ -57,6 +57,19 @@ def download_from_drive(model, custom_model, run_path):
         if not os.path.exists(sr_path):
             raise Exception("Super-res model failed to download!")
 
+        return (
+            None,
+            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)
+
+
+def download_reconst(run_path):
+    try:
+        d = "https://drive.google.com/uc?id="
+
         # Download VQGAN reconstruction model
         rec_path = os.path.join(
             run_path,
@@ -83,10 +96,6 @@ def download_from_drive(model, custom_model, run_path):
         if not os.path.exists(rec_path):
             raise Exception("Reconstruction HiFi-GAN failed to download!")
 
-        return (
-            None,
-            os.path.join(run_path, "models", drive_id, "TalkNetSpect.nemo"),
-            os.path.join(run_path, "models", drive_id, "hifiganmodel"),
-        )
+        return None
     except Exception as e:
-        return (str(e), None, None)
+        return str(e)