Update app.py

app.py

@@ -1,12 +1,15 @@
-import gradio as gr
 import torch
-from transformers import pipeline
+import gradio as gr
 import yt_dlp as youtube_dl
-import os
-from scipy.io import wavfile
 import numpy as np
 from datasets import Dataset, Audio
+from scipy.io import wavfile
+
+from transformers import pipeline
+from transformers.pipelines.audio_utils import ffmpeg_read
+
 import tempfile
+import os
 import time
 import demucs
 
@@ -25,29 +28,33 @@ pipe = pipeline(
     device=device,
 )
 
-separator = demucs.pretrained.hdemucs()
+separator = demucs.api.Separator(model=DEMUCS_MODEL_NAME)
 
 def separate_vocal(path):
-    origin, separated = separator(path)
-    vocal_path = os.path.splitext(path)[0] + "_vocals.wav"
-    wavfile.write(vocal_path, separator.samplerate, separated[1].numpy())
-    return vocal_path
+    origin, separated = separator.separate_audio_file(path)
+    demucs.api.save_audio(separated["vocals"], path, samplerate=separator.samplerate)
+    return path
 
-def transcribe(inputs_path, task, use_demucs, dataset_name, oauth_token, progress=gr.Progress()):
+def transcribe(inputs_path, task, use_demucs, dataset_name, oauth_token: gr.OAuthToken | None, progress=gr.Progress()):
     if inputs_path is None:
         raise gr.Error("No audio file submitted! Please upload or record an audio file before submitting your request.")
-    if not dataset_name:
+    if dataset_name is None:
         raise gr.Error("No dataset name submitted! Please submit a dataset name. Should be in the format : <user>/<dataset_name> or <org>/<dataset_name>. Also accepts <dataset_name>, which will default to the namespace of the logged-in user.")
-    if oauth_token is None:
-        raise gr.Error("No OAuth token submitted! Please login to use this demo.")
 
+    if oauth_token is None:
+        gr.Warning("Make sure to click and login before using this demo.")
+        return [["transcripts will appear here"]], ""
+    
     total_step = 4
     current_step = 0
     
     current_step += 1
     progress((current_step, total_step), desc="Transcribe using Whisper.")
+
     sampling_rate, inputs = wavfile.read(inputs_path)
+
     out = pipe(inputs_path, batch_size=BATCH_SIZE, generate_kwargs={"task": task}, return_timestamps=True)
+    
     text = out["text"]
     
     current_step += 1
