Update app.py

app.py

@@ -20,6 +20,7 @@ import io
 import numpy as np
 import random
 import uuid
+import opencc
 import spaces
 import nltk
 nltk.download('punkt')
@@ -29,11 +30,20 @@ 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"
 whisper_model, align_model, ssrspeech_model = None, None, None
-_whitespace_re = re.compile(r"\s+")
 
 def get_random_string():
     return "".join(str(uuid.uuid4()).split("-"))
 
+def traditional_to_simplified(segments):
+    converter = opencc.OpenCC('t2s') 
+    seg_num = len(segments)
+    for i in range(seg_num):
+        words = segments[i]['words']
+        for j in range(len(words)):
+            segments[i]['words'][j]['word'] = converter.convert(segments[i]['words'][j]['word'])
+        segments[i]['text'] = converter.convert(segments[i]['text'])
+    return segments
+
 @spaces.GPU
 def seed_everything(seed):
     if seed != -1:
@@ -75,9 +85,9 @@ def get_mask_interval(transcribe_state, word_span):
     
 @spaces.GPU
 class WhisperxAlignModel:
-    def __init__(self):
+    def __init__(self, language):
         from whisperx import load_align_model
-        self.model, self.metadata = load_align_model(language_code="en", device=device)
+        self.model, self.metadata = load_align_model(language_code=language, device=device)
 
     def align(self, segments, audio_path):
         from whisperx import align, load_audio
@@ -86,12 +96,12 @@ class WhisperxAlignModel:
 
 @spaces.GPU
 class WhisperModel:
-    def __init__(self, model_name):
+    def __init__(self, model_name, language):
         from whisper import load_model
-        self.model = load_model(model_name, device)
+        self.model = load_model(model_name, device, language=language)
 
         from whisper.tokenizer import get_tokenizer
