update test

app.py

@@ -65,24 +65,19 @@ def _return_yt_html_embed(yt_url):
     return HTML_str
 
 
-def yt_transcribe(yt_url, return_timestamps=False):
-    ydl_opts = {
-        'format': 'bestaudio/best',
-        'outtmpl': 'audio.mp3',
-        'noplaylist': True
-    }
+@spaces.GPU
+def yt_transcribe(yt_url, task):
+    html_embed_str = _return_yt_html_embed(yt_url)
 
-    try:
-        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
-            ydl.download([yt_url])
-    except Exception as e:
-        return f"Error downloading audio: {str(e)}"
+    with tempfile.TemporaryDirectory() as tmpdirname:
+        filepath = os.path.join(tmpdirname, "audio.mp3")
+        download_yt_audio(yt_url, filepath)
 
-    if not os.path.exists("audio.mp3"):
-        return "Downloaded audio file not found."
+        inputs = ffmpeg_read(filepath, pipe.feature_extractor.sampling_rate)
+        inputs = {"array": inputs, "sampling_rate": pipe.feature_extractor.sampling_rate}
 
-    text = transcribe("audio.mp3", return_timestamps=return_timestamps)
-    return _return_yt_html_embed(yt_url), text
+    text = pipe(inputs, batch_size=BATCH_SIZE, generate_kwargs={"task": task}, return_timestamps=True)["text"]
+    return html_embed_str, text
 
 
 demo = gr.Blocks()
@@ -90,8 +85,8 @@ demo = gr.Blocks()
 mf_transcribe = gr.Interface(
     fn=transcribe,
     inputs=[
-        gr.components.Audio(sources=['upload', 'microphone'], type="filepath"),
-        gr.components.Checkbox(label="Return timestamps"),
+        gr.Audio(sources="microphone", type="filepath"),
+        gr.Radio(["transcribe", "translate"], label="Task", value="transcribe"),
     ],
     outputs="text",
     title="NB-Whisper Demo",
@@ -106,20 +101,20 @@ mf_transcribe = gr.Interface(
 yt_transcribe = gr.Interface(
     fn=yt_transcribe,
     inputs=[
-        gr.components.Textbox(lines=1, placeholder="Paste the URL to a YouTube video here", label="YouTube URL"),
-        gr.components.Checkbox(label="Return timestamps"),
+        gr.Textbox(lines=1, placeholder="Paste the URL to a YouTube video here", label="YouTube URL"),
+        gr.Radio(["transcribe", "translate"], label="Task", value="transcribe"),
     ],
-    examples=[["https://www.youtube.com/watch?v=mukeSSa5GKo"]],
     outputs=["html", "text"],
     title="Whisper Demo: Transcribe YouTube",
     description=(
         "Transcribe long-form YouTube videos with the click of a button! Demo uses the the fine-tuned checkpoint:"
-        f" [{MODEL_NAME}](https://huggingface.co/{MODEL_NAME}) and 🤗 Transformers to transcribe audio files of"
+        f" [{MODEL_NAME}](https://huggingface.co/{MODEL_NAME}) and 🤗 Transformers to transcribe video files of"
         " arbitrary length."
     ),
     allow_flagging="never",
 )
 
+
 with demo:
     gr.TabbedInterface([
         mf_transcribe,