Update app.py

app.py

@@ -10,6 +10,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 inference_scale import inference_one_sample
 import librosa
 import soundfile as sf
@@ -70,37 +72,57 @@ def get_mask_interval(transcribe_state, word_span):
 
     return (start, end)
 
-from whisperx import load_align_model, load_model, load_audio
-from whisperx import align as align_func
+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
 
-ssrspeech_model_name = "English"
-text_tokenizer = TextTokenizer(backend="espeak")
-language = "en"
-transcribe_model_name = "base.en"
 
-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)
+from whisperx import load_align_model, load_model, load_audio
+from whisperx import align as align_func
 
-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)
+# Load models
+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_zh = f"{MODELS_PATH}/Mandarin.pth"
+ckpt_zh = torch.load(ssrspeech_fn_zh)
+model_zh = ssr.SSR_Speech(ckpt_zh["config"])
+model_zh.load_state_dict(ckpt_zh["model"])
+config_zh = model_zh.args
+phn2num_zh = ckpt_zh["phn2num"]
+model_zh.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,
+
+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,
+    "phn2num": phn2num_zh,
+    "model": model_zh,
+    "text_tokenizer": text_tokenizer_zh,
+    "audio_tokenizer": AudioTokenizer(signature=encodec_fn)
+}
 
 
 def get_transcribe_state(segments):
@@ -112,7 +134,9 @@ def get_transcribe_state(segments):
     }
 
 @spaces.GPU
-def transcribe(audio_path):
+def transcribe_en(audio_path):
+    language = "en"
+    transcribe_model_name = "base.en"
     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"]
     for segment in segments:
@@ -126,9 +150,37 @@ def transcribe(audio_path):
         state, success_message
     ]
 
+@spaces.GPU
+def transcribe_zh(audio_path):
+    language = "zh"
+    transcribe_model_name = "base"
+    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"]
+    for segment in segments:
+        segment['text'] = replace_numbers_with_words(segment['text'])
+    _, segments = align(segments, audio_path)
+    state = get_transcribe_state(segments)
+    success_message = "<span style='color:green;'>Success: Transcribe completed successfully!</span>"
+
+    return [
+        state["transcript"], state['segments'],
+        state, success_message
+    ]
 
 @spaces.GPU
-def align(segments, audio_path):
+def align_en(segments, audio_path):
+    language = "en"
+    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
+
+
+@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"]
@@ -155,14 +207,18 @@ 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, 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):
+def run_edit_en(seed, sub_amount, aug_text, cfg_coef, prompt_length,
+        audio_path, original_transcript, transcript):
     