-        tokenizer = get_tokenizer(multilingual=False)
+        tokenizer = get_tokenizer(multilingual=False, language=language)
         self.supress_tokens = [-1] + [
             i
             for i in range(tokenizer.eot)
@@ -103,9 +113,9 @@ class WhisperModel:
 
 @spaces.GPU
 class WhisperxModel:
-    def __init__(self, model_name, align_model: WhisperxAlignModel):
+    def __init__(self, model_name, align_model, language):
         from whisperx import load_model
-        self.model = load_model(model_name, device, asr_options={"suppress_numerals": True, "max_new_tokens": None, "clip_timestamps": None, "hallucination_silence_threshold": None})
+        self.model = load_model(model_name, device, asr_options={"suppress_numerals": True, "max_new_tokens": None, "clip_timestamps": None, "hallucination_silence_threshold": None}, language=language)
         self.align_model = align_model
 
     def transcribe(self, audio_path):
@@ -158,71 +168,45 @@ def load_models(whisper_backend_name, whisper_model_name, alignment_model_name,
         "text_tokenizer": text_tokenizer,
         "audio_tokenizer": AudioTokenizer(signature=encodec_fn)
     }
-    return gr.Accordion()
+    success_message = "<span style='color:green;'>Success: Models loading completed successfully!</span>"
+
+    return [
+        gr.Accordion(),
+        success_message
+    ]
 
 
 def get_transcribe_state(segments):
-    words_info = [word_info for segment in segments for word_info in segment["words"]]
     transcript = " ".join([segment["text"] for segment in segments])
     transcript = transcript[1:] if transcript[0] == " " else transcript
     return {
         "segments": segments,
         "transcript": transcript,
-        "words_info": words_info,
-        "transcript_with_start_time": " ".join([f"{word['start']} {word['word']}" for word in words_info]),
-        "transcript_with_end_time": " ".join([f"{word['word']} {word['end']}" for word in words_info]),
-        "word_bounds": [f"{word['start']} {word['word']} {word['end']}" for word in words_info]
     }
 
 @spaces.GPU
-def transcribe(seed, audio_path):
+def transcribe(audio_path):
+    global transcribe_model
+    
     if transcribe_model is None:
         raise gr.Error("Transcription model not loaded")
-    seed_everything(seed)
 
     segments = transcribe_model.transcribe(audio_path)
     state = get_transcribe_state(segments)
     success_message = "<span style='color:green;'>Success: Transcribe completed successfully!</span>"
 
     return [
-        state["transcript"], state["transcript_with_start_time"], state["transcript_with_end_time"],
+        state["transcript"], state['segments'],
         state, success_message
     ]
 
-@spaces.GPU
-def align_segments(transcript, audio_path):
-    # from aeneas.executetask import ExecuteTask
-    # from aeneas.task import Task
-    # import json
-    # config_string = 'task_language=eng|os_task_file_format=json|is_text_type=plain'
-
-    # tmp_transcript_path = os.path.join(TMP_PATH, f"{get_random_string()}.txt")
-    # tmp_sync_map_path = os.path.join(TMP_PATH, f"{get_random_string()}.json")
-    # with open(tmp_transcript_path, "w") as f:
-    #     f.write(transcript)
-
-    # task = Task(config_string=config_string)
-    # task.audio_file_path_absolute = os.path.abspath(audio_path)
-    # task.text_file_path_absolute = os.path.abspath(tmp_transcript_path)
-    # task.sync_map_file_path_absolute = os.path.abspath(tmp_sync_map_path)
-    # ExecuteTask(task).execute()
-    # task.output_sync_map_file()
-
-    # with open(tmp_sync_map_path, "r") as f:
-    #     return json.load(f)
 
 @spaces.GPU
-def align(seed, transcript, audio_path):
+def align(segments, audio_path):
+    global align_model
     if align_model is None:
         raise gr.Error("Align model not loaded")
-    seed_everything(seed)
-    transcript = replace_numbers_with_words(transcript).replace("  ", " ").replace("  ", " ")
-    fragments = align_segments(transcript, audio_path)
-    segments = [{
-        "start": float(fragment["begin"]),
-        "end": float(fragment["end"]),
-        "text": " ".join(fragment["lines"])
-    } for fragment in fragments["fragments"]]
+    
     segments = align_model.align(segments, audio_path)
     state = get_transcribe_state(segments)
     success_message = "<span style='color:green;'>Success: Alignment completed successfully!</span>"
@@ -255,7 +239,8 @@ def run(seed, sub_amount, ssrspeech_model_choice, codec_audio_sr, codec_sr, top_
         stop_repetition, kvcache, silence_tokens, aug_text, cfg_coef,
         audio_path, transcribe_state, original_transcript, transcript,
         mode, selected_sentence, previous_audio_tensors):
-
+    
+    global ssrspeech_model
     aug_text = True if aug_text == 1 else False
     if ssrspeech_model is None:
         raise gr.Error("ssrspeech model not loaded")
@@ -387,35 +372,16 @@ def load_sentence(selected_sentence, codec_audio_sr, audio_tensors):
     selected_sentence_idx = int(selected_sentence[:colon_position])
     return get_output_audio([audio_tensors[selected_sentence_idx]], codec_audio_sr)
 
-smart_transcript_info = """
-If enabled, the target transcript will be constructed for you:</br>
- - In TTS and Long TTS mode just write the text you want to synthesize.</br>
- - In Edit mode just write the text to replace selected editing segment.</br>
-If disabled, you should write the target transcript yourself:</br>
- - In TTS mode write prompt transcript followed by generation transcript.</br>
- - In Long TTS select split by newline (<b>SENTENCE SPLIT WON'T WORK</b>) and start each line with a prompt transcript.</br>
- - In Edit mode write full prompt</br>
-"""
 
 demo_original_transcript = "Gwynplaine had, besides, for his work and for his feats of strength, round his neck and over his shoulders, an esclavine of leather."
 
 demo_text = {
     "TTS": {
-        "smart": "I cannot believe that the same model can also do text to speech synthesis too!",
         "regular": "Gwynplaine had, besides, for his work and for his feats of strength, I cannot believe that the same model can also do text to speech synthesis too!"
     },
     "Edit": {
-        "smart": "take over the stage for half an hour,",
         "regular": "Gwynplaine had, besides, for his work and for his feats of strength, take over the stage for half an hour, an esclavine of leather."
     },
-    "Long TTS": {
-        "smart": "You can run the model on a big text!\n"
-                 "Just write it line-by-line. Or sentence-by-sentence.\n"
-                 "If some sentences sound odd, just rerun the model on them, no need to generate the whole text again!",
-        "regular": "Gwynplaine had, besides, for his work and for his feats of strength, You can run the model on a big text!\n"
-                   "Gwynplaine had, besides, for his work and for his feats of strength, Just write it line-by-line. Or sentence-by-sentence.\n"
-                   "Gwynplaine had, besides, for his work and for his feats of strength, If some sentences sound odd, just rerun the model on them, no need to generate the whole text again!"
-    }
 }
 
 all_demo_texts = {vv for k, v in demo_text.items() for kk, vv in v.items()}
@@ -425,22 +391,6 @@ demo_words = ['0.069 Gwynplain 0.611', '0.671 had, 0.912', '0.952 besides, 1.414
 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 update_demo(mode, smart_transcript, edit_word_mode, transcript, edit_from_word, edit_to_word):
-    if transcript not in all_demo_texts:
-        return transcript, edit_from_word, edit_to_word
-
-    replace_half = edit_word_mode == "Replace half"
-    change_edit_from_word = edit_from_word == demo_words[2] or edit_from_word == demo_words[3]
-    change_edit_to_word = edit_to_word == demo_words[11] or edit_to_word == demo_words[12]
-    demo_edit_from_word_value = demo_words[2] if replace_half else demo_words[3]
-    demo_edit_to_word_value = demo_words[12] if replace_half else demo_words[11]
-    return [
-        demo_text[mode]["smart" if smart_transcript else "regular"],
-        demo_edit_from_word_value if change_edit_from_word else edit_from_word,
-        demo_edit_to_word_value if change_edit_to_word else edit_to_word,
-    ]
-
-
 def get_app():
     with gr.Blocks() as app:
         with gr.Row():
@@ -472,7 +422,7 @@ def get_app():
 
             with gr.Column(scale=3):
                 with gr.Group():
-                    transcript = gr.Textbox(label="Text", lines=7, value=demo_text["TTS"]["smart"])
+                    transcript = gr.Textbox(label="Text", lines=7, value=demo_text["TTS"]["regular"])
 
                     with gr.Row():
                         mode = gr.Radio(label="Mode", choices=["Edit", "TTS"], value="Edit")
@@ -515,7 +465,7 @@ def get_app():
 
         load_models_btn.click(fn=load_models,
                             inputs=[whisper_backend_choice, whisper_model_choice, align_model_choice, ssrspeech_model_choice],
-                            outputs=[models_selector])
+                            outputs=[models_selector, success_output])
 
 
         transcribe_btn.click(fn=transcribe,
@@ -573,4 +523,4 @@ if __name__ == "__main__":
     MODELS_PATH = args.models_path
 
     app = get_app()
-    app.queue().launch(share=args.share, server_port=args.port)
+    app.queue().launch(share=args.share, server_port=args.port)