Update app.py

app.py

@@ -12,8 +12,8 @@ from data.tokenizer import (
 )
 from edit_utils_en import parse_edit_en
 from edit_utils_en import parse_tts_en
-from edit_utils_zh import parse_edit_zh
-from edit_utils_zh import parse_tts_zh
+# from edit_utils_zh import parse_edit_zh
+# from edit_utils_zh import parse_tts_zh
 from inference_scale import inference_one_sample
 import librosa
 import soundfile as sf
@@ -33,31 +33,31 @@ MODELS_PATH = os.getenv("MODELS_PATH", "./pretrained_models")
 os.makedirs(MODELS_PATH, exist_ok=True)
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# if not os.path.exists(os.path.join(MODELS_PATH, "wmencodec.th")):
-#     # download wmencodec
-#     url = "https://huggingface.co/westbrook/SSR-Speech-English/resolve/main/wmencodec.th"
-#     filename = os.path.join(MODELS_PATH, "wmencodec.th")
-#     response = requests.get(url, stream=True)
-#     response.raise_for_status()
-#     with open(filename, "wb") as file:
-#         for chunk in response.iter_content(chunk_size=8192):
-#             file.write(chunk)
-#     print(f"File downloaded to: {filename}")
-# else:
-#     print("wmencodec model found")
-
-# if not os.path.exists(os.path.join(MODELS_PATH, "English.pth")):
-#     # download english model
-#     url = "https://huggingface.co/westbrook/SSR-Speech-English/resolve/main/English.pth"
-#     filename = os.path.join(MODELS_PATH, "English.pth")
-#     response = requests.get(url, stream=True)
-#     response.raise_for_status()
-#     with open(filename, "wb") as file:
-#         for chunk in response.iter_content(chunk_size=8192):
-#             file.write(chunk)
-#     print(f"File downloaded to: {filename}")
-# else:
-#     print("english model found")
+if not os.path.exists(os.path.join(MODELS_PATH, "wmencodec.th")):
+    # download wmencodec
+    url = "https://huggingface.co/westbrook/SSR-Speech-English/resolve/main/wmencodec.th"
+    filename = os.path.join(MODELS_PATH, "wmencodec.th")
+    response = requests.get(url, stream=True)
+    response.raise_for_status()
+    with open(filename, "wb") as file:
+        for chunk in response.iter_content(chunk_size=8192):
+            file.write(chunk)
+    print(f"File downloaded to: {filename}")
+else:
+    print("wmencodec model found")
+
+if not os.path.exists(os.path.join(MODELS_PATH, "English.pth")):
+    # download english model
+    url = "https://huggingface.co/westbrook/SSR-Speech-English/resolve/main/English.pth"
+    filename = os.path.join(MODELS_PATH, "English.pth")
+    response = requests.get(url, stream=True)
+    response.raise_for_status()
+    with open(filename, "wb") as file:
+        for chunk in response.iter_content(chunk_size=8192):
+            file.write(chunk)
+    print(f"File downloaded to: {filename}")
+else:
+    print("english model found")
 
 # if not os.path.exists(os.path.join(MODELS_PATH, "Mandarin.pth")):
 #     # download mandarin model
@@ -129,19 +129,16 @@ from whisperx import load_align_model, load_model, load_audio
 from whisperx import align as align_func
 
 # Load models
-# text_tokenizer_en = TextTokenizer(backend="espeak")
-text_tokenizer_zh = TextTokenizer(backend="espeak", language='cmn')
-
-text = "食品价格已基本都在一万到两万之间"
-print(text_tokenizer_zh(text))
+text_tokenizer_en = TextTokenizer(backend="espeak")
+# text_tokenizer_zh = TextTokenizer(backend="espeak", language='cmn')
 
-# ssrspeech_fn_en = f"{MODELS_PATH}/English.pth"
-# ckpt_en = torch.load(ssrspeech_fn_en)
-# model_en = ssr.SSR_Speech(ckpt_en["config"])
-# model_en.load_state_dict(ckpt_en["model"])
-# config_en = model_en.args
-# phn2num_en = ckpt_en["phn2num"]
-# model_en.to(device)
+ssrspeech_fn_en = f"{MODELS_PATH}/English.pth"
+ckpt_en = torch.load(ssrspeech_fn_en)
+model_en = ssr.SSR_Speech(ckpt_en["config"])
+model_en.load_state_dict(ckpt_en["model"])
+config_en = model_en.args
+phn2num_en = ckpt_en["phn2num"]
+model_en.to(device)
 
 # ssrspeech_fn_zh = f"{MODELS_PATH}/Mandarin.pth"
 # ckpt_zh = torch.load(ssrspeech_fn_zh)
@@ -151,15 +148,15 @@ print(text_tokenizer_zh(text))
 # phn2num_zh = ckpt_zh["phn2num"]
 # model_zh.to(device)
 
