add yt functionality

.gitignore

@@ -4,4 +4,5 @@ __pycache__
 *.swp
 *.egg-info
 build
+yt-audio
 

app.py

@@ -2,10 +2,12 @@ import os
 from tsmnet import Stretcher
 import gradio as gr
 from gradio import processing_utils
-import torch
+# import torch
 import torchaudio
+import yt_dlp
 
 model_root = './weights'
+yt_dl_dir = 'yt-audio'
 available_models = ['speech', 'pop-music', 'classical-music']
 working_sr = 22050
 
@@ -15,13 +17,40 @@ def prepare_models():
         for weight in available_models
     }
 
+def download_yt_audio(url):
+    # purge outdated audio files (older than 1 days)
+    os.system(f'find {yt_dl_dir} -audio -mtime +1 -delete')
+
+    ydl_opts = {
+        'format': 'm4a/bestaudio/best',
+        'postprocessors': [{  # Extract audio using ffmpeg
+            'key': 'FFmpegExtractAudio',
+            'preferredcodec': 'wav',
+        }],
+        'outtmpl': f"{yt_dl_dir}/%(id)s.%(ext)s"
+    }
+
+    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
+        try:
+            ydl.cache.remove()
+            meta = ydl.extract_info(url, download=False)
+            audio_file = os.path.join(yt_dl_dir, meta['id'] + '.wav')
+            if not os.path.isfile(audio_file):
+                ydl.download(url)
+
+        except yt_dlp.DownloadError as error:
+            raise gr.Error(f'Failed to download from YouTube: {error}')
+
+    return audio_file
+
+
 def prepare_audio_file(rec, audio_file, yt_url):
     if rec is not None:
         return rec
     if audio_file is not None:
         return audio_file
     if yt_url != '':
-        pass
+        return download_yt_audio(yt_url)
     else:
         raise gr.Error('No audio found!')
 
@@ -45,6 +74,7 @@ def run(rec, audio_file, yt_url, speed, model, start_time, end_time):
 # @@@@@@@ Start of the program @@@@@@@@
 
 models = prepare_models()
+os.makedirs(yt_dl_dir, exist_ok=True)
 
 with gr.Blocks() as demo:
     gr.Markdown('# TSM-Net')
@@ -56,7 +86,7 @@ with gr.Blocks() as demo:
             with gr.Tab('From file'):
                 audio_file_box = gr.Audio(label='Audio sample', type='filepath')
             with gr.Tab('From YouTube'):
-                yt_url_box  = gr.Textbox(label='YouTube URL', placeholder='Under Construction', interactive=False)
+                yt_url_box  = gr.Textbox(label='YouTube URL', placeholder='https://youtu.be/q6EoRBvdVPQ')
 
             rec_box.change(lambda: [None] * 2, outputs=[audio_file_box, yt_url_box])
             audio_file_box.change(lambda: [None] * 2, outputs=[rec_box, yt_url_box])
@@ -68,7 +98,7 @@ with gr.Blocks() as demo:
                     # gr.Markdown('### Trim audio sample (sec)')
                     with gr.Row():
                         start_time_box = gr.Number(label='Start', value=0)
-                        end_time_box = gr.Number(label='End', value=20)
+                        end_time_box = gr.Number(label='End', value=60)
                 model_box = gr.Dropdown(label='Model weight', choices=available_models, value=available_models[0])
 
             submit_btn = gr.Button('Submit')
@@ -76,7 +106,7 @@ with gr.Blocks() as demo:
         with gr.Column():
             with gr.Accordion('Hint', open=False):
                 gr.Markdown('You can find more settings under the **Fine-grained settings**')
-                gr.Markdown('- Feeling slow? Try to adjust the start/end timestamp')
+                gr.Markdown('- Waiting too long? Try to adjust the start/end timestamp')
                 gr.Markdown('- Low audio quality? Try to switch to a proper model weight')
             outputs=gr.Audio(label='Output')