Load local model

app.py

@@ -15,7 +15,11 @@ DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 # The model will be loaded in the future
 MODEL_NAME = CHECKPOINTS_DIR / "audiosep_base_4M_steps.ckpt"
-MODEL = None
+MODEL = build_audiosep(
+    config_yaml="config/audiosep_base.yaml",
+    checkpoint_path=MODEL_NAME,
+    device=DEVICE,
+)
 
 
 description = """
@@ -28,15 +32,6 @@ numerous tasks such as audio event separation, musical instrument separation, an
 """
 
 
-def get_model():
-    model = build_audiosep(
-        config_yaml="config/audiosep_base.yaml",
-        checkpoint_path=MODEL_NAME,
-        device=DEVICE,
-    )
-    return model
-
-
 def inference(audio_file_path: str, text: str):
     print(f"Separate audio from [{audio_file_path}] with textual query [{text}]")
     mixture, _ = librosa.load(audio_file_path, sr=32000, mono=True)
@@ -60,36 +55,6 @@ def inference(audio_file_path: str, text: str):
         return 32000, np.round(sep_segment * 32767).astype(np.int16)
 
 
-def download_models():
-    CHECKPOINTS_DIR.mkdir(exist_ok=True)
-    success_file = CHECKPOINTS_DIR / "_SUCCESS"
-
-    models = (
-        (
-            "https://drive.google.com/file/d/1wQuXThdATXrkmkPM2sRGaNapJ4mTqmlY/view?usp=sharing",
-            MODEL_NAME,
-        ),
-        (
-            "https://drive.google.com/file/d/11oj8_tPG6SXgw5fIEsZ5HiWZnJOrvdhw/view?usp=sharing",
-            CHECKPOINTS_DIR / "music_speech_audioset_epoch_15_esc_89.98.pt",
-        ),
-    )
-
-    def download(models):
-        for model_url, model_path in models:
-            gdown.download(model_url, str(model_path), quiet=False, fuzzy=True)
-
-        success_file.touch()
-
-        global MODEL
-        MODEL = get_model()
-        button.update(value="Separate", interactive=True)
-
-    if not success_file.exists():
-        thread = Thread(target=download, args=[models])
-        thread.start()
-
-
 with gr.Blocks(title="AudioSep") as demo:
     gr.Markdown(description)
     with gr.Row():
@@ -100,16 +65,14 @@ with gr.Blocks(title="AudioSep") as demo:
             with gr.Column():
                 output_audio = gr.Audio(scale=10)
                 button = gr.Button(
-                    "Downloading the models...",
+                    "Separate",
                     variant="primary",
                     scale=2,
                     size="lg",
-                    interactive=False,
+                    interactive=True,
                 )
                 button.click(
                     fn=inference, inputs=[input_audio, text], outputs=[output_audio]
                 )
 
-download_models()
-
 demo.queue().launch(share=True)