feat: move spec plot to cpu task

app.py

@@ -36,7 +36,7 @@ def gen_task_id(location: str):
     return md5(video).hexdigest()
 
 
-def extract_audio(video: str) -> Tuple[str, str, str]:
+def extract_audio(video: str) -> Tuple[str, str]:
     task_id = gen_task_id(video)
     os.makedirs(path.join("task", task_id), exist_ok=True)
 
@@ -50,6 +50,14 @@ def extract_audio(video: str) -> Tuple[str, str, str]:
             f"ffmpeg -i {videodest} -vn -ar 48000 task/{task_id}/extracted_48k.wav"
         )
 
+    return (task_id, wav48k)
+
+
+def extract_audio_post(task_id: str) -> str:
+    wav48k = path.join("task", task_id, "extracted_48k.wav")
+    if not path.exists(wav48k):
+        raise gr.Error("Audio file not found")
+
     spec = path.join("task", task_id, "extracted_48k.png")
     if not path.exists(spec):
         y, sr = librosa.load(wav48k, sr=16000)
@@ -57,11 +65,11 @@ def extract_audio(video: str) -> Tuple[str, str, str]:
         fig.savefig(path.join("task", task_id, "extracted_48k.png"))
         plt.close(fig)
 
-    return (task_id, wav48k, spec)
+    return spec
 
 
 @zero()
-def extract_vocals(task_id: str) -> Tuple[str, str]:
+def extract_vocals(task_id: str) -> str:
     audio = path.join("task", task_id, "extracted_48k.wav")
     if not path.exists(audio):
         raise gr.Error("Audio file not found")
@@ -83,6 +91,14 @@ def extract_vocals(task_id: str) -> Tuple[str, str]:
             ]
         )
 
+    return vocals
+
+
+def extract_vocals_post(task_id: str) -> str:
+    vocals = path.join("task", task_id, "htdemucs", "extracted_48k", "vocals.wav")
+    if not path.exists(vocals):
+        raise gr.Error("Vocals file not found")
+
     spec = path.join("task", task_id, "vocals.png")
     if not path.exists(spec):
         y, sr = librosa.load(vocals, sr=16000)
@@ -90,7 +106,7 @@ def extract_vocals(task_id: str) -> Tuple[str, str]:
         fig.savefig(path.join("task", task_id, "vocals.png"))
         plt.close(fig)
 
-    return (vocals, spec)
+    return spec
 
 
 @zero(duration=60 * 2)
@@ -241,14 +257,11 @@ with gr.Blocks() as app:
         speaker_clip_zip = gr.File(label="Download Audio Segments")
 
     def preprocess(video: str):
-        task_id_val, extracted_audio_val, extracted_audio_spec_val = extract_audio(
-            video
-        )
+        task_id_val, extracted_audio_val = extract_audio(video)
         return {
             preprocess_output: gr.Column(visible=True),
             task_id: task_id_val,
             extracted_audio: extracted_audio_val,
-            extracted_audio_spec: extracted_audio_spec_val,
             preprocess_btn_label: gr.Markdown("", visible=False),
         }
 
@@ -259,26 +272,34 @@ with gr.Blocks() as app:
             preprocess_output,
             task_id,
             extracted_audio,
-            extracted_audio_spec,
             preprocess_btn_label,
         ],
         api_name="preprocess",
+    ).success(
+        fn=extract_audio_post,
+        inputs=[task_id],
+        outputs=[extracted_audio_spec],
+        api_name="preprocess-post",
     )
 
     def extract_vocals_fn(task_id: str):
-        vocals_val, vocals_spec_val = extract_vocals(task_id)
+        vocals_val = extract_vocals(task_id)
         return {
             extract_vocals_output: gr.Column(visible=True),
             vocals: vocals_val,
-            vocals_spec: vocals_spec_val,
             extract_vocals_btn_label: gr.Markdown("", visible=False),
         }
 
     extract_vocals_btn.click(
         fn=extract_vocals_fn,
         inputs=[task_id],
-        outputs=[extract_vocals_output, vocals, vocals_spec, extract_vocals_btn_label],
-        api_name="extract_vocals",
+        outputs=[extract_vocals_output, vocals, extract_vocals_btn_label],
+        api_name="extract-vocals",
+    ).success(
+        fn=extract_vocals_post,
+        inputs=[task_id],
+        outputs=[vocals_spec],
+        api_name="extract-vocals-post",
     )
 
     def diarize_fn(task_id: str):