Update app.py

app.py

@@ -125,10 +125,11 @@ class WhisperxModel:
         return self.align_model.align(segments, audio_path)
 
 @spaces.GPU
-def load_models(whisper_backend_name, ssrspeech_model_name):
+def load_models(ssrspeech_model_name):
     global transcribe_model, align_model, ssrspeech_model
     
     alignment_model_name = "whisperX"
+    whisper_backend_name = "whisperX"
     if ssrspeech_model_name == "English":
         ssrspeech_model_name = "English"
         text_tokenizer = TextTokenizer(backend="espeak")
@@ -141,16 +142,8 @@ def load_models(whisper_backend_name, ssrspeech_model_name):
         language = "zh"
         whisper_model_name = "base"
 
-    if alignment_model_name is not None:
-        align_model = WhisperxAlignModel(language)
-
-    if whisper_model_name is not None:
-        if whisper_backend_name == "whisper":
-            transcribe_model = WhisperModel(whisper_model_name, language)
-        else:
-            if align_model is None:
-                raise gr.Error("Align model required for whisperx backend")
-            transcribe_model = WhisperxModel(whisper_model_name, align_model, language)
+    align_model = WhisperxAlignModel(language)
+    transcribe_model = WhisperxModel(whisper_model_name, align_model, language)
 
     ssrspeech_fn = f"{MODELS_PATH}/{ssrspeech_model_name}.pth"
     if not os.path.exists(ssrspeech_fn):