-# encodec_fn = f"{MODELS_PATH}/wmencodec.th"
+encodec_fn = f"{MODELS_PATH}/wmencodec.th"
 
-# ssrspeech_model_en = {
-#     "config": config_en,
-#     "phn2num": phn2num_en,
-#     "model": model_en,
-#     "text_tokenizer": text_tokenizer_en,
-#     "audio_tokenizer": AudioTokenizer(signature=encodec_fn)
-# }
+ssrspeech_model_en = {
+    "config": config_en,
+    "phn2num": phn2num_en,
+    "model": model_en,
+    "text_tokenizer": text_tokenizer_en,
+    "audio_tokenizer": AudioTokenizer(signature=encodec_fn)
+}
 
 # ssrspeech_model_zh = {
 #     "config": config_zh,
@@ -195,21 +192,21 @@ def transcribe_en(audio_path):
         state, success_message
     ]
 
-@spaces.GPU
-def transcribe_zh(audio_path):
-    language = "zh"
-    transcribe_model_name = "medium"
-    transcribe_model = load_model(transcribe_model_name, device, asr_options={"suppress_numerals": True, "max_new_tokens": None, "clip_timestamps": None, "hallucination_silence_threshold": None}, language=language)
-    segments = transcribe_model.transcribe(audio_path, batch_size=8)["segments"]
-    _, segments = align_zh(segments, audio_path)
-    state = get_transcribe_state(segments)
-    success_message = "<span style='color:green;'>Success: Transcribe completed successfully!</span>"
-    converter = opencc.OpenCC('t2s')
-    state["transcript"] = converter.convert(state["transcript"])
-    return [
-        state["transcript"], state['segments'],
-        state, success_message
-    ]
+# @spaces.GPU
+# def transcribe_zh(audio_path):
+#     language = "zh"
+#     transcribe_model_name = "medium"
+#     transcribe_model = load_model(transcribe_model_name, device, asr_options={"suppress_numerals": True, "max_new_tokens": None, "clip_timestamps": None, "hallucination_silence_threshold": None}, language=language)
+#     segments = transcribe_model.transcribe(audio_path, batch_size=8)["segments"]
+#     _, segments = align_zh(segments, audio_path)
+#     state = get_transcribe_state(segments)
+#     success_message = "<span style='color:green;'>Success: Transcribe completed successfully!</span>"
+#     converter = opencc.OpenCC('t2s')
+#     state["transcript"] = converter.convert(state["transcript"])
+#     return [
+#         state["transcript"], state['segments'],
+#         state, success_message
+#     ]
 
 @spaces.GPU
 def align_en(segments, audio_path):
@@ -222,15 +219,15 @@ def align_en(segments, audio_path):
     return state, segments
 
 
-@spaces.GPU
-def align_zh(segments, audio_path):
-    language = "zh"
-    align_model, metadata = load_align_model(language_code=language, device=device)
-    audio = load_audio(audio_path)
-    segments = align_func(segments, align_model, metadata, audio, device, return_char_alignments=False)["segments"]
-    state = get_transcribe_state(segments)
+# @spaces.GPU
+# def align_zh(segments, audio_path):
+#     language = "zh"
+#     align_model, metadata = load_align_model(language_code=language, device=device)
+#     audio = load_audio(audio_path)
+#     segments = align_func(segments, align_model, metadata, audio, device, return_char_alignments=False)["segments"]
+#     state = get_transcribe_state(segments)
 
-    return state, segments
+#     return state, segments
 
 
 def get_output_audio(audio_tensors, codec_audio_sr):
