Update app.py

app.py

@@ -126,10 +126,7 @@ class WhisperxModel:
 
 @spaces.GPU
 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")
@@ -171,6 +168,7 @@ def load_models(ssrspeech_model_name):
 
     return [
         gr.Accordion(),
+        whisper_model, align_model, ssrspeech_model,
         success_message
     ]
 
@@ -184,9 +182,7 @@ def get_transcribe_state(segments):
     }
 
 @spaces.GPU
-def transcribe(audio_path):
-    global transcribe_model
-    
+def transcribe(audio_path, transcribe_model):
     if transcribe_model is None:
         raise gr.Error("Transcription model not loaded")
 
@@ -201,8 +197,7 @@ def transcribe(audio_path):
 
 
 @spaces.GPU
-def align(segments, audio_path):
-    global align_model
+def align(segments, audio_path, align_model):
     if align_model is None:
         raise gr.Error("Align model not loaded")
     
@@ -232,9 +227,8 @@ def replace_numbers_with_words(sentence):
 @spaces.GPU
 def run(seed, sub_amount, ssrspeech_model_choice, codec_audio_sr, codec_sr, top_k, top_p, temperature,
         stop_repetition, kvcache, silence_tokens, aug_text, cfg_coef, prompt_length,
-        audio_path, original_transcript, transcript, mode):
+        audio_path, original_transcript, transcript, mode, whisper_model, align_model, ssrspeech_model):
     
-    global transcribe_model, align_model, ssrspeech_model
     aug_text = True if aug_text == 1 else False
     if ssrspeech_model is None:
         raise gr.Error("ssrspeech model not loaded")
@@ -254,16 +248,16 @@ 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, whisper_model)
     if language == 'zh':
         converter = opencc.OpenCC('t2s')
         orig_transcript = converter.convert(orig_transcript)
-        transcribe_state = align(traditional_to_simplified(segments), audio_path)
+        transcribe_state = align(traditional_to_simplified(segments), audio_path, align_model)
         transcribe_state['segments'] = traditional_to_simplified(transcribe_state['segments'])
     elif language == 'en':
         orig_transcript = orig_transcript.lower()
         target_transcript = target_transcript.lower()
-        transcribe_state = align(segments, audio_path)
+        transcribe_state = align(segments, audio_path, align_model)
     print(orig_transcript)
     print(target_transcript)
 
@@ -282,17 +276,17 @@ 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, whisper_model)
         
         if language == 'zh':
             converter = opencc.OpenCC('t2s')
             orig_transcript = converter.convert(orig_transcript)
-            transcribe_state = align(traditional_to_simplified(segments), audio_path)
+            transcribe_state = align(traditional_to_simplified(segments), audio_path, align_model)
             transcribe_state['segments'] = traditional_to_simplified(transcribe_state['segments'])
         elif language == 'en':
             orig_transcript = orig_transcript.lower()
             target_transcript = target_transcript.lower()
-            transcribe_state = align(segments, audio_path)
+            transcribe_state = align(segments, audio_path, align_model)
         print(orig_transcript)
         target_transcript_copy = target_transcript # for tts cut out
         if language == 'en':
@@ -370,14 +364,14 @@ 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, whisper_model)
         if language == 'zh':
-            transcribe_state = align(traditional_to_simplified(new_segments), audio_path)
+            transcribe_state = align(traditional_to_simplified(new_segments), audio_path, align_model)
             transcribe_state['segments'] = traditional_to_simplified(transcribe_state['segments'])
             tmp1 = transcribe_state['segments'][0]['words'][0]['word']
             tmp2 = target_transcript_copy
         elif language == 'en':
-            transcribe_state = align(new_segments, audio_path)
+            transcribe_state = align(new_segments, audio_path, align_model)
             tmp1 = transcribe_state['segments'][0]['words'][0]['word'].lower()
             tmp2 = target_transcript_copy.lower()
         if tmp1 == tmp2:
@@ -460,7 +454,7 @@ def get_app():
 
         load_models_btn.click(fn=load_models,
                             inputs=[ssrspeech_model_choice],
-                            outputs=[models_selector, success_output])
+                            outputs=[models_selector, whisper_model, align_model, ssrspeech_model, success_output])
 
         semgents = gr.State() # not used
         transcribe_btn.click(fn=transcribe,
@@ -474,7 +468,7 @@ def get_app():
                         top_k, top_p, temperature, stop_repetition, kvcache, silence_tokens, 
                         aug_text, cfg_coef, prompt_length, 
                         input_audio, original_transcript, transcript,
-                        mode
+                        mode, whisper_model, align_model, ssrspeech_model
                     ],
                     outputs=[output_audio, success_output])