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

app.py

@@ -20,24 +20,40 @@ from pydub import AudioSegment
 torch.backends.cuda.matmul.allow_tf32 = False
 torch.backends.cudnn.allow_tf32 = False
 
-# Loading token from .env
+# Passing .env file as environment variable on machine
 load_dotenv()
 
-# Hugging Face authentication token grab
+# Hugging Face authentication token grab from .env
 token = os.getenv("HF_TOKEN")
 
+# GPU processing availability check
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Name of mixture folder
+mixtures = Path(os.getcwd())/"mixtures"
+
+# Name of separated file folder
+sources = Path(os.getcwd()) / "sources"
+
+# Creates mixture directory
+Path.mkdir(mixtures, exist_ok=True)
+
+# Creates separated directory
+Path.mkdir(sources, exist_ok=True)
 
 # Path checking for removing source directory every run
 def reset_dir():
+    remove_dir = [mixtures,sources]
     # 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")
+    for dir in remove_dir:
+        if os.path.exists(dir):
+            try:
+                # Remove directorty recursively
+                shutil.rmtree(sources)
+            except:
+                raise gr.Error("Reset Failed")
+    
+    gr.Info("System Reset",title="Notice")
 
     # To clear all Gradio Blocks
     return None, None, None, None, None
@@ -52,7 +68,8 @@ def chatFormat(transcription_log):
         start = i.get("start")
         end = i.get("end")
 
-        title_format = f"**{name} - [{start}s - {end}s]**"
+        # Title format for gradio chatbot block
+        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}})
@@ -70,14 +87,6 @@ def process_audio(mixture, progress=gr.Progress()):
 
     # Gradio processing alert
     gr.Info("Processing Has Started", title="Notice")
-    # GPU processing availability check
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-
-    # Name of separated file folder
-    sources = Path(os.getcwd()) / "sources"
-
-    # Creates separated directory
-    Path.mkdir(sources, exist_ok=True)
 
 # ---------------------------------------------------------
 # 1. INITIALIZE MODELS
@@ -146,13 +155,46 @@ def process_audio(mixture, progress=gr.Progress()):
     # Calling ReSepformer model onto audio source
     assert separator is not None, "Failed to load the Sepformer model."
 
+    progress(0.40,"Denoising Mixture...")
+
+    #Denoising mixture
+    sr, mixture_np = mixture
+
+    if len(mixture_np.shape) > 1:
+        # Swaps channels location with sample rate in numpy array
+        if mixture_np.shape[0] > mixture_np.shape[1]:
+            mixture_np = mixture_np.T
+
+    # If it has more than 2 channels take the first one
+    if len(mixture_np.shape) > 1 and mixture_np.shape[0] > 2:
+        mixture_np = mixture_np[0] 
+    
+    # Converting numpy to float32 for torchaudio 
+    if mixture_np.dtype != np.float32:
+
+        # Converting integer value to float between -1.0 and 1.0
+        if np.issubdtype(mixture_np.dtype, np.integer):
+            mixture_np = mixture_np.astype(np.float32) / np.iinfo(mixture_np.dtype).max
+        else:
+            mixture_np = mixture_np.astype(np.float32)
+
+    # Apply noise gate onto waveform
+    denoised_mixture_np = nr.reduce_noise(y=mixture_np,sr=sr, prop_decrease=0.5,n_fft=512)
+
+    # Convert numpy data back to torch for saving
+    denoised_mixture = torch.from_numpy(denoised_mixture_np)
+
+    # Saving new denoised mixture
+    torchaudio.save(f"{mixtures}/mixture.wav",denoised_mixture,sr)
+
     progress(0.50,"Separating Mixtures...")
+
     # Running mixed audio onto separation model
     print("Separating Audio Mixture")
-    est_sources = separator.separate_file(mixture)
+
+    est_sources = separator.separate_file(f"{mixtures}/mixture.wav")
     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
@@ -161,25 +203,6 @@ def process_audio(mixture, progress=gr.Progress()):
         f"{sources}/speaker_2.wav", est_sources[:, :, 1].detach().cpu(), 16000
     )
 
-    # 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
@@ -193,7 +216,7 @@ def process_audio(mixture, progress=gr.Progress()):
     # Storing both speakers
     file_paths = sorted(
         # Glob finds all similar files
-        glob.glob(f"{sources}/speaker_denoised_*.wav")
+        glob.glob(f"{sources}/speaker_*.wav")
     )
 
     spk_1 = file_paths[0]
@@ -209,8 +232,16 @@ def process_audio(mixture, progress=gr.Progress()):
         print(f"Diarizing {speaker_key}")
         diarization = diarize(path, num_speakers=1)
 
+        # Minimum timestamp duration in seconds
+        minimum_duration = 0.8
+
         # Storing speaker voice activity timestamps to dictionary
-        for turn, _ in diarization.speaker_diarization:
+        for turn, _ in diarization.exclusive_speaker_diarization:
+            duration = turn.end - turn.start
+            if duration <  minimum_duration:
+                continue
+        
+            # Adding timestamp to list
             unfiltered_log[i].append(
                 {
                     "start": float(round(turn.start, 2)),
@@ -257,6 +288,7 @@ def process_audio(mixture, progress=gr.Progress()):
 
     end_time = perf_counter()
 
+    # Inference time calculation
     inference_time = f"{end_time - start_time:.2f} seconds"
     return (spk_1, spk_2, chat, inference_time)
 
@@ -278,7 +310,7 @@ 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"
+            type="numpy", label="Input Audio Mixture | 2 Speakers"
         )
     with gr.Row(elem_id="centered-row"):
         process_button = gr.Button(
@@ -286,7 +318,7 @@ with gr.Blocks() as demo:
             variant="primary",
             scale=0,
         )
-        reset_button = gr.Button("Reset", variant="stop", interactive=False, scale=0)
+        reset_button = gr.Button("Restart", 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)
@@ -295,7 +327,7 @@ with gr.Blocks() as demo:
             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")
+        logs = gr.Chatbot(label="Conversation Log",layout="bubble")
 
     # Process Button Functionality
     process_button.click(
@@ -333,4 +365,4 @@ with gr.Blocks() as demo:
     )
 
 if __name__ == "__main__":
-    demo.launch(theme=gr.themes.Soft(), css=css)
+    demo.launch(theme=gr.Theme.from_hub("pseudolab/huggingface-korea-theme"), css=css)