JacobLinCool/forced-alignment

1import spaces2import gradio as gr3import json4import torch5import numpy as np6import librosa7from accelerate.utils.imports import is_cuda_available8from iso639 import iter_langs9from ctc_forced_aligner import (10    load_alignment_model,11    generate_emissions,12    preprocess_text,13    get_alignments,14    get_spans,15    postprocess_results,16)17 18device = "cuda" if is_cuda_available() else "cpu"19dtype = torch.float16 if is_cuda_available() else torch.float3220 21 22alignment_model, alignment_tokenizer = load_alignment_model(23    device,24    dtype=dtype,25)26 27 28def process_alignment(audio_waveform, text, language="eng"):29    print(f"{audio_waveform.shape=}, {text=}, {language=}")30    # Generate emissions31    emissions, stride = generate_emissions(32        alignment_model, audio_waveform, batch_size=1633    )34 35    # Preprocess text36    tokens_starred, text_starred = preprocess_text(37        text,38        romanize=True,39        language=language,40    )41 42    # Get alignments43    segments, scores, blank_id = get_alignments(44        emissions,45        tokens_starred,46        alignment_tokenizer,47    )48 49    # Get spans and word timestamps50    spans = get_spans(tokens_starred, segments, blank_id)51    word_timestamps = postprocess_results(text_starred, spans, stride, scores)52 53    return word_timestamps54 55 56def trim_audio(audio_array, sample_rate, word_timestamps):57    start_time = int(word_timestamps[0]["start"] * sample_rate)58    end_time = int(word_timestamps[-1]["end"] * sample_rate)59    print(f"{start_time=}, {end_time=}")60    trimmed_audio = audio_array[start_time:end_time]61    return (sample_rate, trimmed_audio)62 63 64def get_language_choices():65    return [f"{lang.pt3} - {lang.name}" for lang in iter_langs() if lang.pt3]66 67 68@spaces.GPU69def align(audio, text, language="eng - English"):70    # Extract the ISO 639-3 code from the selected language71    iso_code = language.split(" - ")[0]72 73    # Convert the input audio to 16kHz mono74    sample_rate, audio_array = audio75    audio_array = (76        audio_array.astype(np.float32) / 32768.077    )  # Convert to float32 and normalize78    print(f"{sample_rate=}, {audio_array.shape=}")79 80    if len(audio_array.shape) > 1:81        audio_array = audio_array.mean(axis=1)  # Convert to mono if stereo82    audio_array = librosa.resample(audio_array, orig_sr=sample_rate, target_sr=16000)83 84    # Convert to torch tensor and move to the correct device85    audio_waveform = torch.from_numpy(audio_array).to(device=device, dtype=dtype)86 87    # Process the alignment88    word_timestamps = process_alignment(audio_waveform, text, iso_code)89 90    # Trim the audio91    trimmed_audio = trim_audio(audio_array, 16000, word_timestamps)92 93    # Create JSON output94    output_json = {95        "input_text": text,96        "word_timestamps": word_timestamps,97        "language": language,98    }99 100    return trimmed_audio, json.dumps(output_json, indent=2)101 102 103@spaces.GPU104def align_result_only(audio, text, language="eng - English"):105    # Extract the ISO 639-3 code from the selected language106    iso_code = language.split(" - ")[0]107 108    # Convert the input audio to 16kHz mono109    sample_rate, audio_array = audio110    audio_array = (111        audio_array.astype(np.float32) / 32768.0112    )  # Convert to float32 and normalize113    print(f"{sample_rate=}, {audio_array.shape=}")114 115    if len(audio_array.shape) > 1:116        audio_array = audio_array.mean(axis=1)  # Convert to mono if stereo117    audio_array = librosa.resample(audio_array, orig_sr=sample_rate, target_sr=16000)118 119    # Convert to torch tensor and move to the correct device120    audio_waveform = torch.from_numpy(audio_array).to(device=device, dtype=dtype)121 122    # Process the alignment123    word_timestamps = process_alignment(audio_waveform, text, iso_code)124 125    # Create JSON output126    output_json = {127        "input_text": text,128        "word_timestamps": word_timestamps,129        "language": language,130    }131 132    return json.dumps(output_json, indent=2)133 134 135# Create Gradio blocks136with gr.Blocks() as demo:137    gr.Markdown("# Forced Alignment")138 139    gr.Markdown(140        """141    This tool aligns audio with text and provides word-level timestamps.142    143    ## How to use:144    1. Upload an audio file or record audio145    2. Enter the corresponding text146    3. Select the language147    4. Click 'Process' to get the alignment results148    """149    )150 151    with gr.Row():152        with gr.Column():153            audio_input = gr.Audio(label="Input Audio")154            text_input = gr.Textbox(label="Input Text")155            language_input = gr.Dropdown(156                choices=get_language_choices(), label="Language", value="eng - English"157            )158            submit_button = gr.Button(159                "Get Alignment and Trimmed Audio", variant="primary"160            )161            submit_button_result_only = gr.Button(162                "Get Alignment Only", variant="secondary"163            )164 165        with gr.Column():166            audio_output = gr.Audio(label="Trimmed Output Audio")167            json_output = gr.JSON(label="Alignment Results")168 169    submit_button.click(170        fn=align,171        inputs=[audio_input, text_input, language_input],172        outputs=[audio_output, json_output],173    )174 175    submit_button_result_only.click(176        fn=align_result_only,177        inputs=[audio_input, text_input, language_input],178        outputs=[json_output],179    )180 181    gr.Markdown("## Examples")182    gr.Examples(183        examples=[184            ["examples/example1.mp3", "我們搭上公車要回台北了", "zho - Chinese"],185            [186                "examples/example2.wav",187                "ON SATURDAY MORNINGS WHEN THE SODALITY MET IN THE CHAPEL TO RECITE THE LITTLE OFFICE HIS PLACE WAS A CUSHIONED KNEELING DESK AT THE RIGHT OF THE ALTAR FROM WHICH HE LED HIS WING OF BOYS THROUGH THE RESPONSES",188                "eng - English",189            ],190        ],191        inputs=[audio_input, text_input, language_input],192    )193 194# Launch the demo195if __name__ == "__main__":196    demo.launch()197