- Formatted Comments

- Added Custome Gradio Theme
- Replaced Storage Buckets with Model download

.gitignore

@@ -2,16 +2,17 @@
 
 
 # Virtual Environments
-.env
-.venv
-env/
-venv/
+.env
+.venv
+env/
+venv/
 ENV/
 
 
 # Ignore specific output directories
 speechbrain
 uploads
+pretrained_models
 
 # Audio formats from processing
 *.wav

README.md

@@ -4,7 +4,7 @@ emoji: 💻
 colorFrom: red
 colorTo: blue
 sdk: gradio
-sdk_version: 6.24.0
+sdk_version: 6.26.0
 python_version: 3.12.12
 app_file: app.py
 pinned: false
@@ -16,7 +16,7 @@ Check out the configuration reference at https://huggingface.co/docs/hub/spaces-
 
 # PureVoiceAI
 
-An advanced, end-to-end audio processing pipeline designed to improve the quality of transcription/note-taking. PureVoiceAI integrates machine learning models to perform audio source separation, speaker diarization, and speech transcription, turning noisy, overlapping audio into clean, attributed text.
+An advanced, end-to-end audio processing pipeline designed to improve the quality of transcription/note-taking. PureVoiceAI integrates machine learning models to perform audio source separation, voice activity detection, and speech transcription, turning noisy, overlapping audio into clean, attributed text.
 
 ## Key Features
 

app.py

@@ -1,37 +1,24 @@
+import glob
 import os
 import tempfile
+from pathlib import Path
+import gradio as gr
+import spaces
 import torch
 import torchaudio
-import glob
 import transformers
-import gradio as gr
-import spaces
-from pathlib import Path
-from speechbrain.inference.separation import SepformerSeparation
-from pyannote.audio import Pipeline
-from pyannote.core import Segment, Annotation
-from transformers import pipeline, AutoModelForSpeechSeq2Seq, AutoProcessor
+import numpy as np
 from huggingface_hub import sync_bucket
-
-
-# Process Outline
-# 1. Initialize ReSepformer model
-# 2. Initialize PyAnnote Huggingface model
-# 3. Initialize Whisper Huggingface model
-# 4. Recieve POST file from front end
-# 5. Save file to temporary location
-# 6. Loading file into RE-SepFormer and process
-# 7. Output files to temporary location
-# 8. Load each audio file into PyAnnote model
-# 9. Export each audio timings.
-# 10. Load separated audio into Whisper model
-# 11. Export
+from pyannote.audio import Pipeline
+from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
+from speechbrain.inference.separation import SepformerSeparation
+from pydub import AudioSegment
 
 # Disable TF32 for maximum precision (fixes the Pyannote warning)
 torch.backends.cuda.matmul.allow_tf32 = False
 torch.backends.cudnn.allow_tf32 = False
 
-# device  = "cuda:0" if torch.cuda.is_available() else "cpu" # GPU availability check
+device = "cuda:0" if torch.cuda.is_available() else "cpu"  # GPU availability check
 
 
 token = os.environ.get("HF_TOKEN")
@@ -40,19 +27,16 @@ temp_dir = tempfile.gettempdir()
 
 @spaces.GPU
 def process_audio(mixture):
-
     print(f"temp directory created: {temp_dir}")
 
     # Name of separated file folder
     sources = Path(temp_dir) / "sources"
 
-
     # Creates separated directory
-    Path.mkdir(sources,exist_ok=True)
-
+    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" )
+    # sync_bucket("hf://buckets/Jahmori-R/PureVoiceAISpace-storage", "./local")
 
     # ---------------------------------------------------------
     # INITIALIZE MODELS
@@ -60,19 +44,20 @@ def process_audio(mixture):
 
     print("Intializing models...")
 
-    #Loading Resepformer Model
+    # Loading Resepformer Model
     try:
-        separator = SepformerSeparation.from_hparams(source=f"./local/resepformer-librimix2spk",hparams_file="hyperparams.yaml",run_opts={"device": "cuda"})
-
         # separator = SepformerSeparation.from_hparams(
