Update app.py

app.py

@@ -8,9 +8,7 @@ from data.tokenizer import (
     AudioTokenizer,
     TextTokenizer,
 )
-from edit_utils_zh import parse_edit_zh
 from edit_utils_en import parse_edit_en
-from edit_utils_zh import parse_tts_zh
 from edit_utils_en import parse_tts_en
 from inference_scale import inference_one_sample
 import librosa
@@ -29,7 +27,6 @@ DEMO_PATH = os.getenv("DEMO_PATH", "./demo")
 TMP_PATH = os.getenv("TMP_PATH", "./demo/temp")
 MODELS_PATH = os.getenv("MODELS_PATH", "./pretrained_models")
 device = "cuda" if torch.cuda.is_available() else "cpu"
-transcribe_model, align_model, ssrspeech_model = None, None, None
 
 def get_random_string():
     return "".join(str(uuid.uuid4()).split("-"))
@@ -124,56 +121,36 @@ class WhisperxModel:
             segment['text'] = replace_numbers_with_words(segment['text'])
         return self.align_model.align(segments, audio_path)
 
-@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")
-        language = "en"
-        transcribe_model_name = "base.en"
-        
-    elif ssrspeech_model_name == "Mandarin":
-        ssrspeech_model_name = "Mandarin"
-        text_tokenizer = TextTokenizer(backend="espeak", language='cmn')
-        language = "zh"
-        transcribe_model_name = "base"
-
-    align_model = WhisperxAlignModel(language)
-    transcribe_model = WhisperxModel(transcribe_model_name, align_model, language)
-
-    ssrspeech_fn = f"{MODELS_PATH}/{ssrspeech_model_name}.pth"
-    if not os.path.exists(ssrspeech_fn):
-        os.system(f"wget https://huggingface.co/westbrook/SSR-Speech-{ssrspeech_model_name}/resolve/main/{ssrspeech_model_name}.pth -O " + ssrspeech_fn)
-    print(transcribe_model, align_model)
-    ckpt = torch.load(ssrspeech_fn)
-    model = ssr.SSR_Speech(ckpt["config"])
-    model.load_state_dict(ckpt["model"])
-    config = model.args
-    phn2num = ckpt["phn2num"]
-    model.to(device)
-
-    encodec_fn = f"{MODELS_PATH}/wmencodec.th"
-    if not os.path.exists(encodec_fn):
-        os.system(f"wget https://huggingface.co/westbrook/SSR-Speech-English/resolve/main/wmencodec.th -O " + encodec_fn)
-
-    ssrspeech_model = {
-        "config": config,
-        "phn2num": phn2num,
-        "model": model,
-        "text_tokenizer": text_tokenizer,
-        "audio_tokenizer": AudioTokenizer(signature=encodec_fn)
-    }
-    success_message = "<span style='color:green;'>Success: Models loading completed successfully!</span>"
-
-    return [
-        gr.Accordion(),
-        success_message
-    ]
-
+ssrspeech_model_name = "English"
+text_tokenizer = TextTokenizer(backend="espeak")
+language = "en"
+transcribe_model_name = "base.en"
+
+align_model = WhisperxAlignModel(language)
+transcribe_model = WhisperxModel(transcribe_model_name, align_model, language)
+
+ssrspeech_fn = f"{MODELS_PATH}/{ssrspeech_model_name}.pth"
+if not os.path.exists(ssrspeech_fn):
+    os.system(f"wget https://huggingface.co/westbrook/SSR-Speech-{ssrspeech_model_name}/resolve/main/{ssrspeech_model_name}.pth -O " + ssrspeech_fn)
+
+ckpt = torch.load(ssrspeech_fn)
+model = ssr.SSR_Speech(ckpt["config"])
+model.load_state_dict(ckpt["model"])
+config = model.args
+phn2num = ckpt["phn2num"]
+model.to(device)
+
+encodec_fn = f"{MODELS_PATH}/wmencodec.th"
+if not os.path.exists(encodec_fn):
+    os.system(f"wget https://huggingface.co/westbrook/SSR-Speech-English/resolve/main/wmencodec.th -O " + encodec_fn)
+
+ssrspeech_model = {
+    "config": config,
+    "phn2num": phn2num,
+    "model": model,
+    "text_tokenizer": text_tokenizer,
+    "audio_tokenizer": AudioTokenizer(signature=encodec_fn)
+}
 
 def get_transcribe_state(segments):
     transcript = " ".join([segment["text"] for segment in segments])
@@ -185,8 +162,6 @@ def get_transcribe_state(segments):
 
 @spaces.GPU
 def transcribe(audio_path):
-    global transcribe_model
-    
     if transcribe_model is None:
         raise gr.Error("Transcription model not loaded")
 
@@ -202,7 +177,6 @@ def transcribe(audio_path):
 
 @spaces.GPU
 def align(segments, audio_path):
-    global align_model
     if align_model is None:
         raise gr.Error("Align model not loaded")
     
@@ -230,21 +204,15 @@ def replace_numbers_with_words(sentence):
     return re.sub(r'\b\d+\b', replace_with_words, sentence) # Regular expression that matches numbers
 
 @spaces.GPU