@@ -445,210 +442,210 @@ def run_tts_en(seed, sub_amount, aug_text, cfg_coef, cfg_stride, prompt_length,
     return output_audio, success_message
 
 
-@spaces.GPU
-def run_edit_zh(seed, sub_amount, aug_text, cfg_coef, cfg_stride, prompt_length,
-        audio_path, original_transcript, transcript):
+# @spaces.GPU
+# def run_edit_zh(seed, sub_amount, aug_text, cfg_coef, cfg_stride, prompt_length,
+#         audio_path, original_transcript, transcript):
     
-    codec_audio_sr = 16000
-    codec_sr = 50
-    top_k = 0
-    top_p = 0.8
-    temperature = 1
-    kvcache = 1
-    stop_repetition = 2
+#     codec_audio_sr = 16000
+#     codec_sr = 50
+#     top_k = 0
+#     top_p = 0.8
+#     temperature = 1
+#     kvcache = 1
+#     stop_repetition = 2
     
-    aug_text = True if aug_text == 1 else False
+#     aug_text = True if aug_text == 1 else False
     
-    seed_everything(seed)
+#     seed_everything(seed)
 
-    # resample audio
-    audio, _ = librosa.load(audio_path, sr=16000)
-    sf.write(audio_path, audio, 16000)
+#     # resample audio
+#     audio, _ = librosa.load(audio_path, sr=16000)
+#     sf.write(audio_path, audio, 16000)
     
-    # text normalization
-    target_transcript = transcript.replace("  ", " ").replace("  ", " ").replace("\n", " ")
-    orig_transcript = original_transcript.replace("  ", " ").replace("  ", " ").replace("\n", " ")
+#     # text normalization
+#     target_transcript = transcript.replace("  ", " ").replace("  ", " ").replace("\n", " ")
+#     orig_transcript = original_transcript.replace("  ", " ").replace("  ", " ").replace("\n", " ")
 
-    [orig_transcript, segments, _, _] = transcribe_zh(audio_path)
+    # [orig_transcript, segments, _, _] = transcribe_zh(audio_path)
 
-    converter = opencc.OpenCC('t2s')
-    orig_transcript = converter.convert(orig_transcript)
-    transcribe_state,_ = align_zh(traditional_to_simplified(segments), audio_path)
-    transcribe_state['segments'] = traditional_to_simplified(transcribe_state['segments'])
+    # converter = opencc.OpenCC('t2s')
+    # orig_transcript = converter.convert(orig_transcript)
+    # transcribe_state,_ = align_zh(traditional_to_simplified(segments), audio_path)
+    # transcribe_state['segments'] = traditional_to_simplified(transcribe_state['segments'])
 
-    print(orig_transcript)
-    print(target_transcript)
+    # print(orig_transcript)
+    # print(target_transcript)
 
-    operations, orig_spans = parse_edit_zh(orig_transcript, target_transcript)
-    print(operations)
-    print("orig_spans: ", orig_spans)
+    # operations, orig_spans = parse_edit_zh(orig_transcript, target_transcript)
+    # print(operations)
+    # print("orig_spans: ", orig_spans)
     
-    if len(orig_spans) > 3:
-        raise gr.Error("Current model only supports maximum 3 editings")
+    # if len(orig_spans) > 3:
+    #     raise gr.Error("Current model only supports maximum 3 editings")
         
-    starting_intervals = []
-    ending_intervals = []
-    for orig_span in orig_spans:
-        start, end = get_mask_interval(transcribe_state, orig_span)
-        starting_intervals.append(start)
-        ending_intervals.append(end)
-
-    print("intervals: ", starting_intervals, ending_intervals)
-
-    info = torchaudio.info(audio_path)
-    audio_dur = info.num_frames / info.sample_rate
-    
-    def combine_spans(spans, threshold=0.2):
-        spans.sort(key=lambda x: x[0])
-        combined_spans = []
-        current_span = spans[0]
-
-        for i in range(1, len(spans)):
-            next_span = spans[i]
-            if current_span[1] >= next_span[0] - threshold:
-                current_span[1] = max(current_span[1], next_span[1])
-            else:
-                combined_spans.append(current_span)
-                current_span = next_span
-        combined_spans.append(current_span)
-        return combined_spans
-    
-    morphed_span = [[max(start - sub_amount, 0), min(end + sub_amount, audio_dur)]
-                    for start, end in zip(starting_intervals, ending_intervals)] # in seconds
-    morphed_span = combine_spans(morphed_span, threshold=0.2)
-    print("morphed_spans: ", morphed_span)
-    mask_interval = [[round(span[0]*codec_sr), round(span[1]*codec_sr)] for span in morphed_span]
-    mask_interval = torch.LongTensor(mask_interval) # [M,2], M==1 for now
-    
-    decode_config = {'top_k': top_k, 'top_p': top_p, 'temperature': temperature, 'stop_repetition': stop_repetition, 'kvcache': kvcache, "codec_audio_sr": codec_audio_sr, "codec_sr": codec_sr}
-    
-    new_audio = inference_one_sample(
-        ssrspeech_model_zh["model"],
-        ssrspeech_model_zh["config"],
-        ssrspeech_model_zh["phn2num"],
-        ssrspeech_model_zh["text_tokenizer"], 
-        ssrspeech_model_zh["audio_tokenizer"],
-        audio_path, orig_transcript, target_transcript, mask_interval,
-        cfg_coef, cfg_stride, aug_text, False, True, False,
-        device, decode_config
-    )
-    audio_tensors = []
-    # save segments for comparison
-    new_audio = new_audio[0].cpu()
-    torchaudio.save(audio_path, new_audio, codec_audio_sr)
-    audio_tensors.append(new_audio)
-    output_audio = get_output_audio(audio_tensors, codec_audio_sr)
-    
-    success_message = "<span style='color:green;'>Success: Inference successfully!</span>"
-    return output_audio, success_message
-
-
-@spaces.GPU
-def run_tts_zh(seed, sub_amount, aug_text, cfg_coef, cfg_stride, prompt_length,
-        audio_path, original_transcript, transcript):
-    
-    codec_audio_sr = 16000
-    codec_sr = 50
-    top_k = 0
-    top_p = 0.8
-    temperature = 1
-    kvcache = 1
-    stop_repetition = 2
-    
-    aug_text = True if aug_text == 1 else False
-    
-    seed_everything(seed)
-
-    # resample audio
-    audio, _ = librosa.load(audio_path, sr=16000)
-    sf.write(audio_path, audio, 16000)
-    
-    # text normalization
-    target_transcript = transcript.replace("  ", " ").replace("  ", " ").replace("\n", " ")
-    orig_transcript = original_transcript.replace("  ", " ").replace("  ", " ").replace("\n", " ")
-
-    [orig_transcript, segments, _, _] = transcribe_zh(audio_path)
-
-    converter = opencc.OpenCC('t2s')
-    orig_transcript = converter.convert(orig_transcript)
-    transcribe_state,_ = align_zh(traditional_to_simplified(segments), audio_path)
-    transcribe_state['segments'] = traditional_to_simplified(transcribe_state['segments'])
-
-    print(orig_transcript)
-    print(target_transcript)
-
-    info = torchaudio.info(audio_path)
-    duration = info.num_frames / info.sample_rate
-    cut_length = duration
-    # Cut long audio for tts
-    if duration > prompt_length:
-        seg_num = len(transcribe_state['segments'])
-        for i in range(seg_num):
-            words = transcribe_state['segments'][i]['words']
-            for item in words:
-                if item['end'] >= prompt_length:
-                    cut_length = min(item['end'], cut_length)
-
-    audio, _ = librosa.load(audio_path, sr=16000, duration=cut_length)
-    sf.write(audio_path, audio, 16000)
-    [orig_transcript, segments, _, _] = transcribe_zh(audio_path)
+    # starting_intervals = []
+    # ending_intervals = []
+    # for orig_span in orig_spans:
+    #     start, end = get_mask_interval(transcribe_state, orig_span)
+    #     starting_intervals.append(start)
+    #     ending_intervals.append(end)
+
+    # print("intervals: ", starting_intervals, ending_intervals)
+
+    # info = torchaudio.info(audio_path)
+    # audio_dur = info.num_frames / info.sample_rate
+    
+    # def combine_spans(spans, threshold=0.2):
+    #     spans.sort(key=lambda x: x[0])
+    #     combined_spans = []
+    #     current_span = spans[0]
+
+    #     for i in range(1, len(spans)):
+    #         next_span = spans[i]
+    #         if current_span[1] >= next_span[0] - threshold:
+    #             current_span[1] = max(current_span[1], next_span[1])
+    #         else:
+    #             combined_spans.append(current_span)
+    #             current_span = next_span
+    #     combined_spans.append(current_span)
+    #     return combined_spans
+    
+    # morphed_span = [[max(start - sub_amount, 0), min(end + sub_amount, audio_dur)]
+    #                 for start, end in zip(starting_intervals, ending_intervals)] # in seconds
+    # morphed_span = combine_spans(morphed_span, threshold=0.2)
+    # print("morphed_spans: ", morphed_span)
+    # mask_interval = [[round(span[0]*codec_sr), round(span[1]*codec_sr)] for span in morphed_span]
+    # mask_interval = torch.LongTensor(mask_interval) # [M,2], M==1 for now
+    
+    # decode_config = {'top_k': top_k, 'top_p': top_p, 'temperature': temperature, 'stop_repetition': stop_repetition, 'kvcache': kvcache, "codec_audio_sr": codec_audio_sr, "codec_sr": codec_sr}
+    
+    # new_audio = inference_one_sample(
+    #     ssrspeech_model_zh["model"],
+    #     ssrspeech_model_zh["config"],
+    #     ssrspeech_model_zh["phn2num"],
+    #     ssrspeech_model_zh["text_tokenizer"], 
+    #     ssrspeech_model_zh["audio_tokenizer"],
+    #     audio_path, orig_transcript, target_transcript, mask_interval,
+    #     cfg_coef, cfg_stride, aug_text, False, True, False,
+    #     device, decode_config
+    # )
+    # audio_tensors = []
+    # # save segments for comparison
+    # new_audio = new_audio[0].cpu()
+    # torchaudio.save(audio_path, new_audio, codec_audio_sr)
+    # audio_tensors.append(new_audio)
+    # output_audio = get_output_audio(audio_tensors, codec_audio_sr)
+    
+    # success_message = "<span style='color:green;'>Success: Inference successfully!</span>"
+    # return output_audio, success_message
+
+
+# @spaces.GPU
+# def run_tts_zh(seed, sub_amount, aug_text, cfg_coef, cfg_stride, prompt_length,
+#         audio_path, original_transcript, transcript):
+    
+#     codec_audio_sr = 16000
+#     codec_sr = 50
+#     top_k = 0
+#     top_p = 0.8
+#     temperature = 1
+#     kvcache = 1
+#     stop_repetition = 2
+    
+#     aug_text = True if aug_text == 1 else False
+    
+#     seed_everything(seed)
+
+#     # resample audio
+#     audio, _ = librosa.load(audio_path, sr=16000)
+#     sf.write(audio_path, audio, 16000)
+    
+    # # text normalization
+    # target_transcript = transcript.replace("  ", " ").replace("  ", " ").replace("\n", " ")
+    # orig_transcript = original_transcript.replace("  ", " ").replace("  ", " ").replace("\n", " ")
+
+    # [orig_transcript, segments, _, _] = transcribe_zh(audio_path)
+
+    # converter = opencc.OpenCC('t2s')
+    # orig_transcript = converter.convert(orig_transcript)
+    # transcribe_state,_ = align_zh(traditional_to_simplified(segments), audio_path)
+    # transcribe_state['segments'] = traditional_to_simplified(transcribe_state['segments'])
+
+    # print(orig_transcript)
+    # print(target_transcript)
+
+    # info = torchaudio.info(audio_path)
+    # duration = info.num_frames / info.sample_rate
+    # cut_length = duration
+    # # Cut long audio for tts
+    # if duration > prompt_length:
+    #     seg_num = len(transcribe_state['segments'])
+    #     for i in range(seg_num):
+    #         words = transcribe_state['segments'][i]['words']
+    #         for item in words:
+    #             if item['end'] >= prompt_length:
+    #                 cut_length = min(item['end'], cut_length)
+
+    # audio, _ = librosa.load(audio_path, sr=16000, duration=cut_length)
+    # sf.write(audio_path, audio, 16000)
+    # [orig_transcript, segments, _, _] = transcribe_zh(audio_path)
+    
+
+    # converter = opencc.OpenCC('t2s')
+    # orig_transcript = converter.convert(orig_transcript)
+    # transcribe_state,_ = align_zh(traditional_to_simplified(segments), audio_path)
+    # transcribe_state['segments'] = traditional_to_simplified(transcribe_state['segments'])
+    
+    # print(orig_transcript)
+    # target_transcript_copy = target_transcript # for tts cut out
+    # target_transcript_copy = target_transcript_copy[0]
+    # target_transcript = orig_transcript + target_transcript
+    # print(target_transcript)
+
+    
+    # info = torchaudio.info(audio_path)
+    # audio_dur = info.num_frames / info.sample_rate
+    
+    # morphed_span = [(audio_dur, audio_dur)] # in seconds
+    # mask_interval = [[round(span[0]*codec_sr), round(span[1]*codec_sr)] for span in morphed_span]
+    # mask_interval = torch.LongTensor(mask_interval) # [M,2], M==1 for now
+    # print("mask_interval: ", mask_interval)
+
+    # decode_config = {'top_k': top_k, 'top_p': top_p, 'temperature': temperature, 'stop_repetition': stop_repetition, 'kvcache': kvcache, "codec_audio_sr": codec_audio_sr, "codec_sr": codec_sr}
+    
+    # new_audio = inference_one_sample(
+    #     ssrspeech_model_zh["model"],
+    #     ssrspeech_model_zh["config"],
+    #     ssrspeech_model_zh["phn2num"],
+    #     ssrspeech_model_zh["text_tokenizer"], 
+    #     ssrspeech_model_zh["audio_tokenizer"],
+    #     audio_path, orig_transcript, target_transcript, mask_interval,
+    #     cfg_coef, cfg_stride, aug_text, False, True, True,
+    #     device, decode_config
+    # )
+    # audio_tensors = []
+    # # save segments for comparison
+    # new_audio = new_audio[0].cpu()
+    # torchaudio.save(audio_path, new_audio, codec_audio_sr)
+    
+    # [new_transcript, new_segments, _,_] = transcribe_zh(audio_path)
+    
+    # transcribe_state,_ = align_zh(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
     
-
-    converter = opencc.OpenCC('t2s')
-    orig_transcript = converter.convert(orig_transcript)
-    transcribe_state,_ = align_zh(traditional_to_simplified(segments), audio_path)
-    transcribe_state['segments'] = traditional_to_simplified(transcribe_state['segments'])
-    
-    print(orig_transcript)
-    target_transcript_copy = target_transcript # for tts cut out
-    target_transcript_copy = target_transcript_copy[0]
-    target_transcript = orig_transcript + target_transcript
-    print(target_transcript)
-
-    
-    info = torchaudio.info(audio_path)
-    audio_dur = info.num_frames / info.sample_rate
-    
-    morphed_span = [(audio_dur, audio_dur)] # in seconds
-    mask_interval = [[round(span[0]*codec_sr), round(span[1]*codec_sr)] for span in morphed_span]
-    mask_interval = torch.LongTensor(mask_interval) # [M,2], M==1 for now
-    print("mask_interval: ", mask_interval)
-
-    decode_config = {'top_k': top_k, 'top_p': top_p, 'temperature': temperature, 'stop_repetition': stop_repetition, 'kvcache': kvcache, "codec_audio_sr": codec_audio_sr, "codec_sr": codec_sr}
-    
-    new_audio = inference_one_sample(
-        ssrspeech_model_zh["model"],
-        ssrspeech_model_zh["config"],
-        ssrspeech_model_zh["phn2num"],
-        ssrspeech_model_zh["text_tokenizer"], 
-        ssrspeech_model_zh["audio_tokenizer"],
-        audio_path, orig_transcript, target_transcript, mask_interval,
-        cfg_coef, cfg_stride, aug_text, False, True, True,
-        device, decode_config
-    )
-    audio_tensors = []
-    # save segments for comparison
-    new_audio = new_audio[0].cpu()
-    torchaudio.save(audio_path, new_audio, codec_audio_sr)
-    
-    [new_transcript, new_segments, _,_] = transcribe_zh(audio_path)
-    
-    transcribe_state,_ = align_zh(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
-    
-    if tmp1 == tmp2:
-        offset = transcribe_state['segments'][0]['words'][0]['start']
-    else:
-        offset = transcribe_state['segments'][0]['words'][1]['start']
-    
-    new_audio, _ = torchaudio.load(audio_path, frame_offset=int(offset*codec_audio_sr))
-    audio_tensors.append(new_audio)
-    output_audio = get_output_audio(audio_tensors, codec_audio_sr)
-    
-    success_message = "<span style='color:green;'>Success: Inference successfully!</span>"
-    return output_audio, success_message
+    # if tmp1 == tmp2:
+    #     offset = transcribe_state['segments'][0]['words'][0]['start']
+    # else:
+    #     offset = transcribe_state['segments'][0]['words'][1]['start']
+    
+    # new_audio, _ = torchaudio.load(audio_path, frame_offset=int(offset*codec_audio_sr))
+    # audio_tensors.append(new_audio)
+    # output_audio = get_output_audio(audio_tensors, codec_audio_sr)
+    
+    # success_message = "<span style='color:green;'>Success: Inference successfully!</span>"
+    # return output_audio, success_message
 
 
 if __name__ == "__main__":
@@ -818,131 +815,131 @@ if __name__ == "__main__":
                     outputs=[output_audio, success_output]
                 )
                 
-            with gr.Tab("Mandarin Speech Editing"):
+            # with gr.Tab("Mandarin Speech Editing"):
                 
-                with gr.Row():
-                    with gr.Column(scale=2):
-                        input_audio = gr.Audio(value=f"{DEMO_PATH}/aishell3_test.wav", label="Input Audio", type="filepath", interactive=True)
-                        with gr.Group():
-                            original_transcript = gr.Textbox(label="Original transcript", lines=5, value="价格已基本都在三万到六万之间",
-                                                            info="Use whisperx model to get the transcript.")
-                            transcribe_btn = gr.Button(value="Transcribe")
-
-                    with gr.Column(scale=3):
-                        with gr.Group():
-                            transcript = gr.Textbox(label="Text", lines=7, value="价格已基本都在一万到两万之间", interactive=True)
-                            run_btn = gr.Button(value="Run")
-
-                    with gr.Column(scale=2):
-                        output_audio = gr.Audio(label="Output Audio")
+            #     with gr.Row():
+            #         with gr.Column(scale=2):
+            #             input_audio = gr.Audio(value=f"{DEMO_PATH}/aishell3_test.wav", label="Input Audio", type="filepath", interactive=True)
+            #             with gr.Group():
+            #                 original_transcript = gr.Textbox(label="Original transcript", lines=5, value="价格已基本都在三万到六万之间",
+            #                                                 info="Use whisperx model to get the transcript.")
+            #                 transcribe_btn = gr.Button(value="Transcribe")
+
+            #         with gr.Column(scale=3):
+            #             with gr.Group():
+            #                 transcript = gr.Textbox(label="Text", lines=7, value="价格已基本都在一万到两万之间", interactive=True)
+            #                 run_btn = gr.Button(value="Run")
+
+            #         with gr.Column(scale=2):
+            #             output_audio = gr.Audio(label="Output Audio")
                         
-                with gr.Row():
-                    with gr.Accordion("Advanced Settings", open=False):
-                        seed = gr.Number(label="seed", value=-1, precision=0, info="random seeds always works :)")
-                        aug_text = gr.Radio(label="aug_text", choices=[0, 1], value=1,
-                                            info="set to 1 to use classifer-free guidance, change if you don't like the results")
-                        cfg_coef = gr.Number(label="cfg_coef", value=1.5,
-                                            info="cfg guidance scale, 1.5 is a good value, change if you don't like the results")
-                        cfg_stride = gr.Number(label="cfg_stride", value=1,
-                                            info="cfg stride, 1 is a good value for Mandarin, change if you don't like the results")
-                        prompt_length = gr.Number(label="prompt_length", value=3,
-                                            info="used for tts prompt, will automatically cut the prompt audio to this length")
-                        sub_amount = gr.Number(label="sub_amount", value=0.12, info="margin to the left and right of the editing segment, change if you don't like the results")
-
-                success_output = gr.HTML()
-
-                semgents = gr.State() # not used
-                state = gr.State() # not used
-                audio_state = gr.State(value=f"{DEMO_PATH}/aishell3_test.wav")
-                input_audio.change(
-                    lambda audio: audio,
-                    inputs=[input_audio],
-                    outputs=[audio_state]
-                )
+            #     with gr.Row():
+            #         with gr.Accordion("Advanced Settings", open=False):
+            #             seed = gr.Number(label="seed", value=-1, precision=0, info="random seeds always works :)")
+            #             aug_text = gr.Radio(label="aug_text", choices=[0, 1], value=1,
+            #                                 info="set to 1 to use classifer-free guidance, change if you don't like the results")
+            #             cfg_coef = gr.Number(label="cfg_coef", value=1.5,
+            #                                 info="cfg guidance scale, 1.5 is a good value, change if you don't like the results")
+            #             cfg_stride = gr.Number(label="cfg_stride", value=1,
+            #                                 info="cfg stride, 1 is a good value for Mandarin, change if you don't like the results")
+            #             prompt_length = gr.Number(label="prompt_length", value=3,
+            #                                 info="used for tts prompt, will automatically cut the prompt audio to this length")
+            #             sub_amount = gr.Number(label="sub_amount", value=0.12, info="margin to the left and right of the editing segment, change if you don't like the results")
+
+            #     success_output = gr.HTML()
+
+            #     semgents = gr.State() # not used
+            #     state = gr.State() # not used
+            #     audio_state = gr.State(value=f"{DEMO_PATH}/aishell3_test.wav")
+            #     input_audio.change(
+            #         lambda audio: audio,
+            #         inputs=[input_audio],
+            #         outputs=[audio_state]
+            #     )
                 
-                transcribe_btn.click(fn=transcribe_zh,
-                                    inputs=[audio_state],
-                                    outputs=[original_transcript, semgents, state, success_output])
+            #     transcribe_btn.click(fn=transcribe_zh,
+            #                         inputs=[audio_state],
+            #                         outputs=[original_transcript, semgents, state, success_output])
                 
-                run_btn.click(fn=run_edit_zh,
-                            inputs=[
-                                seed, sub_amount,
-                                aug_text, cfg_coef, cfg_stride, prompt_length, 
-                                audio_state, original_transcript, transcript,
-                            ],
-                            outputs=[output_audio, success_output])
-
-                transcript.submit(fn=run_edit_zh,
-                        inputs=[
-                                seed, sub_amount,
-                                aug_text, cfg_coef, cfg_stride, prompt_length, 
-                                audio_state, original_transcript, transcript,
-                        ],
-                    outputs=[output_audio, success_output]
-                )
+            #     run_btn.click(fn=run_edit_zh,
+            #                 inputs=[
+            #                     seed, sub_amount,
+            #                     aug_text, cfg_coef, cfg_stride, prompt_length, 
+            #                     audio_state, original_transcript, transcript,
+            #                 ],
+            #                 outputs=[output_audio, success_output])
+
+            #     transcript.submit(fn=run_edit_zh,
+            #             inputs=[
+            #                     seed, sub_amount,
+            #                     aug_text, cfg_coef, cfg_stride, prompt_length, 
+            #                     audio_state, original_transcript, transcript,
+            #             ],
+            #         outputs=[output_audio, success_output]
+            #     )
                 
-            with gr.Tab("Mandarin TTS"):
+            # with gr.Tab("Mandarin TTS"):
                 
-                with gr.Row():
-                    with gr.Column(scale=2):
-                        input_audio = gr.Audio(value=f"{DEMO_PATH}/aishell3_test.wav", label="Input Audio", type="filepath", interactive=True)
-                        with gr.Group():
-                            original_transcript = gr.Textbox(label="Original transcript", lines=5, value="价格已基本都在三万到六万之间",
-                                                            info="Use whisperx model to get the transcript.")
-                            transcribe_btn = gr.Button(value="Transcribe")
-
-                    with gr.Column(scale=3):
-                        with gr.Group():
-                            transcript = gr.Textbox(label="Text", lines=7, value="我简直不敢相信同一个模型也可以进行文本到语音的生成", interactive=True)
-                            run_btn = gr.Button(value="Run")
-
-                    with gr.Column(scale=2):
-                        output_audio = gr.Audio(label="Output Audio")
+            #     with gr.Row():
+            #         with gr.Column(scale=2):
+            #             input_audio = gr.Audio(value=f"{DEMO_PATH}/aishell3_test.wav", label="Input Audio", type="filepath", interactive=True)
+            #             with gr.Group():
+            #                 original_transcript = gr.Textbox(label="Original transcript", lines=5, value="价格已基本都在三万到六万之间",
+            #                                                 info="Use whisperx model to get the transcript.")
+            #                 transcribe_btn = gr.Button(value="Transcribe")
+
+            #         with gr.Column(scale=3):
+            #             with gr.Group():
+            #                 transcript = gr.Textbox(label="Text", lines=7, value="我简直不敢相信同一个模型也可以进行文本到语音的生成", interactive=True)
+            #                 run_btn = gr.Button(value="Run")
+
+            #         with gr.Column(scale=2):
+            #             output_audio = gr.Audio(label="Output Audio")
                         
-                with gr.Row():
-                    with gr.Accordion("Advanced Settings", open=False):
-                        seed = gr.Number(label="seed", value=-1, precision=0, info="random seeds always works :)")
-                        aug_text = gr.Radio(label="aug_text", choices=[0, 1], value=1,
-                                            info="set to 1 to use classifer-free guidance, change if you don't like the results")
-                        cfg_coef = gr.Number(label="cfg_coef", value=1.5,
-                                            info="cfg guidance scale, 1.5 is a good value, change if you don't like the results")
-                        cfg_stride = gr.Number(label="cfg_stride", value=1,
-                                            info="cfg stride, 1 is a good value for Mandarin, change if you don't like the results")
-                        prompt_length = gr.Number(label="prompt_length", value=3,
-                                            info="used for tts prompt, will automatically cut the prompt audio to this length")
-                        sub_amount = gr.Number(label="sub_amount", value=0.12, info="margin to the left and right of the editing segment, change if you don't like the results")
-
-                success_output = gr.HTML()
-
-                semgents = gr.State() # not used
-                state = gr.State() # not used
-                audio_state = gr.State(value=f"{DEMO_PATH}/aishell3_test.wav")
-                input_audio.change(
-                    lambda audio: audio,
-                    inputs=[input_audio],
-                    outputs=[audio_state]
-                )
+            #     with gr.Row():
+            #         with gr.Accordion("Advanced Settings", open=False):
+            #             seed = gr.Number(label="seed", value=-1, precision=0, info="random seeds always works :)")
+            #             aug_text = gr.Radio(label="aug_text", choices=[0, 1], value=1,
+            #                                 info="set to 1 to use classifer-free guidance, change if you don't like the results")
+            #             cfg_coef = gr.Number(label="cfg_coef", value=1.5,
+            #                                 info="cfg guidance scale, 1.5 is a good value, change if you don't like the results")
+            #             cfg_stride = gr.Number(label="cfg_stride", value=1,
+            #                                 info="cfg stride, 1 is a good value for Mandarin, change if you don't like the results")
+            #             prompt_length = gr.Number(label="prompt_length", value=3,
+            #                                 info="used for tts prompt, will automatically cut the prompt audio to this length")
+            #             sub_amount = gr.Number(label="sub_amount", value=0.12, info="margin to the left and right of the editing segment, change if you don't like the results")
+
+            #     success_output = gr.HTML()
+
+            #     semgents = gr.State() # not used
+            #     state = gr.State() # not used
+            #     audio_state = gr.State(value=f"{DEMO_PATH}/aishell3_test.wav")
+            #     input_audio.change(
+            #         lambda audio: audio,
+            #         inputs=[input_audio],
+            #         outputs=[audio_state]
+            #     )
                 
-                transcribe_btn.click(fn=transcribe_zh,
-                                    inputs=[audio_state],
-                                    outputs=[original_transcript, semgents, state, success_output])
+            #     transcribe_btn.click(fn=transcribe_zh,
+            #                         inputs=[audio_state],
+            #                         outputs=[original_transcript, semgents, state, success_output])
                 
-                run_btn.click(fn=run_tts_zh,
-                            inputs=[
-                                seed, sub_amount,
-                                aug_text, cfg_coef, cfg_stride, prompt_length, 
-                                audio_state, original_transcript, transcript,
-                            ],
-                            outputs=[output_audio, success_output])
-
-                transcript.submit(fn=run_tts_zh,
-                        inputs=[
-                                seed, sub_amount,
-                                aug_text, cfg_coef, cfg_stride, prompt_length, 
-                                audio_state, original_transcript, transcript,
-                        ],
-                    outputs=[output_audio, success_output]
-                )
+            #     run_btn.click(fn=run_tts_zh,
+            #                 inputs=[
+            #                     seed, sub_amount,
+            #                     aug_text, cfg_coef, cfg_stride, prompt_length, 
+            #                     audio_state, original_transcript, transcript,
+            #                 ],
+            #                 outputs=[output_audio, success_output])
+
+            #     transcript.submit(fn=run_tts_zh,
+            #             inputs=[
+            #                     seed, sub_amount,
+            #                     aug_text, cfg_coef, cfg_stride, prompt_length, 
+            #                     audio_state, original_transcript, transcript,
+            #             ],
+            #         outputs=[output_audio, success_output]
+            #     )
 
         # Launch the Gradio demo
         demo.launch()