Update app.py

app.py

@@ -4,11 +4,9 @@ import time
 
 import gradio as gr
 import numpy as np
-import torch
 import yt_dlp as youtube_dl
 from gradio_client import Client
 from pyannote.audio import Pipeline
-from transformers.pipelines.audio_utils import ffmpeg_read
 
 
 YT_LENGTH_LIMIT_S = 36000  # limit to 1 hour YouTube files
@@ -189,11 +187,11 @@ def align(transcription, segments, group_by_speaker=True):
     return transcription
 
 
-def transcribe(audio_path, group_by_speaker=True):
+def transcribe(audio_path, task="transcribe", group_by_speaker=True, progress=gr.Progress()):
     # run Whisper JAX asynchronously using Gradio client (endpoint)
     job = client.submit(
         audio_path,
-        "transcribe",
+        task,
         True,
         api_name="/predict_1",
     )
@@ -211,11 +209,11 @@ def transcribe(audio_path, group_by_speaker=True):
     return transcription
 
 
-def transcribe_yt(yt_url, group_by_speaker=True):
+def transcribe_yt(yt_url, task="transcribe", group_by_speaker=True, progress=gr.Progress()):
     # run Whisper JAX asynchronously using Gradio client (endpoint)
     job = client.submit(
         yt_url,
-        "transcribe",
+        task,
         True,
         api_name="/predict_2",
     )
@@ -224,17 +222,8 @@ def transcribe_yt(yt_url, group_by_speaker=True):
     with tempfile.TemporaryDirectory() as tmpdirname:
         filepath = os.path.join(tmpdirname, "video.mp4")
         download_yt_audio(yt_url, filepath)
+        diarization = diarization_pipeline(filepath)
 
-        with open(filepath, "rb") as f:
-            inputs = f.read()
-
-    inputs = ffmpeg_read(inputs, SAMPLING_RATE)
-    inputs = torch.from_numpy(inputs).float()
-    inputs = inputs.unsqueeze(0)
-
-    diarization = diarization_pipeline(
-        {"waveform": inputs, "sample_rate": SAMPLING_RATE},
-    )
     segments = diarization.for_json()["content"]
 
     # only fetch the transcription result after performing diarization
@@ -257,6 +246,7 @@ microphone = gr.Interface(
     fn=transcribe,
     inputs=[
         gr.inputs.Audio(source="microphone", optional=True, type="filepath"),
+        gr.inputs.Radio(["transcribe", "translate"], label="Task", default="transcribe"),
         gr.inputs.Checkbox(default=True, label="Group by speaker"),
     ],
     outputs=[
@@ -272,6 +262,7 @@ audio_file = gr.Interface(
     fn=transcribe,
     inputs=[
         gr.inputs.Audio(source="upload", optional=True, label="Audio file", type="filepath"),
+        gr.inputs.Radio(["transcribe", "translate"], label="Task", default="transcribe"),
         gr.inputs.Checkbox(default=True, label="Group by speaker"),
     ],
     outputs=[
@@ -287,6 +278,7 @@ youtube = gr.Interface(
     fn=transcribe_yt,
     inputs=[
         gr.inputs.Textbox(lines=1, placeholder="Paste the URL to a YouTube video here", label="YouTube URL"),
+        gr.inputs.Radio(["transcribe", "translate"], label="Task", default="transcribe"),
         gr.inputs.Checkbox(default=True, label="Group by speaker"),
     ],
     outputs=[