Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import whisperx | |
| import torch | |
| import librosa | |
| import logging | |
| import os | |
| import time | |
| import numpy as np | |
| # Configure logging | |
| logging.basicConfig(level=logging.INFO) | |
| logger = logging.getLogger("whisperx_app") | |
| # Device setup (force CPU) | |
| device = "cpu" | |
| compute_type = "int8" | |
| torch.set_num_threads(os.cpu_count()) | |
| # Pre-load models | |
| 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='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 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() | |
| final_result = "" | |
| debug_log = [] | |
| try: | |
| # 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") | |
| # Get the preloaded model and determine batch size | |
| model = models[model_size] | |
| batch_size = 8 if model_size == "tiny" else 4 | |
| # 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"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) | |
| final_result = "Error occurred during transcription" | |
| debug_log.append(f"ERROR: {str(e)}") | |
| if debug: | |
| return final_result, "\n".join(debug_log) | |
| return final_result | |
| # Gradio Interface | |
| with gr.Blocks(title="WhisperX CPU Transcription") as demo: | |
| gr.Markdown("# WhisperX CPU Transcription with Word-Level Timestamps") | |
| with gr.Row(): | |
| with gr.Column(): | |
| audio_input = gr.Audio( | |
| label="Upload Audio File", | |
| type="filepath", | |
| sources=["upload", "microphone"], | |
| interactive=True, | |
| ) | |
| model_selector = gr.Dropdown( | |
| 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") | |
| with gr.Column(): | |
| output_text = gr.Textbox( | |
| label="Transcription Output", | |
| lines=20, | |
| placeholder="Transcription will appear here...", | |
| ) | |
| debug_output = gr.Textbox( | |
| label="Debug Information", | |
| lines=10, | |
| placeholder="Debug logs will appear here...", | |
| visible=False, | |
| ) | |
| # Toggle debug visibility | |
| def toggle_debug(debug_enabled): | |
| return gr.update(visible=debug_enabled) | |
| debug_checkbox.change( | |
| toggle_debug, | |
| inputs=[debug_checkbox], | |
| outputs=[debug_output] | |
| ) | |
| # Process transcription with the new pause_threshold parameter | |
| transcribe_btn.click( | |
| transcribe, | |
| 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() | |