AzALlN/whisperX
0
1import numpy as np2import streamlit as st3from constants import WHISPER_MODELS, language_dict4import streamlit as st5from utils import (6 translate_to_english,7 detect_language,8 write,9 read,10 get_key,11)12import whisperx as whisper13import json14import pandas as pd15from pydub import AudioSegment16import os17import uuid18 19if "btn1" not in st.session_state:20 st.session_state["btn1"] = False21if "btn2" not in st.session_state:22 st.session_state["btn2"] = False23 24 25class ByteEncoder(json.JSONEncoder):26 def default(self, obj):27 if isinstance(obj, bytes):28 return obj.hex()29 return json.JSONEncoder.default(self, obj)30 31 32def disable_btn2():33 st.session_state["btn2"] = True34 35 36def disable_btn1():37 st.session_state["btn1"] = True38 39 40st.set_page_config(page_title="Whisper-X", layout="wide")41import torch42 43if torch.cuda.is_available():44 device = "gpu"45else:46 device = "cpu"47input, output = st.columns(2, gap="medium")48with input:49 st.header("Input")50 audio_file = open("audio.wav", "rb")51 audio_bytes = audio_file.read()52 # st.markdown("""**sample audio**""", unsafe_allow_html=True)53 st.audio(audio_bytes, format="audio/wav")54 # st.markdown("""**your audio file**""", unsafe_allow_html=True)55 audio_uploaded = st.file_uploader(56 label="Upload your file",57 type=["mp3", "wav"],58 help="Your input file",59 # on_change=disable_btn2,60 # disabled=st.session_state["btn1"],61 )62 text_json = st.file_uploader(63 label="Aligned JSON",64 type=["json"],65 help="Your aligned json file (Only if you need to skip transcribe)",66 # disabled=st.session_state["btn2"],67 # on_change=disable_btn1,68 )69 # text_json = None70 71 # st.markdown("""**model**""", unsafe_allow_html=True)72 model_name = st.selectbox(73 label="Choose your model",74 options=WHISPER_MODELS,75 help="Choose a Whisper model.",76 )77 model_name = "base" if model_name == "" else model_name78 # st.markdown("**transcription**", unsafe_allow_html=True)79 transcription = st.selectbox(80 "transcription",81 options=["plain text", "srt", "vtt", "ass", "tsv"],82 help="Choose the format for the transcription",83 )84 translate = st.checkbox(85 "translate", help="Translate the text to English when set to True"86 )87 language = st.selectbox(88 label="language",89 options=list(language_dict.keys()) + list(language_dict.values()),90 help="Translate the text to English when set to True",91 )92 patience = st.number_input(93 label="patience",94 step=0.01,95 value=1.0,96 help="optional patience value to use in beam decoding, as in https://arxiv.org/abs/2204.05424, the default (1.0) is equivalent to conventional beam search",97 )98 temperature = st.number_input(99 label="temperature",100 step=0.01,101 value=1.0,102 help="temperature to use for sampling",103 )104 suppress_tokens = st.text_input(105 "suppress_tokens",106 value="-1",107 help="comma-separated list of token ids to suppress during sampling; '-1' will suppress most special characters except common punctuations",108 )109 initial_prompt = st.text_area(110 label="initial_prompt",111 help="optional text to provide as a prompt for the first window.",112 )113 condition_on_previous_text = st.checkbox(114 "condition_on_previous_text",115 help="if True, provide the previous output of the model as a prompt for the next window; disabling may make the text inconsistent across windows, but the model becomes less prone to getting stuck in a failure loop",116 )117 temperature_increment_on_fallback = st.number_input(118 label="temperature_increment_on_fallback",119 step=0.01,120 value=0.2,121 help="temperature to increase when falling back when the decoding fails to meet either of the thresholds below",122 )123 compression_ratio_threshold = st.number_input(124 label="compression_ratio_threshold",125 value=2.4,126 step=0.01,127 help="if the gzip compression ratio is higher than this value, treat the decoding as failed",128 )129 logprob_threshold = st.number_input(130 label="logprob_threshold",131 value=-1.0,132 step=0.01,133 help="if the average log probability is lower than this value, treat the decoding as failed",134 )135 no_speech_threshold = st.number_input(136 label="no_speech_threshold",137 value=0.6,138 step=0.01,139 help="if the probability of the <|nospeech|> token is higher than this value AND the decoding has failed due to `logprob_threshold`, consider the segment as silence",140 )141 if temperature_increment_on_fallback is not None:142 temperature = tuple(143 np.arange(temperature, 1.0 + 1e-6, temperature_increment_on_fallback)144 )145 else:146 temperature = [temperature]147 submit = st.button("Submit", type="primary")148with output:149 st.header("Output")150 151 segments_pre = st.empty()152 segments_post = st.empty()153 segments_post_json = st.empty()154 segments_post2 = st.empty()155 trans = st.empty()156 lang = st.empty()157 158 name = str(uuid.uuid1())159 if submit:160 if audio_uploaded is None:161 # st.audio(audio_bytes, format="audio/wav")162 audio_uploaded = audio_file163 if audio_uploaded is not None:164 if audio_uploaded.name.endswith(".wav"):165 temp = AudioSegment.from_wav(audio_uploaded)166 temp.export(f"{name}.wav")167 if audio_uploaded.name.endswith(".mp3"):168 169 try:170 171 172 temp = AudioSegment.from_file(audio_uploaded, format="mp3")173 temp.export(f"{name}.wav")174 except:175 176 temp = AudioSegment.from_file(audio_uploaded, format="mp4")177 temp.export(f"{name}.wav")178 if language == "":179 model = whisper.load_model(model_name)180 with st.spinner("Detecting language..."):181 detection = detect_language(f"{name}.wav", model)182 language = detection.get("detected_language")183 del model184 if len(language) > 2:185 language = get_key(language)186 187 if text_json is None:188 189 with st.spinner("Running ... "):190 decode = {"suppress_tokens": suppress_tokens, "beam_size": 5}191 model = whisper.load_model(model_name)192 with st.container():193 with st.spinner(f"Running with {model_name} model"):194 result = model.transcribe(195 f"{name}.wav",196 language=language,197 patience=patience,198 initial_prompt=initial_prompt,199 condition_on_previous_text=condition_on_previous_text,200 temperature=temperature,201 compression_ratio_threshold=compression_ratio_threshold,202 logprob_threshold=logprob_threshold,203 no_speech_threshold=no_speech_threshold,204 **decode,205 )206 207 if translate:208 result = translate_to_english(result, json=False)209 with open("transcription.json", "w") as f:210 json.dump(result["segments"], f, indent=4, cls=ByteEncoder)211 with st.spinner("Running alignment model ..."):212 model_a, metadata = whisper.load_align_model(213 language_code=result["language"], device=device214 )215 result_aligned = whisper.align(216 result["segments"],217 model_a,218 metadata,219 f"{name}.wav",220 device=device,221 )222 write(223 f"{name}.wav",224 dtype=transcription,225 result_aligned=result_aligned,226 )227 trans_text = read(f"{name}.wav", transcription)228 trans.text_area(229 "transcription", trans_text, height=None, max_chars=None, key=None230 )231 char_segments = []232 word_segments = []233 234 for x in range(len(result_aligned["segments"])):235 word_segments.append(236 {237 "word-segments": result_aligned["segments"][x][238 "word-segments"239 ]240 .fillna("")241 .to_dict(orient="records")242 }243 )244 char_segments.append(245 {246 "char-segments": result_aligned["segments"][x][247 "char-segments"248 ]249 .fillna("")250 .to_dict(orient="records")251 }252 )253 254 for x in range(len(result_aligned["segments"])):255 256 result_aligned["segments"][x]["word-segments"] = word_segments[x]257 result_aligned["segments"][x]["char-segments"] = char_segments[x]258 segments_pre.text_area(259 "Segments before alignment",260 result["segments"],261 height=None,262 max_chars=None,263 key=None,264 )265 segments_post.text_area(266 "Word Segments after alignment",267 result_aligned["word_segments"],268 height=None,269 max_chars=None,270 key=None,271 )272 segments_post2.text_area(273 "Segments after alignment",274 result_aligned["segments"],275 height=None,276 max_chars=None,277 key=None,278 )279 lang.text_input(280 "detected language", language_dict.get(language), disabled=True281 )282 os.remove(f"{name}.wav")283 if text_json is not None:284 with st.spinner("Running ... "):285 286 model = whisper.load_model(model_name)287 json_filname = str(uuid.uuid1())288 data = json.load(text_json)289 290 # Close the uploaded file291 text_json.close()292 293 # Write the JSON data to a new file294 with open(f"{json_filname}.json", "w") as outfile:295 json.dump(data, outfile)296 297 # with open("fold.json", "w", encoding="utf-8") as f:298 # json.dump(text_json, f)299 with open(f"{json_filname}.json", "r", encoding="utf-8") as f:300 cont = json.load(f)301 302 with st.spinner("Running alignment model ..."):303 model_a, metadata = whisper.load_align_model(304 language_code=language, device=device305 )306 result_aligned = whisper.align(307 cont,308 model_a,309 metadata,310 f"{name}.wav",311 device=device,312 )313 words_segments = result_aligned["word_segments"]314 write(315 f"{name}.wav",316 dtype=transcription,317 result_aligned=result_aligned,318 )319 trans_text = read(f"{name}.wav", transcription)320 char_segments = []321 word_segments = []322 323 for x in range(len(result_aligned["segments"])):324 word_segments.append(325 {326 "word-segments": result_aligned["segments"][x][327 "word-segments"328 ]329 .fillna("")330 .to_dict(orient="records")331 }332 )333 char_segments.append(334 {335 "char-segments": result_aligned["segments"][x][336 "char-segments"337 ]338 .fillna("")339 .to_dict(orient="records")340 }341 )342 343 for x in range(len(result_aligned["segments"])):344 345 result_aligned["segments"][x]["word-segments"] = word_segments[x]346 result_aligned["segments"][x]["char-segments"] = char_segments[x]347 trans.text_area(348 "transcription", trans_text, height=None, max_chars=None, key=None349 )350 segments_pre.text_area(351 "Segments before alignment",352 cont,353 height=None,354 max_chars=None,355 key=None,356 )357 358 segments_post.text_area(359 "Word Segments after alignment",360 result_aligned["word_segments"],361 height=None,362 max_chars=None,363 key=None,364 )365 366 segments_post2.text_area(367 "Segments after alignment",368 result_aligned["segments"],369 expanded=False,370 height=None,371 max_chars=None,372 key=None,373 )374 lang.text_input(375 "detected language", language_dict.get(language), disabled=True376 )377 os.remove(f"{name}.wav")378 os.remove(f"{json_filname}.json")