Sync from Jahmori-R/PureVoiceAI-HG via hub-sync

app.py

@@ -1,6 +1,8 @@
+from time import perf_counter
 import glob
 import os
-import tempfile
+from dotenv import load_dotenv
+import shutil
 from pathlib import Path
 import gradio as gr
 import spaces
@@ -8,7 +10,7 @@ import torch
 import torchaudio
 import transformers
 import numpy as np
-from huggingface_hub import sync_bucket
+import noisereduce as nr
 from pyannote.audio import Pipeline
 from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
 from speechbrain.inference.separation import SepformerSeparation
@@ -18,13 +20,52 @@ from pydub import AudioSegment
 torch.backends.cuda.matmul.allow_tf32 = False
 torch.backends.cudnn.allow_tf32 = False
 
-# Temporary Directory Intializing
+# Loading token from .env
+load_dotenv()
+
+# Path checking for removing source directory every run
+def reset_dir():
+    # Getting current source directory
+    sources = Path(__file__).parent / "sources"
+    if os.path.exists(sources):
+        try:
+            # Remove directorty recursively
+            shutil.rmtree(sources)
+            gr.Info("Process Reset")
+        except:
+            raise gr.Error("Reset Failed")
+
+    # To clear all Gradio Blocks
+    return None, None, None, None, None
+
+
+# Formatting the transcription log to Gradio chat display
+def chatFormat(transcription_log):
+    chat = []
+    for i in transcription_log:
+        name = i.get("speaker")
+        text = i.get("text")
+        start = i.get("start")
+        end = i.get("end")
+
+        title_format = f"**{name} - [{start}s - {end}s]**"
+        role = "user" if name == "speaker_2" else "assistant"
+
+        chat.append({"role": role, "content":text, "metadata":{"title":title_format}})
+
+    return chat
 
-temp_dir = tempfile.gettempdir()
-print(f"temp directory created: {temp_dir}")
 
 @spaces.GPU
-def process_audio(mixture):
+def process_audio(mixture, progress=gr.Progress()):
+
+    # Inference Time Processing
+    start_time = perf_counter()
+
+    progress(0, desc="Starting")
+
+    # Gradio processing alert
+    gr.Info("Processing Has Started", title="Notice")
     # GPU processing availability check
     device = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -32,20 +73,16 @@ def process_audio(mixture):
     token = os.environ.get("HF_TOKEN")
 
     # Name of separated file folder
-    sources = Path(temp_dir) / "sources"
+    sources = Path(os.getcwd()) / "sources"
 
     # Creates separated directory
     Path.mkdir(sources, exist_ok=True)
 
-    # Downloads RE-SepFormer model from huggingface buckets to space local directory
-    # sync_bucket("hf://buckets/Jahmori-R/PureVoiceAISpace-storage", "./local")
-
-
-
-    # ---------------------------------------------------------
-    # 1. INITIALIZE MODELS
-    # ---------------------------------------------------------
+# ---------------------------------------------------------
+# 1. INITIALIZE MODELS
+# ---------------------------------------------------------
 
+    progress(0.12, "Loading RE-SepFormer Model...")
     print("Intializing models...")
 
     # Loading Resepformer Model
@@ -57,9 +94,9 @@ def process_audio(mixture):
         )
         print("RE-SepFormer model successfully loaded")
     except Exception as e:
-        print(f"Failed to load RE-SepFormer model. Error: {e}")
-        raise SystemExit()
+        raise RuntimeError("Failed to load RE-SepFormer model") from e
 
+    progress(0.25, "Loading Pyannote Model...")
     # Loading Pyannote Model
 
     try:
@@ -70,11 +107,12 @@ def process_audio(mixture):
         diarize.to(torch.device("cuda"))
         print("Pyannote model successfully loaded")
     except Exception as e:
-        print(f"Failed to load Pyannote model. Error: {e}")
-        raise SystemExit()
+        raise RuntimeError("Failed to load Pyannote model") from e
+
+    progress(0.38, "Loading Whisper Model...")
 
     # Loading Whisper Model
-    transformers.logging.set_verbosity_error() # Prevents WhisperAI logging in cmd line
+    transformers.logging.set_verbosity_error()  # Prevents WhisperAI logging in cmd line
     try:
         transcription_model_id = "openai/whisper-large-v3-turbo"
         torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
@@ -98,22 +136,22 @@ def process_audio(mixture):
 
         print("Whisper model successfully loaded")
     except Exception as e:
-        print(f"Failed to load Whisper model. Error: {e}")
-        raise SystemExit()
-
-    # ---------------------------------------------------------
-    # LOAD & SEPARATE AUDIO
-    # ---------------------------------------------------------
+        raise RuntimeError("Failed to load Whisper model") from e
 
+# ---------------------------------------------------------
+# LOAD & SEPARATE AUDIO
+# ---------------------------------------------------------
 
     # Calling ReSepformer model onto audio source
     assert separator is not None, "Failed to load the Sepformer model."
 
+    progress(0.50,"Separating Mixtures...")
     # Running mixed audio onto separation model
     print("Separating Audio Mixture")
     est_sources = separator.separate_file(mixture)
     print("Mixture separated")
-
+    
+    print(type(est_sources))
     # Saving separated sources to directory
     torchaudio.save(
         f"{sources}/speaker_1.wav", est_sources[:, :, 0].detach().cpu(), 16000
@@ -122,22 +160,46 @@ def process_audio(mixture):
         f"{sources}/speaker_2.wav", est_sources[:, :, 1].detach().cpu(), 16000
     )
 
-    # ---------------------------------------------------------
-    # VAD + TRANSCRIPTION ON SOURCES
-    # ---------------------------------------------------------
+    # Applying noise gate to remaining artefacts in sources
+    spk1_wave_torch,sr1 = torchaudio.load(f"{sources}/speaker_1.wav")
+    spk2_wave_torch,sr2 = torchaudio.load(f"{sources}/speaker_2.wav")
+
+    # Convert torch to numpy for noisereduce
+    spk1_wave_np = spk1_wave_torch.numpy()
+    spk2_wave_np = spk2_wave_torch.numpy()
+
+    # Apply noise gate onto waveforms
+    spk1_denoised_np = nr.reduce_noise(y=spk1_wave_np,sr=sr1,n_fft=512)
+    spk2_denoised_np = nr.reduce_noise(y=spk2_wave_np,sr=sr2,n_fft=512)
+
+    # Convert numpy data back to torch for saving
+    spk1_denoised = torch.from_numpy(spk1_denoised_np)
+    spk2_denoised = torch.from_numpy(spk2_denoised_np)
+
+    # Saving new denoised sources
+    torchaudio.save(f"{sources}/speaker_denoised_1.wav",spk1_denoised,sr1)
+    torchaudio.save(f"{sources}/speaker_denoised_2.wav",spk2_denoised,sr2)
+
+# ---------------------------------------------------------
+# VAD + TRANSCRIPTION ON SOURCES
+# ---------------------------------------------------------
 
     # Dictionary to store speaker information
     unfiltered_log = [[], []]
 
+    progress(0.62, "Sorting Sources...")
+
     # Storing both speakers
     file_paths = sorted(
-      # Glob finds all similar files
-        glob.glob(f"{sources}/speaker_*.wav")
+        # Glob finds all similar files
+        glob.glob(f"{sources}/speaker_denoised_*.wav")
     )
 
     spk_1 = file_paths[0]
     spk_2 = file_paths[1]
 
+    progress(0.75, "VAD + Transcribing...")
+
     # Individual Speaker Processing
     for i, path in enumerate(file_paths):
         speaker_key = f"speaker_{i + 1}"
@@ -156,14 +218,13 @@ def process_audio(mixture):
                 }
             )
 
-
         # Creating audio segment from detected voice activity
         audio_seg = AudioSegment.from_file(path)
 
         print("Transcribing Chunks...")
         # Transcribing each voice activity detected by Pyannote
         for segment in unfiltered_log[i]:
-            # Converting segment timestamps to milliseconds for pydub
+            # Converting segment timestamps to milliseconds for pydub segmentation
             start_ms = segment["start"] * 1000
             end_ms = segment["end"] * 1000
             chunk = audio_seg[start_ms:end_ms]
@@ -172,26 +233,35 @@ def process_audio(mixture):
             samples = np.array(chunk.get_array_of_samples())
             samples = samples.astype(np.float32) / 32768.0
 
+            # Whisper Configutation
+            kwargs = {
+                "language": "english",  # Language transcription
+            }
             # Generating transcription given chunk segment
-            transcription = transcribe(
-                samples,
-                generate_kwargs={
-                    "language": "english",
-                    "condition_on_prev_tokens": True,
-                },
-            )
+            transcription = transcribe(samples, generate_kwargs=kwargs)
+
             # Appending each segment transcription to dictionary