-        #     source="Jahmori-R/resepformer-librimix2spk",
-        #     savedir=f"./pretrained_models/resepformer-librimix2spk",
+        #     source=f"./local/resepformer-librimix2spk",
+        #     hparams_file="hyperparams.yaml",
         #     run_opts={"device": "cuda"},
-        #     # fetch_config=FetchConfig(token=True)
-        #     token=token,
-        #
         # )
-        #
+
+        separator = SepformerSeparation.from_hparams(
+            source="Jahmori-R/resepformer-librimix2spk",
+            savedir=f"./pretrained_models/resepformer-librimix2spk",
+            # run_opts={"device": "cuda"},
+            # fetch_config=FetchConfig(token=True)
+        )
 
         print("RE-SepFormer model successfully loaded")
     except Exception as e:
@@ -81,14 +66,16 @@ def process_audio(mixture):
 
     # #Loading Pyannote Model
     try:
-        diarize = Pipeline.from_pretrained("pyannote/speaker-diarization-community-1",token=token)
+        diarize = Pipeline.from_pretrained(
+            "pyannote/speaker-diarization-community-1", token=token
+        )
         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()
 
-    #Loading Whisper Model
+    # Loading Whisper Model
     # Prevents WhisperAI logging in cmd
     transformers.logging.set_verbosity_error()
     try:
@@ -98,14 +85,15 @@ def process_audio(mixture):
             transcription_model_id,
             dtype=torch_dtype,
             low_cpu_mem_usage=True,
