Update app.py

app.py

@@ -12,9 +12,13 @@ from transformers.pipelines.audio_utils import ffmpeg_read
 import tempfile
 import os
 import time
+import demucs.api
+
+os.environ["GRADIO_TEMP_DIR"] = "/home/yoach/spaces/tmp"
 
 
 MODEL_NAME = "openai/whisper-large-v3"
+DEMUCS_MODEL_NAME = "htdemucs_ft"
 BATCH_SIZE = 8
 FILE_LIMIT_MB = 1000
 YT_LENGTH_LIMIT_S = 3600  # limit to 1 hour YouTube files
@@ -28,8 +32,23 @@ pipe = pipeline(
     device=device,
 )
 
+separator = demucs.api.Separator(model = DEMUCS_MODEL_NAME, )
+
+def separate_vocal(path):
+    origin, separated = separator.separate_audio_file(path)
+    demucs.api.save_audio(separated["vocals"], path, samplerate=separator.samplerate)
+    return path
+
+    
+
+# def separate_vocal(path, track_name, output_folder, demucs_model_name = "htdemucs_ft"):
+#     
+#   os.system(f"python3 -m demucs.separate --two-stems=vocals -n {demucs_model_name} {path} -o {output_folder}")
+#   
+#   return os.path.join(output_folder, demucs_model_name, track_name, "vocals.wav")
+
 
-def transcribe(inputs_path, task, dataset_name, oauth_token: gr.OAuthToken):
+def transcribe(inputs_path, task, use_demucs, dataset_name, oauth_token: gr.OAuthToken):
     if inputs_path is None:
         raise gr.Error("No audio file submitted! Please upload or record an audio file before submitting your request.")
     
@@ -39,18 +58,23 @@ def transcribe(inputs_path, task, dataset_name, oauth_token: gr.OAuthToken):
     
     text = out["text"]
     
-    chunks = naive_postprocess_whisper_chunks(out["chunks"])
-        
+    chunks = naive_postprocess_whisper_chunks(out["chunks"], inputs, sampling_rate)
+
     transcripts = []
     audios = []
     with tempfile.TemporaryDirectory() as tmpdirname:
         for i,chunk in enumerate(chunks):
-            begin, end = chunk["timestamp"]
-            begin, end = int(begin*sampling_rate), int(end*sampling_rate)
+            
             # TODO: make sure 1D or 2D?
-            arr = inputs[begin:end]
+            arr = chunk["audio"]
             path = os.path.join(tmpdirname, f"{i}.wav")
             wavfile.write(path, sampling_rate,  arr)
+            
+            if use_demucs == "separate-audio":
+                # use demucs tp separate vocals
+                print(f"Separating vocals #{i}")
+                path = separate_vocal(path)
+                
             audios.append(path)
             transcripts.append(chunk["text"])
             
@@ -102,7 +126,7 @@ def download_yt_audio(yt_url, filename):
             raise gr.Error(str(err))
 
 
-def yt_transcribe(yt_url, task, dataset_name, oauth_token: gr.OAuthToken, max_filesize=75.0, dataset_sampling_rate = 24000):
+def yt_transcribe(yt_url, task, use_demucs, dataset_name, oauth_token: gr.OAuthToken, max_filesize=75.0, dataset_sampling_rate = 24000):
     html_embed_str = _return_yt_html_embed(yt_url)
 
     with tempfile.TemporaryDirectory() as tmpdirname:
@@ -117,21 +141,26 @@ def yt_transcribe(yt_url, task, dataset_name, oauth_token: gr.OAuthToken, max_fi
     out = pipe(inputs, batch_size=BATCH_SIZE, generate_kwargs={"task": task}, return_timestamps=True)
     
     text = out["text"]
-    
-    chunks = naive_postprocess_whisper_chunks(out["chunks"])
-    
+        
     inputs = ffmpeg_read(inputs_path, dataset_sampling_rate)
     
+    chunks = naive_postprocess_whisper_chunks(out["chunks"], inputs, dataset_sampling_rate)
+
     transcripts = []
     audios = []
     with tempfile.TemporaryDirectory() as tmpdirname:
         for i,chunk in enumerate(chunks):
-            begin, end = chunk["timestamp"]
-            begin, end = int(begin*dataset_sampling_rate), int(end*dataset_sampling_rate)
+            
             # TODO: make sure 1D or 2D?
-            arr = inputs[begin:end]
+            arr = chunk["audio"]
             path = os.path.join(tmpdirname, f"{i}.wav")
             wavfile.write(path, dataset_sampling_rate,  arr)
+            
+            if use_demucs == "separate-audio":
+                # use demucs tp separate vocals
+                print(f"Separating vocals #{i}")
+                path = separate_vocal(path)
+                
             audios.append(path)
             transcripts.append(chunk["text"])
             
@@ -144,39 +173,57 @@ def yt_transcribe(yt_url, task, dataset_name, oauth_token: gr.OAuthToken, max_fi
     return html_embed_str, text
 
 
-def naive_postprocess_whisper_chunks(chunks, stop_chars = ".!:;?", min_duration = 5):
-    new_chunks = []
+def naive_postprocess_whisper_chunks(chunks, audio_array, sampling_rate,  stop_chars = ".!:;?", min_duration = 5):
+    # merge chunks as long as merged audio duration is lower than min_duration and that a stop character is not met
+    # return list of dictionnaries (text, audio)
+    # min duration is in seconds
     
+    min_duration = int(min_duration * sampling_rate)
+    
+    new_chunks = []
     while chunks:
         current_chunk = chunks.pop(0)
         begin, end = current_chunk["timestamp"]
+        begin, end = int(begin*sampling_rate), int(end*sampling_rate)
+        
+        current_dur = end-begin
+        
         text = current_chunk["text"]
         
-        while chunks and (text[-1] not in stop_chars or (end-begin<min_duration)):
+        print("new audio", begin/sampling_rate, end/sampling_rate)
+            
+        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)
-            end = ch["timestamp"][1]
+            
+            begin, end = ch["timestamp"]
+            begin, end = int(begin*sampling_rate), int(end*sampling_rate)
+            current_dur += end-begin
+            
             text = "".join([text, ch["text"]])
             
+            # TODO: add silence ?
+            chunk_to_concat.append(audio_array[begin:end])
+            
+            print("adding audio chunk", begin/sampling_rate, end/sampling_rate, len(audio_array[begin:end])/sampling_rate)
+            print(ch["timestamp"])
+            
         new_chunks.append({
             "text": text.strip(),
-            "timestamp": (begin, end),
+            "audio": np.concatenate(chunk_to_concat),
         })
-        print(f"LENGTH CHUNK #{len(new_chunks)}: {end-begin}s")
+        print(f"LENGTH CHUNK #{len(new_chunks)}: {current_dur/sampling_rate}s")
             
     return new_chunks
     
     
     
-    
-    
-
-demo = gr.Blocks()
-
 mf_transcribe = gr.Interface(
     fn=transcribe,
     inputs=[
         gr.Audio(type="filepath"),
         gr.Radio(["transcribe", "translate"], label="Task", value="transcribe"),
+        gr.Radio(["no-post-processing", "separate-audio"], label="Audio separation and cleaning (takes longer - use it if your samples are not cleaned (background noise and music))", value="separate-audio"),
         gr.Textbox(lines=1, placeholder="Place your new dataset name here", label="Dataset name"),
     ],
     outputs="text",
@@ -195,6 +242,7 @@ yt_transcribe = gr.Interface(
     inputs=[
         gr.Textbox(lines=1, placeholder="Paste the URL to a YouTube video here", label="YouTube URL"),
         gr.Radio(["transcribe", "translate"], label="Task", value="transcribe"),
+        gr.Radio(["no-post-processing", "separate-audio"], label="Audio separation and cleaning (takes longer - use it if your samples are not cleaned (background noise and music))", value="separate-audio"),
         gr.Textbox(lines=1, placeholder="Place your new dataset name here", label="Dataset name"),
     ],
     outputs=["html", "text"],
@@ -208,7 +256,7 @@ yt_transcribe = gr.Interface(
     allow_flagging="never",
 )
 
-with demo:
+with gr.Blocks() as demo:
     with gr.Row():
         gr.LoginButton()
         gr.LogoutButton()