@@ -56,6 +63,7 @@ def transcribe(inputs_path, task, use_demucs, dataset_name, oauth_token, progres
 
     current_step += 1
     progress((current_step, total_step), desc="Create dataset.")
+    
     transcripts = []
     audios = []
     with tempfile.TemporaryDirectory() as tmpdirname:
@@ -75,30 +83,75 @@ def transcribe(inputs_path, task, use_demucs, dataset_name, oauth_token, progres
         
         current_step += 1
         progress((current_step, total_step), desc="Push dataset.")
-        dataset.push_to_hub(dataset_name, token=oauth_token)
+        dataset.push_to_hub(dataset_name, token=oauth_token.token if oauth_token else oauth_token)
         
     return [[transcript] for transcript in transcripts], text
 
-def yt_transcribe(yt_url, task, use_demucs, dataset_name, oauth_token, progress=gr.Progress()):
+def _return_yt_html_embed(yt_url):
+    video_id = yt_url.split("?v=")[-1]
+    HTML_str = (
+        f'<center> <iframe width="500" height="320" src="https://www.youtube.com/embed/{video_id}"> </iframe>'
+        " </center>"
+    )
+    return HTML_str
+
+def download_yt_audio(yt_url, filename):
+    info_loader = youtube_dl.YoutubeDL()
+    
+    try:
+        info = info_loader.extract_info(yt_url, download=False)
+    except youtube_dl.utils.DownloadError as err:
+        raise gr.Error(str(err))
+    
+    file_length = info["duration_string"]
+    file_h_m_s = file_length.split(":")
+    file_h_m_s = [int(sub_length) for sub_length in file_h_m_s]
+    
+    if len(file_h_m_s) == 1:
+        file_h_m_s.insert(0, 0)
+    if len(file_h_m_s) == 2:
+        file_h_m_s.insert(0, 0)
+    file_length_s = file_h_m_s[0] * 3600 + file_h_m_s[1] * 60 + file_h_m_s[2]
+    
+    if file_length_s > YT_LENGTH_LIMIT_S:
+        yt_length_limit_hms = time.strftime("%HH:%MM:%SS", time.gmtime(YT_LENGTH_LIMIT_S))
+        file_length_hms = time.strftime("%HH:%MM:%SS", time.gmtime(file_length_s))
+        raise gr.Error(f"Maximum YouTube length is {yt_length_limit_hms}, got {file_length_hms} YouTube video.")
+    
+    ydl_opts = {"outtmpl": filename, "format": "worstvideo[ext=mp4]+bestaudio[ext=m4a]/best[ext=mp4]/best"}
+    
+    with youtube_dl.YoutubeDL(ydl_opts) as ydl:
+        try:
+            ydl.download([yt_url])
+        except youtube_dl.utils.ExtractorError as err:
+            raise gr.Error(str(err))
+
+def yt_transcribe(yt_url, task, use_demucs, dataset_name, oauth_token: gr.OAuthToken | None, max_filesize=75.0, dataset_sampling_rate=24000,
+                  progress=gr.Progress()):
+    
     if yt_url is None:
-        raise gr.Error("No YouTube URL submitted! Please provide a working link.")
-    if not dataset_name:
+        raise gr.Error("No youtube link submitted! Please put a working link.")
+    if dataset_name is None:
         raise gr.Error("No dataset name submitted! Please submit a dataset name. Should be in the format : <user>/<dataset_name> or <org>/<dataset_name>. Also accepts <dataset_name>, which will default to the namespace of the logged-in user.")
-    if oauth_token is None:
-        raise gr.Error("No OAuth token submitted! Please login to use this demo.")
 
     total_step = 5
     current_step = 0
     
     html_embed_str = _return_yt_html_embed(yt_url)
 
+    if oauth_token is None:
+        gr.Warning("Make sure to click and login before using this demo.")
+        return html_embed_str, [["transcripts will appear here"]], ""
+        
     current_step += 1
     progress((current_step, total_step), desc="Load video.")
 
     with tempfile.TemporaryDirectory() as tmpdirname:
         filepath = os.path.join(tmpdirname, "video.mp4")
+        
         download_yt_audio(yt_url, filepath)
-        inputs_path = filepath
+        with open(filepath, "rb") as f:
+            inputs_path = f.read()
 
         inputs = ffmpeg_read(inputs_path, pipe.feature_extractor.sampling_rate)
         inputs = {"array": inputs, "sampling_rate": pipe.feature_extractor.sampling_rate}
@@ -106,6 +159,7 @@ def yt_transcribe(yt_url, task, use_demucs, dataset_name, oauth_token, progress=
         current_step += 1
         progress((current_step, total_step), desc="Transcribe using Whisper.")
         out = pipe(inputs, batch_size=BATCH_SIZE, generate_kwargs={"task": task}, return_timestamps=True)
+        
         text = out["text"]
             
         inputs = ffmpeg_read(inputs_path, dataset_sampling_rate)
@@ -116,6 +170,7 @@ def yt_transcribe(yt_url, task, use_demucs, dataset_name, oauth_token, progress=
 
     current_step += 1
     progress((current_step, total_step), desc="Create dataset.")
+    
     transcripts = []
     audios = []
     with tempfile.TemporaryDirectory() as tmpdirname:
@@ -135,12 +190,13 @@ def yt_transcribe(yt_url, task, use_demucs, dataset_name, oauth_token, progress=
         
         current_step += 1
         progress((current_step, total_step), desc="Push dataset.")
-        dataset.push_to_hub(dataset_name, token=oauth_token)
-        
+        dataset.push_to_hub(dataset_name, token=oauth_token.token if oauth_token else oauth_token)
+
     return html_embed_str, [[transcript] for transcript in transcripts], text
 
 def naive_postprocess_whisper_chunks(chunks, audio_array, sampling_rate, stop_chars=".!:;?", min_duration=5):
     min_duration = int(min_duration * sampling_rate)
+
     new_chunks = []
     while chunks:
         current_chunk = chunks.pop(0)
@@ -148,108 +204,72 @@ def naive_postprocess_whisper_chunks(chunks, audio_array, sampling_rate, stop_ch
         begin, end = int(begin * sampling_rate), int(end * sampling_rate)
         current_dur = end - begin
         text = current_chunk["text"]
-        chunk_to_concat = [audio_array[begin:end]]
-        while chunks and (text[-1] not in stop_chars or (current_dur < min_duration)):
-            ch = chunks.pop(0)
-            begin, end = ch["timestamp"]
-            begin, end = int(begin * sampling_rate), int(end * sampling_rate)
-            current_dur += end - begin
-            text = "".join([text, ch["text"]])
-            chunk_to_concat.append(audio_array[begin:end])
-        new_chunks.append({
-            "text": text,
-            "audio": np.concatenate(chunk_to_concat)
-        })
+        
+        chunk_to_concat = []
+        while chunks and (current_dur < min_duration or text[-1] not in stop_chars):
+            next_chunk = chunks.pop(0)
+            next_text = next_chunk["text"].strip()
+            next_begin, next_end = next_chunk["timestamp"]
+            next_begin, next_end = int(next_begin * sampling_rate), int(next_end * sampling_rate)
+            current_dur += next_end - next_begin
+            text += f" {next_text}"
+            end = next_end
+
+        new_chunks.append(
+            {
+                "audio": np.array(audio_array[begin:end]).astype(np.float32),
+                "text": text,
+            }
+        )
+
     return new_chunks
 
-def _return_yt_html_embed(yt_url):
-    video_id = yt_url.split("?v=")[-1]
-    HTML_str = (
-        f'<center> <iframe width="500" height="320" src="https://www.youtube.com/embed/{video_id}"> </iframe>'
-        " </center>"
+with gr.Blocks() as demo:
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown("### Audio or YouTube Video Transcription")
+            with gr.Row():
+                yt_textbox = gr.Textbox(label="YouTube link")
+                yt_button = gr.Button("Transcribe YouTube video")
+        with gr.Column():
+            gr.Markdown("### Upload or Record Audio")
+            local_audio_input = gr.Audio(type="filepath", label="Upload Audio")
+            local_button = gr.Button("Transcribe Local Audio")
+
+    task = gr.Radio(
+        ["transcribe", "translate"],
+        label="Task",
+        value="transcribe",
     )
-    return HTML_str
 
-def download_yt_audio(yt_url, filename):
-    info_loader = youtube_dl.YoutubeDL()
-    try:
-        info = info_loader.extract_info(yt_url, download=False)
-    except youtube_dl.utils.DownloadError as err:
-        raise gr.Error(str(err))
+    demucs_checkbox = gr.CheckboxGroup(["separate-audio"], label="Apply Demucs (Separate Vocal from Audio)")
+    dataset_name = gr.Textbox(label="Dataset name", placeholder="Dataset name to push to Hugging Face Hub")
     
-    file_length = info["duration_string"]
-    file_h_m_s = file_length.split(":")
-    file_h_m_s = [int(sub_length) for sub_length in file_h_m_s]
-    
-    if len(file_h_m_s) == 1:
-        file_h_m_s.insert(0, 0)
-    if len(file_h_m_s) == 2:
-        file_h_m_s.insert(0, 0)
-    file_length_s = file_h_m_s[0] * 3600 + file_h_m_s[1] * 60 + file_h_m_s[2]
-    
-    if file_length_s > YT_LENGTH_LIMIT_S:
-        yt_length_limit_hms = time.strftime("%HH:%MM:%SS", time.gmtime(YT_LENGTH_LIMIT_S))
-        file_length_hms = time.strftime("%HH:%MM:%SS", time.gmtime(file_length_s))
-        raise gr.Error(f"Maximum YouTube length is {yt_length_limit_hms}, got {file_length_hms} YouTube video.")
+    with gr.Row():
+        login_button = gr.Button("Login")
+        login_output = gr.Markdown()
     
-    ydl_opts = {"outtmpl": filename, "format": "worstvideo[ext=mp4]+bestaudio[ext=m4a]/best[ext=mp4]/best"}
-    
-    with youtube_dl.YoutubeDL(ydl_opts) as ydl:
-        try:
-            ydl.download([yt_url])
-        except youtube_dl.utils.ExtractorError as err:
-            raise gr.Error(str(err))
+    with gr.Row():
+        output_transcriptions = gr.Dataframe(headers=["Transcriptions"])
+        output_text = gr.Markdown()
 
-css = """
-#intro {
-  padding: 20px;
-  background-color: #f0f0f0;
-  margin-bottom: 10px;
-}
-#intro h1 {
-  font-size: 30px;
-}
-"""
-gr.config.css(css)
+    login_button.click(
+        fn=None,
+        inputs=None,
+        outputs=login_output,
+        _js="function() { return window.location = 'https://huggingface.co/login'; }",
+    )
 
-with gr.Blocks() as demo:
-    with gr.Tab("Local file"):
-        with gr.Row():
-            with gr.Column():
-                local_audio_input = gr.Audio(type="filepath", label="Upload Audio")
-                task_input = gr.Dropdown(choices=["transcribe", "translate"], value="transcribe", label="Task")
-                use_demucs_input = gr.Dropdown(choices=["do-nothing", "separate-audio"], value="do-nothing", label="Audio preprocessing")
-                dataset_name_input = gr.Textbox(label="Dataset name")
-                hf_token = gr.Textbox(label="HuggingFace Token")
-                submit_local_button = gr.Button("Transcribe")
-            with gr.Column():
-                local_output_text = gr.Dataframe(label="Transcripts")
-                local_output_full_text = gr.Textbox(label="Full Text")
-    
-        submit_local_button.click(
-            transcribe,
-            inputs=[local_audio_input, task_input, use_demucs_input, dataset_name_input, hf_token],
-            outputs=[local_output_text, local_output_full_text],
-        )
+    yt_button.click(
+        yt_transcribe,
+        inputs=[yt_textbox, task, demucs_checkbox, dataset_name, login_button],
+        outputs=[login_output, output_transcriptions, output_text],
+    )
 
-    with gr.Tab("YouTube video"):
-        with gr.Row():
-            with gr.Column():
-                yt_url_input = gr.Textbox(label="YouTube URL")
-                yt_task_input = gr.Dropdown(choices=["transcribe", "translate"], value="transcribe", label="Task")
-                yt_use_demucs_input = gr.Dropdown(choices=["do-nothing", "separate-audio"], value="do-nothing", label="Audio preprocessing")
-                yt_dataset_name_input = gr.Textbox(label="Dataset name")
-                yt_hf_token = gr.Textbox(label="HuggingFace Token")
-                submit_yt_button = gr.Button("Transcribe")
-            with gr.Column():
-                yt_html_embed_str = gr.HTML()
-                yt_output_text = gr.Dataframe(label="Transcripts")
-                yt_output_full_text = gr.Textbox(label="Full Text")
-    
-        submit_yt_button.click(
-            yt_transcribe,
-            inputs=[yt_url_input, yt_task_input, yt_use_demucs_input, yt_dataset_name_input, yt_hf_token],
-            outputs=[yt_html_embed_str, yt_output_text, yt_output_full_text],
-        )
+    local_button.click(
+        transcribe,
+        inputs=[local_audio_input, task, demucs_checkbox, dataset_name, login_button],
+        outputs=[output_transcriptions, output_text],
+    )
 
-demo.launch(share=True)
+demo.launch()