@@ -261,7 +254,7 @@ def run(seed, sub_amount, ssrspeech_model_choice, codec_audio_sr, codec_sr, top_
     target_transcript = replace_numbers_with_words(transcript).replace("  ", " ").replace("  ", " ").replace("\n", " ")
     orig_transcript = replace_numbers_with_words(original_transcript).replace("  ", " ").replace("  ", " ").replace("\n", " ")
 
-    orig_transcript, segments = transcribe(audio_path)
+    [orig_transcript, segments, _] = transcribe(audio_path)
     if language == 'zh':
         converter = opencc.OpenCC('t2s')
         orig_transcript = converter.convert(orig_transcript)
@@ -289,7 +282,7 @@ def run(seed, sub_amount, ssrspeech_model_choice, codec_audio_sr, codec_sr, top_
 
         audio, _ = librosa.load(audio_path, sr=16000, duration=cut_length)
         sf.write(audio_path, audio, 16000)
-        orig_transcript, segments = transcribe(audio_path)
+        [orig_transcript, segments, _] = transcribe(audio_path)
         
         if language == 'zh':
             converter = opencc.OpenCC('t2s')
@@ -377,7 +370,7 @@ def run(seed, sub_amount, ssrspeech_model_choice, codec_audio_sr, codec_sr, top_
     new_audio = new_audio[0].cpu()
     torchaudio.save(audio_path, new_audio, codec_audio_sr)
     if tts: # remove the start parts
-        new_transcript, new_segments = transcribe(audio_path)
+        [new_transcript, new_segments, _] = transcribe(audio_path)
         if language == 'zh':
             transcribe_state = align(traditional_to_simplified(new_segments), audio_path)
             transcribe_state['segments'] = traditional_to_simplified(transcribe_state['segments'])
@@ -411,13 +404,6 @@ demo_text = {
     },
 }
 
-all_demo_texts = {vv for k, v in demo_text.items() for kk, vv in v.items()}
-
-demo_words = ['0.069 Gwynplain 0.611', '0.671 had, 0.912', '0.952 besides, 1.414', '1.494 for 1.634', '1.695 his 1.835', '1.915 work 2.136', '2.196 and 2.297', '2.337 for 2.517', '2.557 his 2.678', '2.758 feats 3.019', '3.079 of 3.139', '3.2 strength, 3.561', '4.022 round 4.263', '4.303 his 4.444', '4.524 neck 4.705', '4.745 and 4.825', '4.905 over 5.086', '5.146 his 5.266', '5.307 shoulders, 5.768', '6.23 an 6.33', '6.531 esclavine 7.133', '7.213 of 7.293', '7.353 leather. 7.614']
-
-demo_words_info = [{'word': 'Gwynplain', 'start': 0.069, 'end': 0.611, 'score': 0.833}, {'word': 'had,', 'start': 0.671, 'end': 0.912, 'score': 0.879}, {'word': 'besides,', 'start': 0.952, 'end': 1.414, 'score': 0.863}, {'word': 'for', 'start': 1.494, 'end': 1.634, 'score': 0.89}, {'word': 'his', 'start': 1.695, 'end': 1.835, 'score': 0.669}, {'word': 'work', 'start': 1.915, 'end': 2.136, 'score': 0.916}, {'word': 'and', 'start': 2.196, 'end': 2.297, 'score': 0.766}, {'word': 'for', 'start': 2.337, 'end': 2.517, 'score': 0.808}, {'word': 'his', 'start': 2.557, 'end': 2.678, 'score': 0.786}, {'word': 'feats', 'start': 2.758, 'end': 3.019, 'score': 0.97}, {'word': 'of', 'start': 3.079, 'end': 3.139, 'score': 0.752}, {'word': 'strength,', 'start': 3.2, 'end': 3.561, 'score': 0.742}, {'word': 'round', 'start': 4.022, 'end': 4.263, 'score': 0.916}, {'word': 'his', 'start': 4.303, 'end': 4.444, 'score': 0.666}, {'word': 'neck', 'start': 4.524, 'end': 4.705, 'score': 0.908}, {'word': 'and', 'start': 4.745, 'end': 4.825, 'score': 0.882}, {'word': 'over', 'start': 4.905, 'end': 5.086, 'score': 0.847}, {'word': 'his', 'start': 5.146, 'end': 5.266, 'score': 0.791}, {'word': 'shoulders,', 'start': 5.307, 'end': 5.768, 'score': 0.729}, {'word': 'an', 'start': 6.23, 'end': 6.33, 'score': 0.854}, {'word': 'esclavine', 'start': 6.531, 'end': 7.133, 'score': 0.803}, {'word': 'of', 'start': 7.213, 'end': 7.293, 'score': 0.772}, {'word': 'leather.', 'start': 7.353, 'end': 7.614, 'score': 0.896}]
-
-
 def get_app():
     with gr.Blocks() as app:
         with gr.Row():
@@ -428,7 +414,6 @@ def get_app():
                     with gr.Row():
                         ssrspeech_model_choice = gr.Radio(label="ssrspeech model", value="English",
                                                         choices=["English", "Mandarin"])
-                        whisper_backend_choice = gr.Radio(label="Whisper backend", value="whisperX", choices=["whisperX", "whisper"])
 
         with gr.Row():
             with gr.Column(scale=2):
@@ -440,7 +425,7 @@ def get_app():
 
             with gr.Column(scale=3):
                 with gr.Group():
-                    transcript = gr.Textbox(label="Text", lines=7, value=demo_text["TTS"]["regular"])
+                    transcript = gr.Textbox(label="Text", lines=7, value=demo_text["Edit"]["regular"])
 
                     with gr.Row():
                         mode = gr.Radio(label="Mode", choices=["Edit", "TTS"], value="Edit")
@@ -449,9 +434,6 @@ def get_app():
 
             with gr.Column(scale=2):
                 output_audio = gr.Audio(label="Output Audio")
-                with gr.Accordion("Inference transcript", open=False):
-                    inference_transcript = gr.Textbox(label="Inference transcript", lines=5, interactive=False,
-                                                    info="Inference was performed on this transcript.")
 
         with gr.Row():
             with gr.Accordion("Generation Parameters - change these if you are unhappy with the generation", open=False):
@@ -477,7 +459,7 @@ def get_app():
         success_output = gr.HTML()
 
         load_models_btn.click(fn=load_models,
-                            inputs=[whisper_backend_choice, ssrspeech_model_choice],
+                            inputs=[ssrspeech_model_choice],
                             outputs=[models_selector, success_output])
 
         semgents = gr.State() # not used