-def run(seed, sub_amount, ssrspeech_model_choice, codec_audio_sr, codec_sr, top_k, top_p, temperature,
+def run(seed, sub_amount, 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):
     
-    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")
     
     seed_everything(seed)
-    
-    if ssrspeech_model_choice == "English":
-        language = "en"
-    elif ssrspeech_model_choice == "Mandarin":
-        language = "zh"
 
     # resample audio
     audio, _ = librosa.load(audio_path, sr=16000)
@@ -255,15 +223,9 @@ def run(seed, sub_amount, ssrspeech_model_choice, codec_audio_sr, codec_sr, top_
     orig_transcript = replace_numbers_with_words(original_transcript).replace("  ", " ").replace("  ", " ").replace("\n", " ")
 
     [orig_transcript, segments, _] = transcribe(audio_path)
-    if language == 'zh':
-        converter = opencc.OpenCC('t2s')
-        orig_transcript = converter.convert(orig_transcript)
-        transcribe_state = align(traditional_to_simplified(segments), audio_path)
-        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)
+    orig_transcript = orig_transcript.lower()
+    target_transcript = target_transcript.lower()
+    transcribe_state = align(segments, audio_path)
     print(orig_transcript)
     print(target_transcript)
 
@@ -284,26 +246,18 @@ def run(seed, sub_amount, ssrspeech_model_choice, codec_audio_sr, codec_sr, top_
         sf.write(audio_path, audio, 16000)
         [orig_transcript, segments, _] = transcribe(audio_path)
         
-        if language == 'zh':
-            converter = opencc.OpenCC('t2s')
-            orig_transcript = converter.convert(orig_transcript)
-            transcribe_state = align(traditional_to_simplified(segments), audio_path)
-            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)
+        
+        orig_transcript = orig_transcript.lower()
+        target_transcript = target_transcript.lower()
+        transcribe_state = align(segments, audio_path)
         print(orig_transcript)
         target_transcript_copy = target_transcript # for tts cut out
-        if language == 'en':
-            target_transcript_copy = target_transcript_copy.split(' ')[0]
-        elif language == 'zh':
-            target_transcript_copy = target_transcript_copy[0]
-        target_transcript = orig_transcript + ' ' + target_transcript if language == 'en' else orig_transcript + target_transcript
+        target_transcript_copy = target_transcript_copy.split(' ')[0]
+        target_transcript = orig_transcript + ' ' + target_transcript
         print(target_transcript)
 
     if mode == "Edit":
-        operations, orig_spans = parse_edit_en(orig_transcript, target_transcript) if language == 'en' else parse_edit_zh(orig_transcript, target_transcript)
+        operations, orig_spans = parse_edit_en(orig_transcript, target_transcript)
         print(operations)
         print("orig_spans: ", orig_spans)
         
@@ -371,15 +325,9 @@ def run(seed, sub_amount, ssrspeech_model_choice, codec_audio_sr, codec_sr, top_
     torchaudio.save(audio_path, new_audio, codec_audio_sr)
     if tts: # remove the start parts
         [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'])
-            tmp1 = transcribe_state['segments'][0]['words'][0]['word']
-            tmp2 = target_transcript_copy
-        elif language == 'en':
-            transcribe_state = align(new_segments, audio_path)
-            tmp1 = transcribe_state['segments'][0]['words'][0]['word'].lower()
-            tmp2 = target_transcript_copy.lower()
+        transcribe_state = align(new_segments, audio_path)
+        tmp1 = transcribe_state['segments'][0]['words'][0]['word'].lower()
+        tmp2 = target_transcript_copy.lower()
         if tmp1 == tmp2:
             offset = transcribe_state['segments'][0]['words'][0]['start']
         else:
@@ -406,15 +354,6 @@ demo_text = {
 
 def get_app():
     with gr.Blocks() as app:
-        with gr.Row():
-            with gr.Column(scale=2):
-                load_models_btn = gr.Button(value="Load models")
-            with gr.Column(scale=5):
-                with gr.Accordion("Select models", open=False) as models_selector:
-                    with gr.Row():
-                        ssrspeech_model_choice = gr.Radio(label="ssrspeech model", value="English",
-                                                        choices=["English", "Mandarin"])
-
         with gr.Row():
             with gr.Column(scale=2):
                 input_audio = gr.Audio(value=f"{DEMO_PATH}/5895_34622_000026_000002.wav", label="Input Audio", type="filepath", interactive=True)
@@ -458,10 +397,6 @@ def get_app():
 
         success_output = gr.HTML()
 
-        load_models_btn.click(fn=load_models,
-                            inputs=[ssrspeech_model_choice],
-                            outputs=[models_selector, success_output])
-
         semgents = gr.State() # not used
         transcribe_btn.click(fn=transcribe,
                             inputs=[input_audio],
@@ -469,7 +404,7 @@ def get_app():
         
         run_btn.click(fn=run,
                     inputs=[
-                        seed, sub_amount, ssrspeech_model_choice,
+                        seed, sub_amount,
                         codec_audio_sr, codec_sr,
                         top_k, top_p, temperature, stop_repetition, kvcache, silence_tokens, 
                         aug_text, cfg_coef, prompt_length,