-            segment["transcription"] = transcription
+            segment.update(transcription)
 
+    progress(0.88, "Sorting Speaker Logs....")
     # Combined both speaker logs
     combined_log = unfiltered_log[0] + unfiltered_log[1]
 
     # Sorting by start time of timestamps
     filtered_log = sorted(combined_log, key=lambda x: x["start"])
+    chat = chatFormat(filtered_log)
 
+    progress(1.00, "Task Finished")
     print("Task Finished")
-    return (spk_1, spk_2, filtered_log)
 
+    end_time = perf_counter()
+
+    inference_time = f"{end_time - start_time:.2f} seconds"
+    return (spk_1, spk_2, chat, inference_time)
+
+# ---------------------------------------------------------
+# GRADIO INTERFACE CUSTOMIZATION
+# ---------------------------------------------------------
 
 # Custom css for gradio interface
 css = """
@@ -203,31 +273,63 @@ css = """
 """
 
 # Gradio layout configuration
-with gr.Blocks(css=css) as demo:
+with gr.Blocks() as demo:
     gr.Markdown("PureVoiceAI")
     with gr.Row(variant="panel"):
         audio_input = gr.Audio(
             type="filepath", label="Input Audio Mixture | 2 Speakers"
         )
     with gr.Row(elem_id="centered-row"):
-        button = gr.Button(
+        process_button = gr.Button(
             "Process",
             variant="primary",
             scale=0,
         )
-
-    with gr.Row(variant="panel"):
-        speaker_1_output = gr.Audio(label="Speaker 1", interactive=False)
-        speaker_2_output = gr.Audio(label="Speaker 2", interactive=False)
-
-    with gr.Row(variant="panel"):
-        logs = gr.JSON(label="Conversation Log")
-
-    button.click(
+        reset_button = gr.Button("Reset", variant="stop", interactive=False, scale=0)
+    with gr.Row(variant="panel", visible=False) as source_row:
+        with gr.Column():
+            speaker_1_output = gr.Audio(label="Speaker 1", interactive=False)
+            speaker_2_output = gr.Audio(label="Speaker 2", interactive=False)
+        with gr.Column():
+            inference_time_log = gr.Textbox(label="Inference Time", interactive=False)
+
+    with gr.Row(variant="panel") as log_row:
+        logs = gr.Chatbot(label="Conversation Log",layout="panel")
+
+    # Process Button Functionality
+    process_button.click(
         fn=process_audio,
         inputs=audio_input,
-        outputs=[speaker_1_output, speaker_2_output, logs],
+        outputs=[speaker_1_output, speaker_2_output, logs, inference_time_log],
+    ).then(
+        fn=lambda: (
+            gr.Row(visible=True),
+            gr.Button(interactive=False),
+            gr.Button(interactive=True),
+        ),
+        inputs=None,
+        outputs=[source_row, process_button, reset_button],
+    )
+
+    # Reset Button Functionality
+    reset_button.click(
+        fn=reset_dir,
+        outputs=[
+            audio_input,
+            speaker_1_output,
+            speaker_2_output,
+            inference_time_log,
+            logs,
+        ],
+    ).then(
+        fn=lambda: (
+            gr.Row(visible=False),
+            gr.Button(interactive=True),
+            gr.Button(interactive=False),
+        ),
+        inputs=None,
+        outputs=[source_row, process_button, reset_button],
     )
 
 if __name__ == "__main__":
-        demo.launch(allowed_paths=[temp_dir], theme=gr.themes.Soft())
+    demo.launch(theme=gr.themes.Soft(), css=css)

requirements.txt

@@ -29,6 +29,7 @@ cuda-toolkit==13.0.2
 cycler==0.12.1
 cymem==2.0.13
 distlib==0.4.3
+dotenv==0.9.9
 einops==0.8.2
 fastapi==0.141.1
 filelock==3.32.7
@@ -73,6 +74,7 @@ murmurhash==1.0.15
 narwhals==2.26.0
 networkx==3.6.1
 nodeenv==1.10.0
+noisereduce==3.0.3
 numba==0.67.0
 numpy==2.5.3
 nvidia-cublas==13.1.0.3
@@ -127,6 +129,7 @@ pyparsing==3.3.2
 pytest==7.4.0
 python-dateutil==2.9.0.post0
 python-discovery==1.6.1
+python-dotenv==1.2.4
 python-multipart==0.0.32
 pytorch-lightning==2.6.6
 pytorch-metric-learning==2.9.0