-    aug_text = True if aug_text == 1 else False
-    if ssrspeech_model is None:
-        raise gr.Error("ssrspeech model not loaded")
+    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
@@ -173,118 +229,269 @@ def run(seed, sub_amount, codec_audio_sr, codec_sr, top_k, top_p, temperature,
     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_en(audio_path)
     orig_transcript = orig_transcript.lower()
     target_transcript = target_transcript.lower()
-    transcribe_state,_ = align(segments, audio_path)
+    transcribe_state,_ = align_en(segments, audio_path)
     print(orig_transcript)
     print(target_transcript)
 
-    if mode == "TTS":
-        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(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
-        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)
-        print(operations)
-        print("orig_spans: ", orig_spans)
-        
-        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
+    operations, orig_spans = parse_edit_en(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")
         
-        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
+    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_en["model"],
+        ssrspeech_model_en["config"],
+        ssrspeech_model_en["phn2num"],
+        ssrspeech_model_en["text_tokenizer"], 
+        ssrspeech_model_en["audio_tokenizer"],
+        audio_path, orig_transcript, target_transcript, mask_interval,
+        cfg_coef, 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_en(seed, sub_amount, aug_text, cfg_coef, 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 = 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_en(audio_path)
+    orig_transcript = orig_transcript.lower()
+    target_transcript = target_transcript.lower()
+    transcribe_state,_ = align_en(segments, audio_path)
+    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_en(audio_path)
+    
+    
+    orig_transcript = orig_transcript.lower()
+    target_transcript = target_transcript.lower()
+    transcribe_state,_ = align_en(segments, audio_path)
+    print(orig_transcript)
+    target_transcript_copy = target_transcript # for tts cut out
+    target_transcript_copy = target_transcript_copy.split(' ')[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_en["model"],
+        ssrspeech_model_en["config"],
+        ssrspeech_model_en["phn2num"],
+        ssrspeech_model_en["text_tokenizer"], 
+        ssrspeech_model_en["audio_tokenizer"],
+        audio_path, orig_transcript, target_transcript, mask_interval,
+        cfg_coef, 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_en(audio_path)
+    transcribe_state,_ = align_en(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:
-        info = torchaudio.info(audio_path)
-        audio_dur = info.num_frames / info.sample_rate
+        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
+
+
+@spaces.GPU
+def run_edit_zh(seed, sub_amount, aug_text, cfg_coef, 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 = 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_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)
+
+    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")
         
-        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)
+    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}
     
-    tts = True if mode == "TTS" else False
     new_audio = inference_one_sample(
-        ssrspeech_model["model"],
-        ssrspeech_model["config"],
-        ssrspeech_model["phn2num"],
-        ssrspeech_model["text_tokenizer"], 
-        ssrspeech_model["audio_tokenizer"],
+        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, aug_text, False, True, tts,
+        cfg_coef, 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)
-    if tts: # remove the start parts
-        [new_transcript, new_segments, _, _] = transcribe(audio_path)
-        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:
-            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)
     
@@ -292,88 +499,112 @@ def run(seed, sub_amount, codec_audio_sr, codec_sr, top_k, top_p, temperature,
     return output_audio, success_message
 
 
-demo_original_transcript = "Gwynplain had besides for his work and for his feats of strength, round his neck and over his shoulders, an esclavine of leather."
+@spaces.GPU
+def run_tts_zh(seed, sub_amount, aug_text, cfg_coef, 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)
 
-demo_text = {
-    "TTS": {
-        "regular": "Gwynplain 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": {
-        "regular": "Gwynplain had besides for his work and feats of strength, hanging from his neck and shoulders, an esclavine of leather."
-    },
-}
+    # resample audio
+    audio, _ = librosa.load(audio_path, sr=16000)
+    sf.write(audio_path, audio, 16000)
+    
+    # text normalization
+    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_zh(audio_path)
 
-def get_app():
-    with gr.Blocks() as app:
-        gr.Markdown("""
-            # SSR-Speech: High-quality Speech Editor and Text-to-Speech Synthesizer
-            Generate and edit speech from text. Adjust advanced settings for more control.
-            
-            Learn more about 🟣**SSR-Speech** on the [SSR-Speech Homepage](https://wanghelin1997.github.io/SSR-Speech-Demo/).
-            
-            🚀 The **SSR-Speech (Mandarin)** demo is now live! Try it on [🤗SSR-Speech-Mandarin Space](https://huggingface.co/spaces/OpenSound/SSR-Speech-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)
-                with gr.Group():
-                    original_transcript = gr.Textbox(label="Original transcript", lines=5, value=demo_original_transcript,
-                                                    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=demo_text["Edit"]["regular"])
-
-                    with gr.Row():
-                        mode = gr.Radio(label="Mode", choices=["Edit", "TTS"], value="Edit")
-
-                    run_btn = gr.Button(value="Run")
-
-            with gr.Column(scale=2):
-                output_audio = gr.Audio(label="Output Audio")
-
-        with gr.Row():
-            with gr.Accordion("Generation Parameters - change these if you are unhappy with the generation", open=False):
-                stop_repetition = gr.Radio(label="stop_repetition", choices=[-1, 1, 2, 3, 4], value=2,
-                                        info="if there are long silence in the generated audio, reduce the stop_repetition to 2 or 1. -1 = disabled")
-                seed = gr.Number(label="seed", value=-1, precision=0, info="random seeds always works :)")
-                kvcache = gr.Radio(label="kvcache", choices=[0, 1], value=1,
-                                    info="set to 0 to use less VRAM, but with slower inference")
-                aug_text = gr.Radio(label="aug_text", choices=[0, 1], value=1,
-                                    info="set to 1 to use cfg")
-                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")
-                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")
-                top_p = gr.Number(label="top_p", value=0.8, info="0.9 is a good value, 0.8 is also good")
-                temperature = gr.Number(label="temperature", value=1, info="haven't try other values, do not change")
-                top_k = gr.Number(label="top_k", value=0, info="0 means we don't use topk sampling, because we use topp sampling")
-                codec_audio_sr = gr.Number(label="codec_audio_sr", value=16000, info='encodec specific, do not change')
-                codec_sr = gr.Number(label="codec_sr", value=50, info='encodec specific, do not change')
-                silence_tokens = gr.Textbox(label="silence tokens", value="[1388,1898,131]", info="encodec specific, do not change")
-
-        success_output = gr.HTML()
-
-        semgents = gr.State() # not used
-        transcribe_btn.click(fn=transcribe,
-                            inputs=[input_audio],
-                            outputs=[original_transcript, semgents, success_output])
-        
-        run_btn.click(fn=run,
-                    inputs=[
-                        seed, sub_amount,
-                        codec_audio_sr, codec_sr,
-                        top_k, top_p, temperature, stop_repetition, kvcache, silence_tokens, 
-                        aug_text, cfg_coef, prompt_length, 
-                        input_audio, original_transcript, transcript,
-                        mode
-                    ],
-                    outputs=[output_audio, success_output])
+    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)
 
-    return app
+    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, 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 __name__ == "__main__":
@@ -393,5 +624,244 @@ if __name__ == "__main__":
     TMP_PATH = args.tmp_path
     MODELS_PATH = args.models_path
 
-    app = get_app()
-    app.queue().launch(share=args.share, server_port=args.port)
+    # app = get_app()
+    # app.queue().launch(share=args.share, server_port=args.port)
+    
+    # CSS styling (optional)
+    css = """
+    #col-container {
+        margin: 0 auto;
+        max-width: 1280px;
+    }
+    """
+    
+    # Gradio Blocks layout
+    with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
+        with gr.Column(elem_id="col-container"):
+            gr.Markdown("""
+                # SSR-Speech: High-quality Speech Editor and Text-to-Speech Synthesizer
+                Generate and edit speech from text. Adjust advanced settings for more control.
+                
+                Learn more about 🚀**SSR-Speech** on the [SSR-Speech Homepage](https://wanghelin1997.github.io/SSR-Speech-Demo/).
+            """)
+
+
+            # Tabs for Generate and Edit
+            with gr.Tab("English Speech Editing"):
+                
+                with gr.Row():
+                    with gr.Column(scale=2):
+                        input_audio = gr.Audio(value=f"{DEMO_PATH}/84_121550_000074_000000.wav", label="Input Audio", type="filepath", interactive=True)
+                        with gr.Group():
+                            original_transcript = gr.Textbox(label="Original transcript", lines=5, value="Debug",
+                                                            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="Debug", 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")
+                        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
+                transcribe_btn.click(fn=transcribe_en,
+                                    inputs=[input_audio],
+                                    outputs=[original_transcript, semgents, state, success_output])
+                
+                run_btn.click(fn=run_edit_en,
+                            inputs=[
+                                seed, sub_amount,
+                                aug_text, cfg_coef, prompt_length, 
+                                input_audio, original_transcript, transcript,
+                            ],
+                            outputs=[output_audio, success_output])
+
+                transcript.submit(fn=run_edit_en,
+                        inputs=[
+                                seed, sub_amount,
+                                aug_text, cfg_coef, prompt_length, 
+                                input_audio, original_transcript, transcript,
+                        ],
+                    outputs=[output_audio, success_output]
+                )
+
+            with gr.Tab("English TTS"):
+                
+                with gr.Row():
+                    with gr.Column(scale=2):
+                        input_audio = gr.Audio(value=f"{DEMO_PATH}/84_121550_000074_000000.wav", label="Input Audio", type="filepath", interactive=True)
+                        with gr.Group():
+                            original_transcript = gr.Textbox(label="Original transcript", lines=5, value="Debug",
+                                                            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="Debug", 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")
+                        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
+                transcribe_btn.click(fn=transcribe_en,
+                                    inputs=[input_audio],
+                                    outputs=[original_transcript, semgents, state, success_output])
+                
+                run_btn.click(fn=run_tts_en,
+                            inputs=[
+                                seed, sub_amount,
+                                aug_text, cfg_coef, prompt_length, 
+                                input_audio, original_transcript, transcript,
+                            ],
+                            outputs=[output_audio, success_output])
+
+                transcript.submit(fn=run_tts_en,
+                        inputs=[
+                                seed, sub_amount,
+                                aug_text, cfg_coef, prompt_length, 
+                                input_audio, original_transcript, transcript,
+                        ],
+                    outputs=[output_audio, success_output]
+                )
+                
+            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="Debug",
+                                                            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="Debug", 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")
+                        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
+                transcribe_btn.click(fn=transcribe_zh,
+                                    inputs=[input_audio],
+                                    outputs=[original_transcript, semgents, state, success_output])
+                
+                run_btn.click(fn=run_edit_zh,
+                            inputs=[
+                                seed, sub_amount,
+                                aug_text, cfg_coef, prompt_length, 
+                                input_audio, original_transcript, transcript,
+                            ],
+                            outputs=[output_audio, success_output])
+
+                transcript.submit(fn=run_edit_zh,
+                        inputs=[
+                                seed, sub_amount,
+                                aug_text, cfg_coef, prompt_length, 
+                                input_audio, original_transcript, transcript,
+                        ],
+                    outputs=[output_audio, success_output]
+                )
+                
+            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="Debug",
+                                                            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="Debug", 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")
+                        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
+                transcribe_btn.click(fn=transcribe_zh,
+                                    inputs=[input_audio],
+                                    outputs=[original_transcript, semgents, state, success_output])
+                
+                run_btn.click(fn=run_tts_zh,
+                            inputs=[
+                                seed, sub_amount,
+                                aug_text, cfg_coef, prompt_length, 
+                                input_audio, original_transcript, transcript,
+                            ],
+                            outputs=[output_audio, success_output])
+
+                transcript.submit(fn=run_tts_zh,
+                        inputs=[
+                                seed, sub_amount,
+                                aug_text, cfg_coef, prompt_length, 
+                                input_audio, original_transcript, transcript,
+                        ],
+                    outputs=[output_audio, success_output]
+                )
+
+        # Launch the Gradio demo
+        demo.launch()