-            use_safetensors=True
+            use_safetensors=True,
         )
 
         processor = AutoProcessor.from_pretrained(transcription_model_id)
 
         transcribe = pipeline(
             "automatic-speech-recognition",
-            model=model,tokenizer=processor.tokenizer,
+            model=model,
+            tokenizer=processor.tokenizer,
             feature_extractor=processor.feature_extractor,
             dtype=torch_dtype,
             device="cuda",
@@ -116,101 +104,142 @@ def process_audio(mixture):
         print(f"Failed to load Whisper model. Error: {e}")
         raise SystemExit()
 
-
     # ---------------------------------------------------------
     # LOAD & SEPARATE AUDIO
     # ---------------------------------------------------------
 
-    #Importing mixed audio source
+    # Importing mixed audio source
 
-    #Calling ReSepformer model onto audio source to separate
+    # Calling ReSepformer model onto audio source to separate
     assert separator is not None, "Failed to load the Sepformer model."
 
     print("Separating Audio Mixture")
-    est_sources  = separator.separate_file(mixture)
+    est_sources = separator.separate_file(mixture)
     print("Mixture separated")
 
-    torchaudio.save(f"{sources}/speaker_1.wav", est_sources[:,:,0].detach().cpu(), 16000)
-    torchaudio.save(f"{sources}/speaker_2.wav", est_sources[:,:,1].detach().cpu(), 16000)
-
-    # Saving each speaker as individual file
-    # for i in range(num_speakers):
-    #     speaker_audio = est_sources[:, :, i]
-    #     temp_pth = os.path.join(sources, f"speaker_{i}.flac")
-    #     torchaudio.save(temp_pth, speaker_audio.cpu(), 16000)
-    #     print(f"Saved: {temp_pth}")
+    torchaudio.save(
+        f"{sources}/speaker_1.wav", est_sources[:, :, 0].detach().cpu(), 16000
+    )
+    torchaudio.save(
+        f"{sources}/speaker_2.wav", est_sources[:, :, 1].detach().cpu(), 16000
+    )
 
     # ---------------------------------------------------------
     # DIARIZATION ON FILES
     # ---------------------------------------------------------
 
     # Dictionary to store speaker information
-    log = {}
-    full_timestamps = []
-
+    unfiltered_log = [[], []]
+    overlap_durations = []
 
     # Storing both speakers
-    file_paths = sorted(glob.glob(f"{sources}/speaker_*.wav")) # Glob finds all similar files
-
-
+    file_paths = sorted(
+        glob.glob(f"{sources}/speaker_*.wav")
+    )  # Glob finds all similar files
 
     spk_1 = file_paths[0]
     spk_2 = file_paths[1]
 
-
+    # Individual Speaker Processing
     for i, path in enumerate(file_paths):
-        speaker_key = f"speaker_{i}"
-        log[speaker_key] = {}
-
-        # waveform, sample_rate = torchaudio.load(path)
-
-        try:
-            # diarization = diarize({"waveform": waveform, "sample_rate": sample_rate, "speakers": 1})
-            diarization = diarize(path, num_speakers=1)
-            timestamps = []
-
-            for turn,_ in diarization.speaker_diarization:
-                timestamps.append({
-                    'start': float(round(turn.start,2)),
-                    "end": float(round(turn.end,2)),
-                })
-            log[speaker_key]["timestamps"] = timestamps
-        except Exception as e:
-            print(f"Pyannote error: {e}")
-            raise SystemExit()
-
-
-
-        # ---------------------------------------------------------
-        # TRANSCRIPTION
-        # ---------------------------------------------------------
-
-        print("Transcribing...")
-
-
-        transcriptions =(transcribe(path,return_timestamps=True))
-
-
-        log[speaker_key]["transcriptions"] = transcriptions
-
-        full_timestamps.append(list(transcriptions["chunks"]))
-
-
-
-
-
-
-
+        speaker_key = f"speaker_{i + 1}"
+
+        # Diarize speaker audio for time segments
+        print(f"Diarizing {speaker_key}")
+        diarization = diarize(path, num_speakers=1)
+
+        # Storing speaker voice activity timestamp to dictionary
+        for turn, _ in diarization.speaker_diarization:
+            unfiltered_log[i].append(
+                {
+                    "start": float(round(turn.start, 2)),
+                    "end": float(round(turn.end, 2)),
+                    "speaker": speaker_key,
+                }
+            )
+
+        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
+            start_ms = segment["start"] * 1000
+            end_ms = segment["end"] * 1000
+            chunk = audio_seg[start_ms:end_ms]
+
+            # Converting pydub AudioSegment to numpy array
+            samples = np.array(chunk.get_array_of_samples())
+            samples = samples.astype(np.float32) / 32768.0
+
+            transcription = transcribe(
+                samples,
+                generate_kwargs={
+                    "language": "english",
+                    "condition_on_prev_tokens": True,
+                },
+            )
+            segment["transcription"] = transcription
+
+    # 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"])
+
+    # Calculating the overlap between each segment to assign appropriate speaker to segment
+    # for log, log_next in zip(filtered_log, filtered_log[1:]):
+    #     current_start = log["start"]
+    #     current_end = log["end"]
+
+    #     next_start = log_next["start"]
+    #     next_end = log_next["end"]
+
+    #     overlap_start = max(current_start, next_start)
+    #     overlap_end = min(current_end, next_end)
+    #     overlap = max(0.0, overlap_end - overlap_start)
+    #     overlap_durations.append(overlap)
+
+    # Overlkap formula = Min(end 1, end 2) - Max(start 1, start 2)
+
+    print("Task Finished")
+    return (spk_1, spk_2, filtered_log)
+
+
+# css for interface modification
+css = """
+
+#centered-row {
+    justify-content: center;
+    display: flex;
+}
+"""
+
+with gr.Blocks(css=css) 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",
+            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)
 
-    return spk_1, spk_2, log, full_timestamps
+    with gr.Row(variant="panel"):
+        logs = gr.JSON(label="Conversation Log")
 
-demo = gr.Interface(
-    fn=process_audio,
-    inputs = gr.Audio(type="filepath", label = "Upload Mixture Audio"),
-    outputs = [gr.Audio(label="Speaker 1"),gr.Audio(label="Speaker 2"),gr.JSON(label="Logs"),gr.JSON(label="Whisper Logs")],
-    title="PureVoiceAI Project 2026", description = "Source separated end-to-end transcription processing pipeline"
-)
+    button.click(
+        fn=process_audio,
+        inputs=audio_input,
+        outputs=[speaker_1_output, speaker_2_output, logs],
+    )
 
 if __name__ == "__main__":
-    demo.launch(allowed_paths=[temp_dir])
+    demo.launch(allowed_paths=[temp_dir], theme=gr.themes.Soft())

requirements.txt

@@ -41,3 +41,4 @@ pytest==7.4.0
 ruff==v0.12.4
 yamllint==1.35.1
 spacy
+pydub