Update inference_gradio.py

inference_gradio.py

@@ -1,576 +1,1197 @@
-import gc
-import os
-import platform
-import psutil
-import tempfile
-from glob import glob
-import traceback
-import click
+import os, gc
+import re, time
+import logging
+from num2words import num2words
 import gradio as gr
-import torch
-import torchaudio
-import soundfile as sf
-from pathlib import Path
-
-from cached_path import cached_path
+import torch, torchaudio
+import numpy as np
+import random
+from scipy.io import wavfile
+import onnx
+import onnxruntime as ort
+import copy
+import uroman as ur
+import jieba, zhconv
+from pypinyin.core import Pinyin
+from pypinyin import Style
 
 from lemas_tts.api import TTS, PRETRAINED_ROOT, CKPTS_ROOT
-
-# Global variables
-tts_api = None
-last_checkpoint = ""
-last_device = ""
-last_ema = None
-
-# Device detection
-device = (
-    "cuda"
-    if torch.cuda.is_available()
-    else "xpu"
-    if torch.xpu.is_available()
-    else "mps"
-    if torch.backends.mps.is_available()
-    else "cpu"
+from lemas_tts.infer.edit_multilingual import gen_wav_multilingual
+from lemas_tts.infer.text_norm.txt2pinyin import (
+    MyConverter,
+    _PAUSE_SYMBOL,
+    change_tone_in_bu_or_yi,
+    get_phoneme_from_char_and_pinyin,
 )
+from lemas_tts.infer.text_norm.cn_tn import NSWNormalizer
+# import io
+# import uuid
+_JIEBA_DICT = os.path.join(
+    os.path.dirname(__file__),
+    "lemas_tts",
+    "infer",
+    "text_norm",
+    "jieba_dict.txt",
+)
+if os.path.isfile(_JIEBA_DICT):
+    jieba.set_dictionary(_JIEBA_DICT)
 
-REPO_ROOT = Path(__file__).resolve().parent
+# from inference_tts_scale import inference_one_sample as inference_tts
+import langid
+langid.set_languages(['es','pt','zh','en','de','fr','it', 'ru', 'id', 'vi'])
 
-# HF location for large TTS checkpoints (too big for Space storage)
-HF_PRETRAINED_ROOT = "hf://LEMAS-Project/LEMAS-TTS/pretrained_models"
 
-# 指向 `pretrained_models` 里的 espeak-ng-data（本地自带的字典）
-# 动态库交给系统安装的 espeak-ng 来提供（通过 apt），不强行指定 PHONEMIZER_ESPEAK_LIBRARY，
-# 避免本地复制的 .so 与 Space 基础镜像不兼容。
-ESPEAK_DATA_DIR = Path(PRETRAINED_ROOT) / "espeak-ng-data"
-os.environ["ESPEAK_DATA_PATH"] = str(ESPEAK_DATA_DIR)
-os.environ["ESPEAKNG_DATA_PATH"] = str(ESPEAK_DATA_DIR)
+os.environ['CURL_CA_BUNDLE'] = ''
+DEMO_PATH = os.getenv("DEMO_PATH", "./demo")
+TMP_PATH = os.getenv("TMP_PATH", "./demo/temp")
+MODELS_PATH = os.getenv("MODELS_PATH", "./pretrained_models")
 
+device = "cuda" if torch.cuda.is_available() else "cpu"
+ASR_DEVICE = "cpu"  # force whisperx/pyannote to CPU to avoid cuDNN issues
+whisper_model, align_model = None, None
+tts_edit_model = None
 
-class UVR5:
-    """Small wrapper around the bundled uvr5 implementation for denoising."""
+_whitespace_re = re.compile(r"\s+")
+alpha_pattern = re.compile(r"[a-zA-Z]")
+
+formatter = ("%(asctime)s [%(levelname)s] %(filename)s:%(lineno)d || %(message)s")
+logging.basicConfig(format=formatter, level=logging.INFO)
+
+# def get_random_string():
+#     return "".join(str(uuid.uuid4()).split("-"))
+
+def seed_everything(seed):
+    if seed != -1:
+        os.environ['PYTHONHASHSEED'] = str(seed)
+        random.seed(seed)
+        np.random.seed(seed)
+        torch.manual_seed(seed)
+        torch.cuda.manual_seed(seed)
+        torch.backends.cudnn.benchmark = False
+        torch.backends.cudnn.deterministic = True
 
-    def __init__(self, model_dir: Path, code_dir: Path):
-        self.model = self.load_model(str(model_dir), str(code_dir))
 
-    def load_model(self, model_dir: str, code_dir: str):
-        import sys
-        import json
+class UVR5:
+    """Small wrapper around the bundled uvr5 implementation for denoising."""
 
+    def __init__(self, model_dir):
+        code_dir = os.path.join(os.path.dirname(__file__), "uvr5")
+        self.model = self.load_model(model_dir, code_dir)
+    
+    def load_model(self, model_dir, code_dir):
+        import sys, json
         if code_dir not in sys.path:
             sys.path.append(code_dir)
-
         from multiprocess_cuda_infer import ModelData, Inference
-
         model_path = os.path.join(model_dir, "Kim_Vocal_1.onnx")
         config_path = os.path.join(model_dir, "MDX-Net-Kim-Vocal1.json")
         with open(config_path, "r", encoding="utf-8") as f:
             configs = json.load(f)
         model_data = ModelData(
             model_path=model_path,
-            audio_path=model_dir,
-            result_path=model_dir,
-            device="cpu",
-            process_method="MDX-Net",
-            base_dir=model_dir,  # keep base_dir and model_dir the same (paths under `pretrained_models`)
-            **configs,
+            audio_path = model_dir,
+            result_path = model_dir,
+            device = 'cpu',
+            process_method = "MDX-Net",
+            base_dir=model_dir,
+            **configs
         )
 
-        uvr5_model = Inference(model_data, "cpu")
+        uvr5_model = Inference(model_data, 'cpu')
         uvr5_model.load_model(model_path, 1)
         return uvr5_model
-
+        
     def denoise(self, audio_info):
-        print("denoise UVR5: ", audio_info)
         input_audio = load_wav(audio_info, sr=44100, channel=2)
-        output_audio = self.model.demix_base({0: input_audio.squeeze()}, is_match_mix=False)
+        output_audio = self.model.demix_base({0:input_audio.squeeze()}, is_match_mix=False)
+        # transform = torchaudio.transforms.Resample(44100, 16000)
+        # output_audio = transform(output_audio)
         return output_audio.squeeze().T.numpy(), 44100
 
-denoise_model = UVR5(
-    model_dir=str(Path(PRETRAINED_ROOT) / "uvr5"),
-    code_dir=str(REPO_ROOT / "uvr5"),
-)
+
+class DeepFilterNet:
+    def __init__(self, model_path):
+        self.hop_size = 480
+        self.fft_size = 960
+        self.model = self.load_model(model_path)
+
+
+    def load_model(self, model_path, threads=1):
+        sess_options = ort.SessionOptions()
+        sess_options.intra_op_num_threads = threads
+        sess_options.graph_optimization_level = (ort.GraphOptimizationLevel.ORT_ENABLE_EXTENDED)
+        sess_options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
+
+        model = onnx.load_model(model_path)
+        ort_session = ort.InferenceSession(
+            model.SerializeToString(),
+            sess_options,
+            providers=["CPUExecutionProvider"], # ["CUDAExecutionProvider"], #
+        )
+
+        input_names = ["input_frame", "states", "atten_lim_db"]
+        output_names = ["enhanced_audio_frame", "new_states", "lsnr"]
+        return ort_session
+        
+
+    def denoise(self, audio_info):
+        wav = load_wav(audio_info, 48000)
+        orig_len = wav.shape[-1]
+        hop_size_divisible_padding_size = (self.hop_size - orig_len % self.hop_size) % self.hop_size
+        orig_len += hop_size_divisible_padding_size
+        wav = torch.nn.functional.pad(
+            wav, (0, self.fft_size + hop_size_divisible_padding_size)
+        )
+        chunked_audio = torch.split(wav, self.hop_size)
+        # chunked_audio = torch.split(wav, int(wav.shape[-1]/2))
+
+        state = np.zeros(45304,dtype=np.float32)
+        atten_lim_db = np.zeros(1,dtype=np.float32)
+        enhanced = []
+        for frame in chunked_audio:
+            out = self.model.run(None,input_feed={"input_frame":frame.numpy(),"states":state,"atten_lim_db":atten_lim_db})
+            enhanced.append(torch.tensor(out[0]))
+            state = out[1]
+
+        enhanced_audio = torch.cat(enhanced).unsqueeze(0)  # [t] -> [1, t] typical mono format
+
+        d = self.fft_size - self.hop_size
+        enhanced_audio = enhanced_audio[:, d: orig_len + d]
+
+        return enhanced_audio.squeeze().numpy(), 48000
+
+
+class TextNorm():
+    def __init__(self):
+        my_pinyin = Pinyin(MyConverter())
+        self.pinyin_parser = my_pinyin.pinyin
+
+    def sil_type(self, time_s):
+        if round(time_s) < 0.4:
+            return ""
+        elif round(time_s) >= 0.4 and round(time_s) < 0.8:
+            return "#1"
+        elif round(time_s) >= 0.8 and round(time_s) < 1.5:
+            return "#2"
+        elif round(time_s) >= 1.5 and round(time_s) < 3.0:
+            return "#3"
+        elif round(time_s) >= 3.0:
+            return "#4"
+
+
+    def add_sil_raw(self, sub_list, start_time, end_time, target_transcript):
+        txt = []
+        txt_list = [x["word"] for x in sub_list]
+        sil = self.sil_type(sub_list[0]["start"])
+        if len(sil) > 0:
+            txt.append(sil)
+        txt.append(txt_list[0])
+        for i in range(1, len(sub_list)):
+            if sub_list[i]["start"] >= start_time and sub_list[i]["end"] <= end_time:
+                txt.append(target_transcript)
+                target_transcript = ""
+            else:
+                sil = self.sil_type(sub_list[i]["start"] - sub_list[i-1]["end"])
+                if len(sil) > 0:
+                    txt.append(sil)
+                txt.append(txt_list[i])
+        return ' '.join(txt)
+
+    def add_sil(self, sub_list, start_time, end_time, target_transcript, src_lang, tar_lang):
+        txts = []
+        txt_list = [x["word"] for x in sub_list]
+        sil = self.sil_type(sub_list[0]["start"])
+        if len(sil) > 0:
+            txts.append([src_lang, sil])
+
+        if sub_list[0]["start"] < start_time:
+            txts.append([src_lang, txt_list[0]])
+        for i in range(1, len(sub_list)):
+            if sub_list[i]["start"] >= start_time and sub_list[i]["end"] <= end_time:
+                txts.append([tar_lang, target_transcript])
+                target_transcript = ""
+            else:
+                sil = self.sil_type(sub_list[i]["start"] - sub_list[i-1]["end"])
+                if len(sil) > 0:
+                    txts.append([src_lang, sil])
+                txts.append([src_lang, txt_list[i]])
+                
+        target_txt = [txts[0]]
+        for txt in txts[1:]:
+            if txt[1] == "":
+                continue
+            if txt[0] != target_txt[-1][0]:
+                target_txt.append([txt[0], ""])
+            target_txt[-1][-1] += " " + txt[1]
+        
+        return target_txt
+
+
+    def get_prompt(self, sub_list, start_time, end_time, src_lang):
+        txts = []
+        txt_list = [x["word"] for x in sub_list]
+
+        if start_time <= sub_list[0]["start"]:
+            sil = self.sil_type(sub_list[0]["start"])
+            if len(sil) > 0:
+                txts.append([src_lang, sil])
+            txts.append([src_lang, txt_list[0]])
+        
+        for i in range(1, len(sub_list)):
+            # if sub_list[i]["start"] <= start_time and sub_list[i]["end"] <= end_time:
+            #     txts.append([tar_lang, target_transcript])
+            #     target_transcript = ""
+            if sub_list[i]["start"] >= start_time and sub_list[i]["end"] <= end_time:
+                sil = self.sil_type(sub_list[i]["start"] - sub_list[i-1]["end"])
+                if len(sil) > 0:
+                    txts.append([src_lang, sil])
+                txts.append([src_lang, txt_list[i]])
+
+        target_txt = [txts[0]]
+        for txt in txts[1:]:
+            if txt[1] == "":
+                continue
+            if txt[0] != target_txt[-1][0]:
+                target_txt.append([txt[0], ""])
+            target_txt[-1][-1] += " " + txt[1]
+        return target_txt
+
+
+    def txt2pinyin(self, text):
+        txts, phonemes = [], []
+        texts = re.split(r"(#\d)", text)
+        print("before norm: ", texts)
+        for text in texts:
+            if text in {'#1', '#2', '#3', '#4'}:
+                txts.append(text)
+                phonemes.append(text)
+                continue
+            text = NSWNormalizer(text.strip()).normalize()
+            
+            text_list = list(jieba.cut(text))
+            print("jieba cut: ", text, text_list)
+            for words in text_list:
+                if words in _PAUSE_SYMBOL:
+                    # phonemes.append('#2')
+                    phonemes[-1] += _PAUSE_SYMBOL[words]
+                    txts[-1] += words
+                elif re.search("[\u4e00-\u9fa5]+", words):
+                    pinyin = self.pinyin_parser(words, style=Style.TONE3, errors="ignore")
+                    new_pinyin = []
+                    for x in pinyin:
+                        x = "".join(x)
+                        if "#" not in x:
+                            new_pinyin.append(x)
+                        else:
+                            phonemes.append(words)
+                            continue
+                    new_pinyin = change_tone_in_bu_or_yi(words, new_pinyin) if len(words)>1 and words[-1] not in {"一","不"} else new_pinyin
+                    phoneme = get_phoneme_from_char_and_pinyin(words, new_pinyin)
+                    phonemes += phoneme
+                    txts += list(words)
+                elif re.search(r"[a-zA-Z]", words) or re.search(r"#[1-4]", words):
+                    phonemes.append(words)
+                    txts.append(words)
+                    # phonemes.append("#1")
+        # phones = " ".join(phonemes)
+        return txts, phonemes
+    
+
+
+def chunk_text(text, max_chars=135):
+    """
+    Splits the input text into chunks, each with a maximum number of characters.
+
+    Args:
+        text (str): The text to be split.
+        max_chars (int): The maximum number of characters per chunk.
+
+    Returns:
+        List[str]: A list of text chunks.
+    """
+    chunks = []
+    current_chunk = ""
+    # Split the text into sentences based on punctuation followed by whitespace
+    sentences = re.split(r"(?<=[;:,.!?])\s+|(?<=[；：，。！？])", text)
+
+    for sentence in sentences:
+        if len(current_chunk.encode("utf-8")) + len(sentence.encode("utf-8")) <= max_chars:
+            current_chunk += sentence + " " if sentence and len(sentence[-1].encode("utf-8")) == 1 else sentence
+        else:
+            if current_chunk:
+                chunks.append(current_chunk.strip())
+            current_chunk = sentence + " " if sentence and len(sentence[-1].encode("utf-8")) == 1 else sentence
+
+    if current_chunk:
+        chunks.append(current_chunk.strip())
+
+    return chunks
+
+
+class MMSAlignModel:
+    def __init__(self):
+        from torchaudio.pipelines import MMS_FA as bundle
+        self.mms_model = bundle.get_model()
+        self.mms_model.to(device)
+        self.mms_tokenizer = bundle.get_tokenizer()
+        self.mms_aligner = bundle.get_aligner()
+        self.text_normalizer = ur.Uroman() 
+
+
+    def text_normalization(self, text_list):
+        text_normalized = []
+        for word in text_list:
+            text_char = ''
+            for c in word:
+                if c.isalpha() or c=="'":
+                    text_char += c.lower()
+                elif c == "-":
+                    text_char += '*'
+            text_char = text_char if len(text_char) > 0 else "*"
+            text_normalized.append(text_char)
+        assert len(text_normalized) == len(text_list), f"normalized text len != raw text len: {len(text_normalized)} != {text_list}"
+        return text_normalized
+
+    def compute_alignments(self, waveform: torch.Tensor, tokens):
+        with torch.inference_mode():
+            emission, _ = self.mms_model(waveform.to(device))
+            token_spans = self.mms_aligner(emission[0], tokens)
+        return emission, token_spans
+
+
+    def align(self, data, wav):
+        waveform = load_wav(wav, 16000).unsqueeze(0)
+        raw_text = data['text'][0]
+        text = " ".join(data['text'][1]).replace("-", " ")
+        text = re.sub("\s+", " ", text)
+        text_normed = self.text_normalizer.romanize_string(text, lcode=data["lang"])
+        # text_normed = re.sub("[\d_.,!$£%?#−/]", '', text_normed)
+        fliter = re.compile("[^a-z^*^'^ ]")
+        text_normed = fliter.sub('', text_normed.lower())
+        text_normed = re.sub("\s+", " ", text_normed)
+        text_normed = text_normed.split()
+        assert len(text_normed) == len(raw_text), f"normalized text len != raw text len: {len(text_normed)} != {len(raw_text)}"
+        tokens = self.mms_tokenizer(text_normed)
+        with torch.inference_mode():
+            emission, _ = self.mms_model(waveform.to(device))
+            token_spans = self.mms_aligner(emission[0], tokens)
+        num_frames = emission.size(1)
+        ratio = waveform.size(1) / num_frames
+        res = []
+        for i in range(len(token_spans)):
+            score = round(sum([x.score for x in token_spans[i]]) / len(token_spans[i]), ndigits=3)
+            start = round(waveform.size(-1) * token_spans[i][0].start / num_frames / 16000, ndigits=3)
+            end = round(waveform.size(-1) * token_spans[i][-1].end / num_frames / 16000, ndigits=3)
+            res.append({"word": raw_text[i], "start": start, "end": end, "score": score})
+
+        res = {"lang":data["lang"], "start": 0, "end": round(waveform.shape[-1]/16000, ndigits=3), "text_raw":data["text_raw"], "text": text, "words": res}
+        return res
+
+
+class WhisperxModel:
+    def __init__(self, model_name):
+        from whisperx import load_model
+        from pathlib import Path
+        prompt = None  # "This might be a blend of Simplified Chinese and English speech, do not translate, only transcription be allowed."
+
+        # Prefer a local VAD model (to avoid network download / 301 issues)
+        vad_fp = Path(MODELS_PATH) / "whisperx-vad-segmentation.bin"
+        if not vad_fp.is_file():
+            logging.warning(
+                "Local whisperx VAD not found at %s, falling back to default download path.",
+                vad_fp,
+            )
+            vad_fp = None
+
+        self.model = load_model(
+            model_name,
+            ASR_DEVICE,
+            compute_type="float32",
+            asr_options={
+                "suppress_numerals": False,
+                "max_new_tokens": None,
+                "clip_timestamps": None,
+                "initial_prompt": prompt,
+                "append_punctuations": ".。,，!！?？:：、",
+                "hallucination_silence_threshold": None,
+                "multilingual": True,
+                "hotwords": None
+            },
+            vad_model_fp=str(vad_fp) if vad_fp is not None else None,
+        )
+
+    def transcribe(self, audio_info, lang=None):
+        audio = load_wav(audio_info).numpy()
+        if lang is None:
+            lang = self.model.detect_language(audio)
+        
+        segments = self.model.transcribe(audio, batch_size=8, language=lang)["segments"]
+        transcript = " ".join([segment["text"] for segment in segments])
+
+        if lang not in {'es','pt','zh','en','de','fr','it', 'ar', 'ru', 'ja', 'ko', 'hi', 'th', 'id', 'vi'}:
+            lang = langid.classify(transcript)[0]
+            segments = self.model.transcribe(audio, batch_size=8, language=lang)["segments"]
+            transcript = " ".join([segment["text"] for segment in segments])
+        logging.debug(f"whisperx: {segments}")
+        
+        transcript = zhconv.convert(transcript, 'zh-hans')
+        transcript = transcript.replace("-", " ")
+        transcript = re.sub(_whitespace_re, " ", transcript)
+        transcript = transcript[1:] if transcript[0] == " " else transcript
+        segments = {'lang':lang, 'text_raw':transcript}
+        if lang == "zh":
+            segments["text"] = text_norm.txt2pinyin(transcript)
+        else:
+            transcript = replace_numbers_with_words(transcript, lang=lang).split(' ')
+            segments["text"] = (transcript, transcript)
+
+        return align_model.align(segments, audio_info)
+
 
 def load_wav(audio_info, sr=16000, channel=1):
-    print("load audio:", audio_info)
-    audio, raw_sr = torchaudio.load(audio_info)
+    raw_sr, audio = audio_info
     audio = audio.T if len(audio.shape) > 1 and audio.shape[1] == 2 else audio
-    audio = audio / torch.max(torch.abs(audio))
-    audio = audio.squeeze().float()
+    audio = audio / np.max(np.abs(audio))
+    audio = torch.from_numpy(audio).squeeze().float()
     if channel == 1 and len(audio.shape) == 2:  # stereo to mono
         audio = audio.mean(dim=0, keepdim=True)
     elif channel == 2 and len(audio.shape) == 1:
         audio = torch.stack((audio, audio)) # mono to stereo
-    if raw_sr != sr:
+    if raw_sr != sr: 
         audio = torchaudio.functional.resample(audio.squeeze(), raw_sr, sr)
     audio = torch.clip(audio, -0.999, 0.999).squeeze()
     return audio
 
 
-def denoise(audio_info):
-    save_path = "./denoised_audio.wav"
-    denoised_audio, sr = denoise_model.denoise(audio_info)
-    sf.write(save_path, denoised_audio, sr, format='wav', subtype='PCM_24')
-    print("save denoised audio:", save_path)
-    return save_path
+def update_word_time(lst, cut_time, edit_start, edit_end):
+    for i in range(len(lst)):
+        lst[i]["start"] = round(lst[i]["start"] - cut_time, ndigits=3)
+        lst[i]["end"] = round(lst[i]["end"] - cut_time,  ndigits=3)
+    edit_start = max(round(edit_start - cut_time, ndigits=3), 0)
+    edit_end = round(edit_end - cut_time, ndigits=3)
+    return lst, edit_start, edit_end
 
-def cancel_denoise(audio_info):
-    return audio_info
+
+# def update_word_time2(lst, cut_time, edit_start, edit_end):
+#     for i in range(len(lst)):
+#         lst[i]["start"] = round(lst[i]["start"] + cut_time, ndigits=3)
+#     return lst, edit_start, edit_end
 
 
-def get_checkpoints_project(project_name=None, is_gradio=True):
-    """Get available checkpoint files"""
-    checkpoint_dir = [str(CKPTS_ROOT)]
-    # Remote ckpt locations on HF (used when local ckpts are not present)
-    remote_ckpts = {
-        "multilingual_grl": f"{HF_PRETRAINED_ROOT}/ckpts/multilingual_grl/multilingual_grl.safetensors",
-        "multilingual_prosody": f"{HF_PRETRAINED_ROOT}/ckpts/multilingual_prosody/multilingual_prosody.safetensors",
+def get_audio_slice(audio, words_info, start_time, end_time, max_len=10, sr=16000, code_sr=50):
+    audio_dur = audio.shape[-1] / sr
+    sub_list = []
+    # 如果尾部小于5s则保留后面全部，并截取前半段音频
+    if audio_dur - end_time <= max_len/2:
+        for word in reversed(words_info):
+            if word['start'] > start_time or audio_dur - word['start'] < max_len:
+                sub_list = [word] + sub_list
+
+    # 如果头部小于5s则保留前面全部，并截取后半段音频
+    elif start_time <=max_len/2:
+        for word in words_info:
+            if word['end'] < max(end_time, max_len):
+                sub_list += [word]
+                
+    # 如果前后都大于5s，则前后各留5s
+    else:
+        for word in words_info:
+            if word['start'] > start_time - max_len/2 and word['end'] < end_time + max_len/2:
+                sub_list += [word]
+    audio = audio.squeeze()
+
+    start = int(sub_list[0]['start']*sr)
+    end = int(sub_list[-1]['end']*sr)
+    # print("wav cuts:", start, end, (end-start) % int(sr/code_sr))
+    end -= (end-start) % int(sr/code_sr) # chunk取整
+
+    sub_list, start_time, end_time = update_word_time(sub_list, sub_list[0]['start'], start_time, end_time)
+    audio = audio.squeeze()
+    # print("after update_word_time:", sub_list, start_time, end_time, (end-start)/sr)
+
+    return (audio[:start], audio[start:end], audio[end:]), (sub_list, start_time, end_time)
+
+
+def load_models(lemas_model_name, whisper_model_name, alignment_model_name, denoise_model_name):  # , audiosr_name):
+
+    global transcribe_model, align_model, denoise_model, text_norm, tts_edit_model
+    # if voicecraft_model:
+    #     del denoise_model
+    #     del transcribe_model
+    #     del align_model
+    #     del voicecraft_model
+    #     del audiosr
+    torch.cuda.empty_cache()
+    gc.collect()
+
+    if denoise_model_name == "UVR5":
+        # Prefer the generic MODELS_PATH root for denoiser assets so that
+        # HF Spaces (where pretrained models are often mounted separately)
+        # and local runs share the same layout.
+        denoise_root = MODELS_PATH  # e.g. "./pretrained_models" or env override
+        denoise_model = UVR5(os.path.join(denoise_root, "uvr5"))
+    elif denoise_model_name == "DeepFilterNet":
+        denoise_model = DeepFilterNet("./pretrained_models/denoiser_model.onnx")
+
+    if alignment_model_name == "MMS":
+        align_model = MMSAlignModel()
+    else:
+        align_model = WhisperxAlignModel()
+    
+    text_norm = TextNorm()
+
+    transcribe_model = WhisperxModel(whisper_model_name)
+
+    # Load LEMAS-TTS editing model (selected multilingual variant)
+    from pathlib import Path
+
+    ckpt_dir = Path(CKPTS_ROOT) / lemas_model_name
+    ckpt_candidates = sorted(
+        list(ckpt_dir.glob("*.safetensors")) + list(ckpt_dir.glob("*.pt"))
+    )
+    if not ckpt_candidates:
+        raise gr.Error(f"No LEMAS-TTS ckpt found under {ckpt_dir}")
+    ckpt_file = str(ckpt_candidates[-1])
+
+    vocab_file = Path(PRETRAINED_ROOT) / "data" / lemas_model_name / "vocab.txt"
+    if not vocab_file.is_file():
+        raise gr.Error(f"Vocab file not found: {vocab_file}")
+
+    prosody_cfg = Path(CKPTS_ROOT) / "prosody_encoder" / "pretssel_cfg.json"
+    prosody_ckpt = Path(CKPTS_ROOT) / "prosody_encoder" / "prosody_encoder_UnitY2.pt"
+
+    # Decide whether to enable the prosody encoder:
+    # - multilingual_prosody: True (if assets exist)
+    # - multilingual_grl: False (GRL-only variant)
+    # - others: fall back to presence of assets.
+    if lemas_model_name.endswith("prosody"):
+        use_prosody = prosody_cfg.is_file() and prosody_ckpt.is_file()
+    elif lemas_model_name.endswith("grl"):
+        use_prosody = False
+    else:
+        use_prosody = prosody_cfg.is_file() and prosody_ckpt.is_file()
+
+    tts_edit_model = TTS(
+        model=lemas_model_name,
+        ckpt_file=ckpt_file,
+        vocab_file=str(vocab_file),
+        device=device,
+        use_prosody_encoder=use_prosody,
+        prosody_cfg_path=str(prosody_cfg) if use_prosody else "",
+        prosody_ckpt_path=str(prosody_ckpt) if use_prosody else "",
+        ode_method="euler",
+        use_ema=True,
+        frontend="phone",
+    )
+    logging.info(f"Loaded LEMAS-TTS edit model from {ckpt_file}")
+
+    return gr.Accordion()
+
+
+def get_transcribe_state(segments):
+    logging.info("===========After Align===========")
+    logging.info(segments)
+    return {
+        "segments": segments,
+        "transcript": segments["text_raw"],
+        "words_info": segments["words"], 
+        "transcript_with_start_time": " ".join([f"{word['start']} {word['word']}" for word in segments["words"]]),
+        "transcript_with_end_time": " ".join([f"{word['word']} {word['end']}" for word in segments["words"]]),
+        "word_bounds": [f"{word['start']} {word['word']} {word['end']}" for word in segments["words"]]
     }
 
-    if project_name is None:
-        # Look for checkpoints in local directory
-        files_checkpoints = []
-        for path in checkpoint_dir:
-            if os.path.isdir(path):
-                files_checkpoints.extend(glob(os.path.join(path, "**/*.pt"), recursive=True))
-                files_checkpoints.extend(glob(os.path.join(path, "**/*.safetensors"), recursive=True))
-                break
-        # Fallback to remote ckpts if none found locally
-        if not files_checkpoints:
-            files_checkpoints = list(remote_ckpts.values())
+
+def transcribe(seed, audio_info):
+    if transcribe_model is None:
+        raise gr.Error("Transcription model not loaded")
+    seed_everything(seed)
+
+    segments = transcribe_model.transcribe(audio_info)
+    state = get_transcribe_state(segments)
+
+    return [
+        state["transcript"], state["transcript_with_start_time"], state["transcript_with_end_time"],
+        # gr.Dropdown(value=state["word_bounds"][-1], choices=state["word_bounds"], interactive=True), # prompt_to_word
+        gr.Dropdown(value=state["word_bounds"][0], choices=state["word_bounds"], interactive=True), # edit_from_word
+        gr.Dropdown(value=state["word_bounds"][-1], choices=state["word_bounds"], interactive=True), # edit_to_word
+        state
+    ]
+
+def align(transcript, audio_info, state):
+    lang = state["segments"]["lang"]
+    # print("realign: ", transcript, state)
+    transcript = re.sub(_whitespace_re, " ", transcript)
+    transcript = transcript[1:] if transcript[0] == " " else transcript
+    segments = {'lang':lang, 'text':transcript, 'text_raw':transcript}
+    if lang == "zh":
+        segments["text"] = text_norm.txt2pinyin(transcript)
     else:
-        files_checkpoints = []
-        if os.path.isdir(checkpoint_dir[0]):
-            files_checkpoints = glob(os.path.join(checkpoint_dir[0], project_name, "*.pt"))
-            files_checkpoints.extend(glob(os.path.join(checkpoint_dir[0], project_name, "*.safetensors")))
-        # If no local ckpts for this project, try remote mapping
-        if not files_checkpoints:
-            ckpt = remote_ckpts.get(project_name)
-            files_checkpoints = [ckpt] if ckpt is not None else []
-    print("files_checkpoints:", project_name, files_checkpoints)
-    # Separate pretrained and regular checkpoints
-    pretrained_checkpoints = [f for f in files_checkpoints if "pretrained_" in os.path.basename(f)]
-    regular_checkpoints = [
-        f
-        for f in files_checkpoints
-        if "pretrained_" not in os.path.basename(f) and "model_last.pt" not in os.path.basename(f)
+        transcript = replace_numbers_with_words(transcript)
+        segments["text"] = (transcript.split(' '), transcript.split(' '))
+    # print("text:", segments["text"])
+    segments = align_model.align(segments, audio_info)
+    
+    state = get_transcribe_state(segments)
+
+    return [
+        state["transcript"], state["transcript_with_start_time"], state["transcript_with_end_time"],
+        # gr.Dropdown(value=state["word_bounds"][-1], choices=state["word_bounds"], interactive=True), # prompt_to_word
+        gr.Dropdown(value=state["word_bounds"][0], choices=state["word_bounds"], interactive=True), # edit_from_word
+        gr.Dropdown(value=state["word_bounds"][-1], choices=state["word_bounds"], interactive=True), # edit_to_word
+        state
     ]
-    last_checkpoint = [f for f in files_checkpoints if "model_last.pt" in os.path.basename(f)]
 
-    # Sort regular checkpoints by number
-    try:
-        regular_checkpoints = sorted(
-            regular_checkpoints, key=lambda x: int(os.path.basename(x).split("_")[1].split(".")[0])
-        )
-    except (IndexError, ValueError):
-        regular_checkpoints = sorted(regular_checkpoints)
 
-    # Combine in order: pretrained, regular, last
-    files_checkpoints = pretrained_checkpoints + regular_checkpoints + last_checkpoint
+def denoise(audio_info):
+    denoised_audio, sr = denoise_model.denoise(audio_info)
+    denoised_audio = denoised_audio # .squeeze().numpy()
+    return (sr, denoised_audio)
+
+def cancel_denoise(audio_info):
+    return audio_info
 
-    select_checkpoint = None if not files_checkpoints else files_checkpoints[-1]
+def get_output_audio(audio_tensors, sr):
+    result = torch.cat(audio_tensors, -1)
+    result = result.squeeze().cpu().numpy()
+    result = (result * np.iinfo(np.int16).max).astype(np.int16)
+    print("save result:", result.shape)
+    # wavfile.write(os.path.join(TMP_PATH, "output.wav"), sr, result)
+    return (int(sr), result)
+    
 
-    if is_gradio:
-        return gr.update(choices=files_checkpoints, value=select_checkpoint)
+def get_edit_audio_part(audio_info, edit_start, edit_end):
+    sr, raw_wav = audio_info
+    raw_wav = raw_wav[int(edit_start*sr):int(edit_end*sr)]
+    return (sr, raw_wav)
 
-    return files_checkpoints, select_checkpoint
 
+def crossfade_concat(chunk1, chunk2, overlap):
+    # 计算淡入和淡出系数
+    fade_out = torch.cos(torch.linspace(0, torch.pi / 2, overlap)) ** 2
+    fade_in = torch.cos(torch.linspace(torch.pi / 2, 0, overlap)) ** 2
+    chunk2[:overlap] = chunk1[-overlap:] * fade_out + chunk2[:overlap] * fade_in
+    chunk = torch.cat((chunk1[:-overlap], chunk2), dim=0)
+    return chunk
 
-def get_available_projects():
-    """Get available project names from data directory"""
-    data_paths = [
-        str(Path(PRETRAINED_ROOT) / "data"),
-    ]
+def replace_numbers_with_words(sentence, lang="en"):
+    sentence = re.sub(r'(\d+)', r' \1 ', sentence) # add spaces around numbers
+    def replace_with_words(match):
+        num = match.group(0)
+        try:
+            return num2words(num, lang=lang) # Convert numbers to words
+        except:
+            return num # In case num2words fails (unlikely with digits but just to be safe)
+    return re.sub(r'\b\d+\b', replace_with_words, sentence) # Regular expression that matches numbers
+
+
+def run(seed, nfe_step, speed, cfg_strength, sway_sampling_coef, ref_ratio,
+        audio_info, denoised_audio, transcribe_state, transcript, smart_transcript,
+        mode, start_time, end_time,
+        split_text, selected_sentence, audio_tensors):
+    if tts_edit_model is None:
+        raise gr.Error("LEMAS-TTS edit model not loaded")
+    if smart_transcript and (transcribe_state is None):
+        raise gr.Error("Can't use smart transcript: whisper transcript not found")
+
+    # if mode == "Rerun":
+    #     colon_position = selected_sentence.find(':')
+    #     selected_sentence_idx = int(selected_sentence[:colon_position])
+    #     sentences = [selected_sentence[colon_position + 1:]]
+
+    # Choose base audio (denoised if duration matches)
+    audio_base = audio_info
+    audio_dur = round(audio_info[1].shape[0] / audio_info[0], ndigits=3)
+    if denoised_audio is not None:
+        denoised_dur = round(denoised_audio[1].shape[0] / denoised_audio[0], ndigits=3)
+        if audio_dur == denoised_dur or (
+            denoised_audio[0] != audio_info[0] and abs(audio_dur - denoised_dur) < 0.1
+        ):
+            audio_base = denoised_audio
+            logging.info("use denoised audio")
+
+    raw_sr, raw_wav = audio_base
+    print("audio_dur: ", audio_dur, raw_sr, raw_wav.shape, start_time, end_time)
     
-    project_list = []
-    for data_path in data_paths:
-        if os.path.isdir(data_path):
-            for folder in os.listdir(data_path):
-                path_folder = os.path.join(data_path, folder)
-                if "test" not in folder:
-                    project_list.append(folder)
+    # Build target text by replacing the selected span with `transcript`
+    words = transcribe_state["words_info"]
+    if not words:
+        raise gr.Error("No word-level alignment found; please run Transcribe first.")
+
+    start_time = float(start_time)
+    end_time = float(end_time)
+    if end_time <= start_time:
+        raise gr.Error("Edit end time must be greater than start time.")
+
+    # Find word indices covering the selected region
+    start_idx = 0
+    for i, w in enumerate(words):
+        if w["end"] > start_time:
+            start_idx = i
             break
-    # Fallback: if no local data dir, default to known HF projects
-    if not project_list:
-        project_list = ["multilingual_grl", "multilingual_prosody"]
-    project_list.sort()
-    print("project_list:", project_list)
-    return project_list
-
-
-def infer(
-    project, file_checkpoint, exp_name, ref_text, ref_audio, denoise_audio, gen_text, nfe_step, use_ema, separate_langs, frontend, speed, cfg_strength, use_acc_grl, ref_ratio, no_ref_audio, sway_sampling_coef, use_prosody_encoder, seed
-):
-    global last_checkpoint, last_device, tts_api, last_ema
-
-    # Resolve checkpoint path (local or HF URL)
-    ckpt_path = file_checkpoint
-    if isinstance(ckpt_path, str) and ckpt_path.startswith("hf://"):
-        try:
-            ckpt_resolved = str(cached_path(ckpt_path))
-        except Exception as e:
-            traceback.print_exc()
-            return None, f"Error downloading checkpoint: {str(e)}", ""
-    else:
-        ckpt_resolved = ckpt_path
 
-    if not os.path.isfile(ckpt_resolved):
-        return None, "Checkpoint not found!", ""
+    end_idx = len(words)
+    for i in range(len(words) - 1, -1, -1):
+        if words[i]["start"] < end_time:
+            end_idx = i + 1
+            break
+    if end_idx <= start_idx:
+        end_idx = min(start_idx + 1, len(words))
+
+    word_start_sec = float(words[start_idx]["start"])
+    word_end_sec = float(words[end_idx - 1]["end"])
+
+    # Edit span in seconds (relative to full utterance)
+    edit_start = max(0.0, word_start_sec - 0.1)
+    edit_end = min(word_end_sec + 0.1, audio_dur)
+    parts_to_edit = [(edit_start, edit_end)]
+
+    display_text = transcribe_state["segments"]["text_raw"].strip()
+    txt_list = display_text.split(" ") if display_text else [w["word"] for w in words]
+
+    prefix = " ".join(txt_list[:start_idx]).strip()
+    suffix = " ".join(txt_list[end_idx:]).strip()
+    new_phrase = transcript.strip()
+
+    pieces = []
+    if prefix:
+        pieces.append(prefix)
+    if new_phrase:
+        pieces.append(new_phrase)
+    if suffix:
+        pieces.append(suffix)
+    target_text = " ".join(pieces)
+
+    logging.info(
+        "target_text: %s (start_idx=%d, end_idx=%d, parts_to_edit=%s)",
+        target_text,
+        start_idx,
+        end_idx,
+        parts_to_edit,
+    )
 
-    if denoise_audio:
-        ref_audio = denoise_audio
+    # Prepare audio for LEMAS-TTS editing (mono, target SR)
+    segment_audio = load_wav(audio_base, sr=tts_edit_model.target_sample_rate)
+
+    seed_val = None if seed == -1 else int(seed)
+
+    # Decide whether to use prosody encoder at inference based on how TTS was built
+    use_prosody_flag = bool(getattr(tts_edit_model, "use_prosody_encoder", False))
+
+    wav_out, _ = gen_wav_multilingual(
+        tts_edit_model,
+        segment_audio,
+        tts_edit_model.target_sample_rate,
+        target_text,
+        parts_to_edit,
+        speed=float(speed),
+        nfe_step=int(nfe_step),
+        cfg_strength=float(cfg_strength),
+        sway_sampling_coef=float(sway_sampling_coef),
+        ref_ratio=float(ref_ratio),
+        no_ref_audio=False,
+        use_acc_grl=False,
+        use_prosody_encoder_flag=use_prosody_flag,
+        seed=seed_val,
+    )
 
-    device_test = device  # Use the global device
+    wav_np = wav_out.cpu().numpy()
+    wav_np = np.clip(wav_np, -0.999, 0.999)
+    wav_int16 = (wav_np * np.iinfo(np.int16).max).astype(np.int16)
+    out_sr = int(tts_edit_model.target_sample_rate)
 
-    if last_checkpoint != ckpt_resolved or last_device != device_test or last_ema != use_ema or tts_api is None:
-        if last_checkpoint != ckpt_resolved:
-            last_checkpoint = ckpt_resolved
+    output_audio = (out_sr, wav_int16)
+    sentences = [f"0: {target_text}"]
+    audio_tensors = [torch.from_numpy(wav_np)]
 
-        if last_device != device_test:
-            last_device = device_test
+    component = gr.Dropdown(choices=sentences, value=sentences[0])
+    return output_audio, target_text, component, audio_tensors
+        
 
-        if last_ema != use_ema:
-            last_ema = use_ema
+def update_input_audio(audio_info):
+    if audio_info is None:
+        return 0, 0, 0
+    elif type(audio_info) is str:
+        info = torchaudio.info(audio_path)
+        max_time = round(info.num_frames / info.sample_rate, 2)
+    elif type(audio_info) is tuple:
+        max_time = round(audio_info[1].shape[0] / audio_info[0], 2)
+    return [
+        # gr.Slider(maximum=max_time, value=max_time),
+        gr.Slider(maximum=max_time, value=0),
+        gr.Slider(maximum=max_time, value=max_time),
+    ]
 
-        # Automatically enable prosody encoder when using the prosody checkpoint
-        use_prosody_encoder = True if "prosody" in str(ckpt_resolved) else False
 
-        # Resolve vocab file (local)
-        local_vocab = Path(PRETRAINED_ROOT) / "data" / project / "vocab.txt"
-        if not local_vocab.is_file():
-            return None, "Vocab file not found!", ""
-        vocab_file = str(local_vocab)
+def change_mode(mode):
+    # tts_mode_controls, edit_mode_controls, edit_word_mode, split_text, long_tts_sentence_editor
+    return [
+        gr.Group(visible=mode != "Edit"),
+        gr.Group(visible=mode == "Edit"),
+        gr.Radio(visible=mode == "Edit"),
+        gr.Radio(visible=mode == "Long TTS"),
+        gr.Group(visible=mode == "Long TTS"),
+    ]
 
-        # Resolve prosody encoder config & weights (local)
-        local_prosody_cfg = Path(CKPTS_ROOT) / "prosody_encoder" / "pretssel_cfg.json"
-        local_prosody_ckpt = Path(CKPTS_ROOT) / "prosody_encoder" / "prosody_encoder_UnitY2.pt"
-        if not local_prosody_cfg.is_file() or not local_prosody_ckpt.is_file():
-            return None, "Prosody encoder files not found!", ""
-        prosody_cfg_path = str(local_prosody_cfg)
-        prosody_ckpt_path = str(local_prosody_ckpt)
 
-        try:
-            tts_api = TTS(
-                model=exp_name,
-                ckpt_file=ckpt_resolved,
-                vocab_file=vocab_file,
-                device=device_test,
-                use_ema=use_ema,
-                frontend=frontend,
-                use_prosody_encoder=use_prosody_encoder,
-                prosody_cfg_path=prosody_cfg_path,
-                prosody_ckpt_path=prosody_ckpt_path,
-            )
-        except Exception as e:
-            traceback.print_exc()
-            return None, f"Error loading model: {str(e)}", ""
-
-        print("Model loaded >>", device_test, file_checkpoint, use_ema)
-
-    if seed == -1:  # -1 used for random
-        seed = None
-
-    try:
-        with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as f:
-            tts_api.infer(
-                ref_file=ref_audio,
-                ref_text=ref_text.strip(),
-                gen_text=gen_text.strip(),
-                nfe_step=nfe_step,
-                separate_langs=separate_langs,
-                speed=speed,
-                cfg_strength=cfg_strength,
-                sway_sampling_coef=sway_sampling_coef,
-                use_acc_grl=use_acc_grl,
-                ref_ratio=ref_ratio,
-                no_ref_audio=no_ref_audio,
-                use_prosody_encoder=use_prosody_encoder,
-                file_wave=f.name,
-                seed=seed,
-            )
-            return f.name, f"Device: {tts_api.device}", str(tts_api.seed)
-    except Exception as e:
-        traceback.print_exc()
-        return None, f"Inference error: {str(e)}", ""
-
-
-def get_gpu_stats():
-    """Get GPU statistics"""
-    gpu_stats = ""
-
-    if torch.cuda.is_available():
-        gpu_count = torch.cuda.device_count()
-        for i in range(gpu_count):
-            gpu_name = torch.cuda.get_device_name(i)
-            gpu_properties = torch.cuda.get_device_properties(i)
-            total_memory = gpu_properties.total_memory / (1024**3)  # in GB
-            allocated_memory = torch.cuda.memory_allocated(i) / (1024**2)  # in MB
-            reserved_memory = torch.cuda.memory_reserved(i) / (1024**2)  # in MB
-
-            gpu_stats += (
-                f"GPU {i} Name: {gpu_name}\n"
-                f"Total GPU memory (GPU {i}): {total_memory:.2f} GB\n"
-                f"Allocated GPU memory (GPU {i}): {allocated_memory:.2f} MB\n"
-                f"Reserved GPU memory (GPU {i}): {reserved_memory:.2f} MB\n\n"
-            )
-    elif torch.xpu.is_available():
-        gpu_count = torch.xpu.device_count()
-        for i in range(gpu_count):
-            gpu_name = torch.xpu.get_device_name(i)
-            gpu_properties = torch.xpu.get_device_properties(i)
-            total_memory = gpu_properties.total_memory / (1024**3)  # in GB
-            allocated_memory = torch.xpu.memory_allocated(i) / (1024**2)  # in MB
-            reserved_memory = torch.xpu.memory_reserved(i) / (1024**2)  # in MB
-
-            gpu_stats += (
-                f"GPU {i} Name: {gpu_name}\n"
-                f"Total GPU memory (GPU {i}): {total_memory:.2f} GB\n"
-                f"Allocated GPU memory (GPU {i}): {allocated_memory:.2f} MB\n"
-                f"Reserved GPU memory (GPU {i}): {reserved_memory:.2f} MB\n\n"
-            )
-    elif torch.backends.mps.is_available():
-        gpu_count = 1
-        gpu_stats += "MPS GPU\n"
-        total_memory = psutil.virtual_memory().total / (
-            1024**3
-        )  # Total system memory (MPS doesn't have its own memory)
-        allocated_memory = 0
-        reserved_memory = 0
-
-        gpu_stats += (
-            f"Total system memory: {total_memory:.2f} GB\n"
-            f"Allocated GPU memory (MPS): {allocated_memory:.2f} MB\n"
-            f"Reserved GPU memory (MPS): {reserved_memory:.2f} MB\n"
-        )
+def load_sentence(selected_sentence, audio_tensors):
+    if selected_sentence is None:
+        return None
+    colon_position = selected_sentence.find(':')
+    selected_sentence_idx = int(selected_sentence[:colon_position])
+    # Use LEMAS-TTS target sample rate if available, otherwise default to 16000
+    sr = getattr(tts_edit_model, "target_sample_rate", 16000)
+    return get_output_audio([audio_tensors[selected_sentence_idx]], sr)
 
-    else:
-        gpu_stats = "No GPU available"
 
-    return gpu_stats
+def update_bound_word(is_first_word, selected_word, edit_word_mode):
+    if selected_word is None:
+        return None
 
+    word_start_time = float(selected_word.split(' ')[0])
+    word_end_time = float(selected_word.split(' ')[-1])
+    if edit_word_mode == "Replace half":
+        bound_time = (word_start_time + word_end_time) / 2
+    elif is_first_word:
+        bound_time = word_start_time
+    else:
+        bound_time = word_end_time
 
-def get_cpu_stats():
-    """Get CPU statistics"""
-    cpu_usage = psutil.cpu_percent(interval=1)
-    memory_info = psutil.virtual_memory()
-    memory_used = memory_info.used / (1024**2)
-    memory_total = memory_info.total / (1024**2)
-    memory_percent = memory_info.percent
+    return bound_time
 
-    pid = os.getpid()
-    process = psutil.Process(pid)
-    nice_value = process.nice()
 
-    cpu_stats = (
-        f"CPU Usage: {cpu_usage:.2f}%\n"
-        f"System Memory: {memory_used:.2f} MB used / {memory_total:.2f} MB total ({memory_percent}% used)\n"
-        f"Process Priority (Nice value): {nice_value}"
-    )
+def update_bound_words(from_selected_word, to_selected_word, edit_word_mode):
+    return [
+        update_bound_word(True, from_selected_word, edit_word_mode),
+        update_bound_word(False, to_selected_word, edit_word_mode),
+    ]
 
-    return cpu_stats
 
+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 = ""
+
+demo_text = {
+    "TTS": {
+        "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."
+    },
+    "Edit": {
+        "smart": "Just write it line-by-line.",
+        "regular": "照片、医疗记录、神经重塑的易损性，这是某种数据库啊！还有PRELESS的脑部扫描、生物管型、神经重塑."
+    },
+    "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!"
+    }
+}
 
-def get_combined_stats():
-    """Get combined system stats"""
-    gpu_stats = get_gpu_stats()
-    cpu_stats = get_cpu_stats()
-    combined_stats = f"### GPU Stats\n{gpu_stats}\n\n### CPU Stats\n{cpu_stats}"
-    return combined_stats
+all_demo_texts = {vv for k, v in demo_text.items() for kk, vv in v.items()}
 
+demo_words = ['0.069 Gwynplain 0.611', '0.671 had, 0.912', '0.952 besides, 1.414', '1.494 for 1.634', '1.695 his 1.835', '1.915 work 2.136', '2.196 and 2.297', '2.337 for 2.517', '2.557 his 2.678', '2.758 feats 3.019', '3.079 of 3.139', '3.2 strength, 3.561', '4.022 round 4.263', '4.303 his 4.444', '4.524 neck 4.705', '4.745 and 4.825', '4.905 over 5.086', '5.146 his 5.266', '5.307 shoulders, 5.768', '6.23 an 6.33', '6.531 esclavine 7.133', '7.213 of 7.293', '7.353 leather. 7.614']
 
-# Create Gradio interface
-with gr.Blocks(title="LEMAS-TTS Inference") as app:
-    gr.Markdown(
-        """
-        # Zero-Shot TTS
+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}]
 
-        Set seed to -1 for random generation.
-        """
-    )
-    with gr.Accordion("Model configuration", open=False):
-    # Model configuration
-        with gr.Row():
-            exp_name = gr.Radio(
-                label="Model",
-                choices=["multilingual_grl", "multilingual_prosody"],
-                value="multilingual_grl",
-                visible=False,
-            )
-        # Project selection
-        available_projects = get_available_projects()
 
-        # Get initial checkpoints
-        list_checkpoints, checkpoint_select = get_checkpoints_project(available_projects[0] if available_projects else None, False)
+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
 
-        with gr.Row():
-            with gr.Column(scale=1):
-                # load_models_btn = gr.Button(value="Load models")
-                cm_project = gr.Dropdown(
-                    choices=available_projects, 
-                    value=available_projects[0] if available_projects else None,
-                    label="Project", 
-                    allow_custom_value=True, 
-                    scale=4
-                )
-                
-            with gr.Column(scale=5):
-                cm_checkpoint = gr.Dropdown(
-                    choices=list_checkpoints, value=checkpoint_select, label="Checkpoints", allow_custom_value=True # scale=4, 
-)
-            bt_checkpoint_refresh = gr.Button("Refresh", scale=1)
+    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():
-            ch_use_ema = gr.Checkbox(label="Use EMA", visible=False, value=True, scale=2, info="Turn off at early stage might offer better results")
-            frontend = gr.Radio(label="Frontend", visible=False, choices=["phone", "char", "bpe"], value="phone", scale=3)
-            separate_langs = gr.Checkbox(label="Separate Languages", visible=False, value=True, scale=2, info="separate language tokens")
+            with gr.Column(scale=2):
+                load_models_btn = gr.Button(value="Load models")
+            with gr.Column(scale=5):
+                with gr.Accordion("Select models", open=False) as models_selector:
+                    # For LEMAS-TTS editing, we expose a simple model selector
+                    # between the two multilingual variants.
+                    with gr.Row():
+                        lemas_model_choice = gr.Radio(
+                            label="Edit Model",
+                            choices=["multilingual_grl", "multilingual_prosody"],
+                            value="multilingual_grl",
+                            interactive=True,
+                            scale=3,
+                        )
+                        denoise_model_choice = gr.Radio(label="Denoise Model", scale=2, value="UVR5", choices=["UVR5", "DeepFilterNet"]) # "830M", "330M_TTSEnhanced", "830M_TTSEnhanced"])
+                        # whisper_backend_choice = gr.Radio(label="Whisper backend", value="", choices=["whisperX", "whisper"])
+                        whisper_model_choice = gr.Radio(label="Whisper model", scale=3, value="medium", choices=["base", "small", "medium", "large"])
+                        align_model_choice = gr.Radio(label="Forced alignment model", scale=2, value="MMS", choices=["whisperX", "MMS"], visible=False)
 
-        # Inference parameters
         with gr.Row():
-            nfe_step = gr.Number(label="NFE Step", scale=1, value=64)
-            speed = gr.Slider(label="Speed", scale=3, value=1.0, minimum=0.5, maximum=1.5, step=0.1)
-            cfg_strength = gr.Slider(label="CFG Strength", scale=2, value=5.0, minimum=0.0, maximum=10.0, step=1)
-            sway_sampling_coef = gr.Slider(label="Sway Sampling Coef", scale=2, value=3, minimum=2, maximum=5, step=0.1)
-            ref_ratio = gr.Slider(label="Ref Ratio", scale=2, value=1.0, minimum=0.0, maximum=1.0, step=0.1)
-            no_ref_audio = gr.Checkbox(label="No Reference Audio", visible=False, value=False, scale=1, info="No mel condition")
-            use_acc_grl = gr.Checkbox(label="Use accent grl condition", visible=False, value=True, scale=1, info="Use accent grl condition")
-            use_prosody_encoder = gr.Checkbox(label="Use prosody encoder", visible=False, value=False, scale=1, info="Use prosody encoder")
-            seed = gr.Number(label="Random Seed", scale=1, value=-1, minimum=-1)
-
-
-    # Input fields
-    ref_text = gr.Textbox(label="Reference Text", placeholder="Enter the text for the reference audio...")
-    ref_audio = gr.Audio(label="Reference Audio", type="filepath", interactive=True, show_download_button=True, editable=True)
+            with gr.Column(scale=2):
+                # Use a numpy waveform as default value to avoid Gradio's
+                # InvalidPathError with local filesystem paths.
+                _demo_value = None
+                demo_candidates = [
+                    os.path.join(DEMO_PATH, "V-00013_en-US.wav"),
+                    os.path.join(os.path.dirname(__file__), "..", "VoiceCraft", "demo", "V-00013_en-US.wav"),
+                ]
+                for demo_path in demo_candidates:
+                    try:
+                        if not os.path.isfile(demo_path):
+                            continue
+                        _demo_wav, _demo_sr = torchaudio.load(demo_path)
+                        if _demo_wav.dim() > 1 and _demo_wav.shape[0] > 1:
+                            _demo_wav = _demo_wav.mean(dim=0, keepdim=True)
+                        _demo_value = (_demo_sr, _demo_wav.squeeze(0).numpy())
+                        break
+                    except Exception:
+                        continue
+
+                input_audio = gr.Audio(
+                    value=_demo_value,
+                    label="Input Audio",
+                    interactive=True,
+                    type="numpy",
+                )
 
+                with gr.Row():
+                    transcribe_btn = gr.Button(value="Transcribe")
+                    align_btn = gr.Button(value="ReAlign")
+                with gr.Group():
+                    original_transcript = gr.Textbox(label="Original transcript", lines=5, interactive=True, value=demo_original_transcript,
+                                                    info="Use whisperx model to get the transcript. Fix and align it if necessary.")
+                    with gr.Accordion("Word start time", open=False, visible=False):
+                        transcript_with_start_time = gr.Textbox(label="Start time", lines=5, interactive=False, info="Start time before each word")
+                    with gr.Accordion("Word end time", open=False, visible=False):
+                        transcript_with_end_time = gr.Textbox(label="End time", lines=5, interactive=False, info="End time after each word")
+                
+                with gr.Row():
+                    denoise_btn = gr.Button(value="Denoise")
+                    cancel_btn = gr.Button(value="Cancel Denoise")
+                denoise_audio = gr.Audio(label="Denoised Audio", value=None, interactive=False, type="numpy")
+
+            with gr.Column(scale=3):
+                with gr.Group():
+                    transcript_inbox = gr.Textbox(label="Text", lines=5, value=demo_text["Edit"]["smart"])
+                    with gr.Row(visible=False):
+                        smart_transcript = gr.Checkbox(label="Smart transcript", value=True)
+                        with gr.Accordion(label="?", open=False):
+                            info = gr.Markdown(value=smart_transcript_info)
+
+                    mode = gr.Radio(label="Mode", choices=["Edit"], value="Edit", visible=False)
+                    with gr.Row(visible=False):
+                        split_text = gr.Radio(label="Split text", choices=["Newline", "Sentence"], value="Newline",
+                                            info="Split text into parts and run TTS for each part.", visible=True)
+                        edit_word_mode = gr.Radio(label="Edit word mode", choices=["Replace half", "Replace all"], value="Replace all",
+                                                info="What to do with first and last word", visible=False)
+
+                    # with gr.Group(visible=False) as tts_mode_controls:
+                        # with gr.Row():
+                        #     edit_from_word = gr.Dropdown(label="First word in prompt", choices=demo_words, value=demo_words[12], interactive=True)
+                        #     edit_to_word = gr.Dropdown(label="Last word in prompt", choices=demo_words, value=demo_words[18], interactive=True)
+                        # with gr.Row():
+                        #     edit_start_time = gr.Slider(label="Prompt start time", minimum=0, maximum=7.614, step=0.001, value=4.022)
+                        #     edit_end_time = gr.Slider(label="Prompt end time", minimum=0, maximum=7.614, step=0.001, value=5.768)
+                        # with gr.Row():
+                        #     check_btn = gr.Button(value="Check prompt",scale=1)
+                        #     edit_audio = gr.Audio(label="Prompt Audio", scale=3)
+
+                    # with gr.Group() as edit_mode_controls:
+                    with gr.Row():
+                        edit_from_word = gr.Dropdown(label="First word to edit", choices=demo_words, value=demo_words[12], interactive=True)
+                        edit_to_word = gr.Dropdown(label="Last word to edit", choices=demo_words, value=demo_words[18], interactive=True)
+                    with gr.Row():
+                        edit_start_time = gr.Slider(label="Edit from time", minimum=0, maximum=7.614, step=0.001, value=4.022)
+                        edit_end_time = gr.Slider(label="Edit to time", minimum=0, maximum=7.614, step=0.001, value=5.768)
+                    # Put the button and audio in separate columns so that
+                    # the tall audio widget does not overlap the clickable
+                    # area of the button.
+                    with gr.Row():
+                        with gr.Column(scale=1):
+                            check_btn = gr.Button(value="Check edit words")
+                        with gr.Column(scale=3):
+                            edit_audio = gr.Audio(label="Edit word(s)", scale=3, type="numpy")
+
+                    run_btn = gr.Button(value="Run", variant="primary")
+
+            with gr.Column(scale=2):
+                output_audio = gr.Audio(label="Output Audio", type="numpy")
+                with gr.Accordion("Inference transcript", open=True):
+                    inference_transcript = gr.Textbox(label="Inference transcript", lines=5, interactive=False, info="Inference was performed on this transcript.")
+                with gr.Group(visible=False) as long_tts_sentence_editor:
+                    sentence_selector = gr.Dropdown(label="Sentence", value=None,
+                                                    info="Select sentence you want to regenerate")
+                    sentence_audio = gr.Audio(label="Sentence Audio", scale=2, type="numpy")
+                    rerun_btn = gr.Button(value="Rerun")
 
-    with gr.Accordion("Denoise audio (Optional / Recommend)", open=True):
         with gr.Row():
-            denoise_btn = gr.Button(value="Denoise")
-            cancel_btn = gr.Button(value="Cancel Denoise")
-        denoise_audio = gr.Audio(label="Denoised Audio", value=None, type="filepath", interactive=True, show_download_button=True, editable=True)
-
-    gen_text = gr.Textbox(label="Text to Generate", placeholder="Enter the text you want to generate...")
-
-    # Inference button and outputs
-    with gr.Row():
-        txt_info_gpu = gr.Textbox("", label="Device Info")
-        seed_info = gr.Textbox(label="Used Random Seed")
-        check_button_infer = gr.Button("Generate Audio", variant="primary")
-
-    gen_audio = gr.Audio(label="Generated Audio", type="filepath", interactive=True, show_download_button=True, editable=True)
-
-    # Examples
-    def _resolve_example(name: str) -> str:
-        local = Path(PRETRAINED_ROOT) / "data" / "test_examples" / name
-        return str(local) if local.is_file() else ""
-
-    examples = gr.Examples(
-        examples=[
-            ["em, #1 I have a list of YouTubers, and I'm gonna be going to their houses and raiding them by.",
-            _resolve_example("en.wav"),
-            "我有一份 YouTuber 名单，我打算去他们家，对他们进行突袭。",
-            ],
-            ["Te voy a dar un tip #1 que le copia a John Rockefeller, uno de los empresarios más picudos de la historia.",
-            _resolve_example("es.wav"),
-            "我要给你一个从历史上最精明的商人之一约翰·洛克菲勒那里抄来的秘诀。",
-            ],
-            ["Nova, #1 dia 25 desse mês vai rolar operação the last Frontier.",
-            _resolve_example("pt.wav"),
-            "新消息，本月二十五日，'最后的边疆行动'将启动。",
-            ],
-        ],
-        inputs=[
-            ref_text,
-            ref_audio,
-            gen_text,
-        ],
-        outputs=[gen_audio, txt_info_gpu, seed_info],
-        fn=infer,
-        cache_examples=False
-    )
-
-    # System Info section at the bottom
-    gr.Markdown("---")
-    gr.Markdown("## System Information")
-    with gr.Accordion("Update System Stats", open=False):
-        update_button = gr.Button("Update System Stats", scale=1)
-        output_box = gr.Textbox(label="GPU and CPU Information", lines=5, scale=5)
+            with gr.Accordion("Generation Parameters - change these if you are unhappy with the generation", open=False):
+                with gr.Row():
+                    nfe_step = gr.Number(
+                        label="NFE Step",
+                        value=64,
+                        precision=0,
+                        info="Number of function evaluations (sampling steps).",
+                    )
+                    speed = gr.Slider(
+                        label="Speed",
+                        minimum=0.5,
+                        maximum=1.5,
+                        step=0.1,
+                        value=1.0,
+                        info="Placeholder for future use; currently not applied.",
+                    )
+                    cfg_strength = gr.Slider(
+                        label="CFG Strength",
+                        minimum=0.0,
+                        maximum=10.0,
+                        step=0.5,
+                        value=5.0,
+                        info="Classifier-free guidance strength.",
+                    )
+
+                with gr.Row():
+                    sway_sampling_coef = gr.Slider(
+                        label="Sway",
+                        minimum=2.0,
+                        maximum=5.0,
+                        step=0.1,
+                        value=3.0,
+                        info="Sampling sway coefficient.",
+                    )
+                    ref_ratio = gr.Slider(
+                        label="Ref Ratio",
+                        minimum=0.0,
+                        maximum=1.0,
+                        step=0.05,
+                        value=1.0,
+                        info="How much to rely on reference audio (if used).",
+                    )
+                    seed = gr.Number(
+                        label="Seed",
+                        value=-1,
+                        precision=0,
+                        info="-1 for random, otherwise fixed seed.",
+                    )
+
+
+        audio_tensors = gr.State()
+        transcribe_state = gr.State(value={"words_info": demo_words_info, "lang":"zh"})
+
+
+        edit_word_mode.change(fn=update_demo,
+                            inputs=[mode, smart_transcript, edit_word_mode, transcript_inbox, edit_from_word, edit_to_word],
+                            outputs=[transcript_inbox, edit_from_word, edit_to_word])
+        smart_transcript.change(
+            fn=update_demo,
+            inputs=[mode, smart_transcript, edit_word_mode, transcript_inbox, edit_from_word, edit_to_word],
+            outputs=[transcript_inbox, edit_from_word, edit_to_word],
+        )
 
-    def update_stats():
-        return get_combined_stats()
-        
-    
-    denoise_btn.click(fn=denoise,
-                        inputs=[ref_audio],
+        load_models_btn.click(fn=load_models,
+                            inputs=[lemas_model_choice, whisper_model_choice, align_model_choice, denoise_model_choice],  # audiosr_choice],
+                            outputs=[models_selector])
+
+        input_audio.upload(fn=update_input_audio,
+                        inputs=[input_audio],
+                        outputs=[edit_start_time, edit_end_time]) # prompt_end_time
+
+        transcribe_btn.click(fn=transcribe,
+                            inputs=[seed, input_audio],
+                            outputs=[original_transcript, transcript_with_start_time, transcript_with_end_time,
+                                    edit_from_word, edit_to_word, transcribe_state]) # prompt_to_word
+        align_btn.click(fn=align,
+                        inputs=[original_transcript, input_audio, transcribe_state],
+                        outputs=[original_transcript, transcript_with_start_time, transcript_with_end_time,
+                                edit_from_word, edit_to_word, transcribe_state]) # prompt_to_word
+
+        denoise_btn.click(fn=denoise,
+                        inputs=[input_audio],
                         outputs=[denoise_audio])
 
-    cancel_btn.click(fn=cancel_denoise,
-                        inputs=[ref_audio],
+        cancel_btn.click(fn=cancel_denoise,
+                        inputs=[input_audio],
                         outputs=[denoise_audio])
 
-    # Event handlers
-    check_button_infer.click(
-        fn=infer,
-        inputs=[
-            cm_project,
-            cm_checkpoint,
-            exp_name,
-            ref_text,
-            ref_audio,
-            denoise_audio,
-            gen_text,
-            nfe_step,
-            ch_use_ema,
-            separate_langs,
-            frontend,
-            speed,
-            cfg_strength,
-            use_acc_grl,
-            ref_ratio,
-            no_ref_audio,
-            sway_sampling_coef,
-            use_prosody_encoder,
-            seed,
-        ],
-        outputs=[gen_audio, txt_info_gpu, seed_info],
-    )
-
-    bt_checkpoint_refresh.click(fn=get_checkpoints_project, inputs=[cm_project], outputs=[cm_checkpoint])
-    cm_project.change(fn=get_checkpoints_project, inputs=[cm_project], outputs=[cm_checkpoint])
-
-    ref_audio.change(
-            fn=lambda x: None,
-            inputs=[ref_audio],
-            outputs=[denoise_audio]
+        # mode.change(fn=change_mode,
+        #             inputs=[mode],
+        #             outputs=[tts_mode_controls, edit_mode_controls, edit_word_mode, split_text, long_tts_sentence_editor])
+
+        check_btn.click(fn=get_edit_audio_part,
+                        inputs=[input_audio, edit_start_time, edit_end_time],
+                        outputs=[edit_audio])
+
+        run_btn.click(fn=run,
+                    inputs=[
+                        seed, nfe_step, speed, cfg_strength, sway_sampling_coef, ref_ratio,
+                        input_audio, denoise_audio, transcribe_state, transcript_inbox, smart_transcript,
+                        mode, edit_start_time, edit_end_time,
+                        split_text, sentence_selector, audio_tensors
+                    ],
+                    outputs=[output_audio, inference_transcript, sentence_selector, audio_tensors])
+
+        sentence_selector.change(
+            fn=load_sentence,
+            inputs=[sentence_selector, audio_tensors],
+            outputs=[sentence_audio],
         )
-
-    update_button.click(fn=update_stats, outputs=output_box)
-    
-    # Auto-load system stats on startup
-    app.load(fn=update_stats, outputs=output_box)
-
-
-@click.command()
-@click.option("--port", "-p", default=7860, type=int, help="Port to run the app on")
-@click.option("--host", "-H", default="0.0.0.0", help="Host to run the app on")
-@click.option(
-    "--share",
-    "-s",
-    default=False,
-    is_flag=True,
-    help="Share the app via Gradio share link",
-)
-@click.option("--api", "-a", default=True, is_flag=True, help="Allow API access")
-def main(port, host, share, api):
-    global app
-    print("Starting LEMAS-TTS Inference Interface...")
-    print(f"Device: {device}")
-    app.queue(api_open=api).launch(
-        server_name=host,
-        server_port=port,
-        share=share,
-        show_api=api,
-        allowed_paths=[str(Path(PRETRAINED_ROOT) / "data")],
-    )
+        rerun_btn.click(fn=run,
+                        inputs=[
+                            seed, nfe_step, speed, cfg_strength, sway_sampling_coef, ref_ratio,
+                            input_audio, denoise_audio, transcribe_state, transcript_inbox, smart_transcript,
+                            gr.State(value="Rerun"), edit_start_time, edit_end_time,
+                            split_text, sentence_selector, audio_tensors
+                        ],
+                        outputs=[output_audio, inference_transcript, sentence_audio, audio_tensors])
+
+        # prompt_to_word.change(fn=update_bound_word,
+        #                     inputs=[gr.State(False), prompt_to_word, gr.State("Replace all")],
+        #                     outputs=[prompt_end_time])
+        edit_from_word.change(fn=update_bound_word,
+                            inputs=[gr.State(True), edit_from_word, edit_word_mode],
+                            outputs=[edit_start_time])
+        edit_to_word.change(fn=update_bound_word,
+                            inputs=[gr.State(False), edit_to_word, edit_word_mode],
+                            outputs=[edit_end_time])
+        edit_word_mode.change(fn=update_bound_words,
+                            inputs=[edit_from_word, edit_to_word, edit_word_mode],
+                            outputs=[edit_start_time, edit_end_time])
+
+    return app
 
 
 if __name__ == "__main__":
-    main()
+    import argparse
+
+    parser = argparse.ArgumentParser(description="VoiceCraft gradio app.")
+    
+    parser.add_argument("--demo-path", default="./demo", help="Path to demo directory")
+    parser.add_argument("--tmp-path", default="/cto_labs/vistring/zhaozhiyuan/outputs/voicecraft/tmp", help="Path to tmp directory")
+    parser.add_argument("--models-path", default="/cto_labs/vistring/zhaozhiyuan/outputs/voicecraft/pretrain/VoiceCraft", help="Path to voicecraft models directory")
+    parser.add_argument("--port", default=41020, type=int, help="App port")
+    parser.add_argument("--share", action="store_true", help="Launch with public url")
+    parser.add_argument("--server_name", default="0.0.0.0", type=str, help="Server name for launching the app. 127.0.0.1 for localhost; 0.0.0.0 to allow access from other machines in the local network. Might also give access to external users depends on the firewall settings.")
+
+    os.environ["USER"] = os.getenv("USER", "user")
+    args = parser.parse_args()
+    DEMO_PATH = args.demo_path
+    TMP_PATH = args.tmp_path
+    MODELS_PATH = args.models_path
+    
+    app = get_app()
+    app.queue().launch(share=args.share, server_name=args.server_name, server_port=args.port)