Update app.py

app.py

@@ -5,6 +5,7 @@ import librosa
 import logging
 import os
 import time
+import numpy as np
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -19,51 +20,105 @@ torch.set_num_threads(os.cpu_count())
 models = {
     "tiny": whisperx.load_model("tiny", device, compute_type=compute_type, vad_method='silero'),
     "base": whisperx.load_model("base", device, compute_type=compute_type, vad_method='silero'),
-    "small": whisperx.load_model("small", device, compute_type=compute_type, vad_method='silero'),
+    "small": whisperx.load_model("small", device, compute_type=compute_type, vad_method='siliro'),
     "large": whisperx.load_model("large", device, compute_type=compute_type, vad_method='silero'),
     "large-v2": whisperx.load_model("large-v2", device, compute_type=compute_type, vad_method='silero'),
     "large-v3": whisperx.load_model("large-v3", device, compute_type=compute_type, vad_method='silero'),
 }
 
-def transcribe(audio_file, model_size="base", debug=False):
+def split_audio_by_pause(audio, sr, pause_threshold, top_db=30):
+    """
+    Splits the audio into segments using librosa's non-silent detection.
+    Adjacent non-silent intervals are merged if the gap between them is less than the pause_threshold.
+    Returns a list of (start_sample, end_sample) tuples.
+    """
+    # Get non-silent intervals based on an amplitude threshold (in dB)
+    intervals = librosa.effects.split(audio, top_db=top_db)
+    if intervals.size == 0:
+        return [(0, len(audio))]
+    
+    merged_intervals = []
+    current_start, current_end = intervals[0]
+    
+    for start, end in intervals[1:]:
+        # Compute the gap duration (in seconds) between the current interval and the next one
+        gap_duration = (start - current_end) / sr
+        if gap_duration < pause_threshold:
+            # Merge intervals if gap is less than the threshold
+            current_end = end
+        else:
+            merged_intervals.append((current_start, current_end))
+            current_start, current_end = start, end
+    merged_intervals.append((current_start, current_end))
+    return merged_intervals
+
+def transcribe(audio_file, model_size="base", debug=False, pause_threshold=0.0):
     start_time = time.time()
-    result = ""
+    final_result = ""
     debug_log = []
     
     try:
-        # Load audio file
+        # Load audio file at 16kHz
         audio, sr = librosa.load(audio_file, sr=16000)
+        debug_log.append(f"Audio loaded: {len(audio)/sr:.2f} seconds long at {sr} Hz")
         
-        # Run inference
+        # Get the preloaded model and determine batch size
         model = models[model_size]
         batch_size = 8 if model_size == "tiny" else 4
-        transcript = model.transcribe(audio, batch_size=batch_size)
         
-        # Align whisper output
-        model_a, metadata = whisperx.load_align_model(
-            language_code=transcript["language"], device=device
-        )
-        transcript_aligned = whisperx.align(
-            transcript["segments"], model_a, metadata, audio, device
-        )
-        
-        # Format word-level output
-        for segment in transcript_aligned["segments"]:
-            for word in segment["words"]:
-                result += f"[{word['start']:5.2f}s-{word['end']:5.2f}s] {word['word']}\n"
+        # If pause_threshold > 0, split audio into segments based on silence pauses
+        if pause_threshold > 0:
+            segments = split_audio_by_pause(audio, sr, pause_threshold)
+            debug_log.append(f"Audio split into {len(segments)} segment(s) using a pause threshold of {pause_threshold}s")
+            # Process each audio segment individually
+            for seg_idx, (seg_start, seg_end) in enumerate(segments):
+                audio_segment = audio[seg_start:seg_end]
+                seg_duration = (seg_end - seg_start) / sr
+                debug_log.append(f"Segment {seg_idx+1}: start={seg_start/sr:.2f}s, duration={seg_duration:.2f}s")
+                
+                # Transcribe this segment
+                transcript = model.transcribe(audio_segment, batch_size=batch_size)
+                
+                # Load alignment model for the detected language in this segment
+                model_a, metadata = whisperx.load_align_model(
+                    language_code=transcript["language"], device=device
+                )
+                transcript_aligned = whisperx.align(
+                    transcript["segments"], model_a, metadata, audio_segment, device
+                )
+                
+                # Format word-level output with adjusted timestamps (adding segment offset)
+                for segment in transcript_aligned["segments"]:
+                    for word in segment["words"]:
+                        # Adjust start and end times by the segment's start time (in seconds)
+                        adjusted_start = word['start'] + seg_start/sr
+                        adjusted_end = word['end'] + seg_start/sr
+                        final_result += f"[{adjusted_start:5.2f}s-{adjusted_end:5.2f}s] {word['word']}\n"
+        else:
+            # Process the entire audio without splitting
+            transcript = model.transcribe(audio, batch_size=batch_size)
+            model_a, metadata = whisperx.load_align_model(
+                language_code=transcript["language"], device=device
+            )
+            transcript_aligned = whisperx.align(
+                transcript["segments"], model_a, metadata, audio, device
+            )
+            for segment in transcript_aligned["segments"]:
+                for word in segment["words"]:
+                    final_result += f"[{word['start']:5.2f}s-{word['end']:5.2f}s] {word['word']}\n"
         
-        debug_log.append(f"Processed in {time.time()-start_time:.2f}s")
         debug_log.append(f"Language detected: {transcript['language']}")
         debug_log.append(f"Batch size: {batch_size}")
+        debug_log.append(f"Processed in {time.time()-start_time:.2f}s")
         
     except Exception as e:
         logger.error("Error during transcription:", exc_info=True)
-        result = "Error occurred during transcription"
+        final_result = "Error occurred during transcription"
         debug_log.append(f"ERROR: {str(e)}")
     
     if debug:
-        return result, "\n".join(debug_log)
-    return result
+        return final_result, "\n".join(debug_log)
+    return final_result
 
 # Gradio Interface
 with gr.Blocks(title="WhisperX CPU Transcription") as demo:
@@ -78,11 +133,18 @@ with gr.Blocks(title="WhisperX CPU Transcription") as demo:
                 interactive=True,
             )
             model_selector = gr.Dropdown(
-                choices=models.keys(),
+                choices=list(models.keys()),
                 value="base",
                 label="Model Size",
                 interactive=True,
             )
+            # New input: pause threshold in seconds (set to 0 to disable splitting)
+            pause_threshold_slider = gr.Slider(
+                minimum=0, maximum=5, step=0.1, value=0,
+                label="Pause Threshold (seconds)",
+                interactive=True,
+                info="Set a pause duration threshold. Audio pauses longer than this will be used to split the audio into segments."
+            )
             debug_checkbox = gr.Checkbox(label="Enable Debug Mode", value=False)
             transcribe_btn = gr.Button("Transcribe", variant="primary")
             
@@ -109,13 +171,13 @@ with gr.Blocks(title="WhisperX CPU Transcription") as demo:
         outputs=[debug_output]
     )
     
-    # Process transcription
+    # Process transcription with the new pause_threshold parameter
     transcribe_btn.click(
         transcribe,
-        inputs=[audio_input, model_selector, debug_checkbox],
+        inputs=[audio_input, model_selector, debug_checkbox, pause_threshold_slider],
         outputs=[output_text, debug_output]
     )
 
 # Launch configuration
 if __name__ == "__main__":
-    demo.queue(max_size=4).launch()
+    demo.queue(max_size=4).launch()