import gc import os import platform import psutil import tempfile from glob import glob import traceback import click import gradio as gr import torch import sys from pathlib import Path # Add the local code directory so that `lemas_tts` can be imported when running this # script directly without installing the package. THIS_FILE = Path(__file__).resolve() SRC_ROOT = THIS_FILE.parents[2] # .../code sys.path.append(str(SRC_ROOT)) def _find_repo_root(start: Path) -> Path: """Locate the repo root by looking for a `pretrained_models` folder upwards.""" for p in [start, *start.parents]: if (p / "pretrained_models").is_dir(): return p cwd = Path.cwd() if (cwd / "pretrained_models").is_dir(): return cwd return start REPO_ROOT = _find_repo_root(THIS_FILE) PRETRAINED_ROOT = REPO_ROOT / "pretrained_models" CKPTS_ROOT = PRETRAINED_ROOT / "ckpts" DATA_ROOT = PRETRAINED_ROOT / "data" UVR5_CODE_DIR = REPO_ROOT / "code" / "uvr5" UVR5_MODEL_DIR = PRETRAINED_ROOT / "uvr5" / "models" / "MDX_Net_Models" / "model_data" from lemas_tts.api import F5TTS import torch, torchaudio import soundfile as sf # 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" ) class UVR5: def __init__(self, model_dir): code_dir = str(UVR5_CODE_DIR) self.model = self.load_model(str(model_dir), code_dir) def load_model(self, model_dir, code_dir): import sys, json, os 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') configs = json.loads(open(config_path, 'r', encoding='utf-8').read()) model_data = ModelData( model_path=model_path, audio_path = model_dir, result_path = model_dir, device = 'cpu', process_method = "MDX-Net", base_dir=code_dir, **configs ) 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) # transform = torchaudio.transforms.Resample(44100, 16000) # output_audio = transform(output_audio) return output_audio.squeeze().T.numpy(), 44100 denoise_model = UVR5(UVR5_MODEL_DIR) def load_wav(audio_info, sr=16000, channel=1): print("load audio:", audio_info) audio, raw_sr = torchaudio.load(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() 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: 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 cancel_denoise(audio_info): return audio_info def get_checkpoints_project(project_name=None, is_gradio=True): """Get available checkpoint files""" checkpoint_dir = [str(CKPTS_ROOT)] if project_name is None: # Look for checkpoints in common locations 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 else: # project_name = project_name.replace("_pinyin", "").replace("_char", "") project_name = "_".join(["F5TTS_v1_Base", "vocos", "custom", project_name.replace("_custom", "")]) if project_name != "F5TTS_v1_Base" else project_name 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"))) else: files_checkpoints = [] 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) ] 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 select_checkpoint = None if not files_checkpoints else files_checkpoints[-1] if is_gradio: return gr.update(choices=files_checkpoints, value=select_checkpoint) return files_checkpoints, select_checkpoint def get_available_projects(): """Get available project names from data directory""" data_path = str(DATA_ROOT) project_list = [] if os.path.isdir(data_path): for folder in os.listdir(data_path): if "test" in folder: continue project_list.append(folder) # Fallback to a sensible default if no projects are found if not project_list: project_list = ["multilingual_acc_grl_custom"] 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 if not os.path.isfile(file_checkpoint): return None, "Checkpoint not found!", "" if denoise_audio: ref_audio = denoise_audio device_test = device # Use the global device if last_checkpoint != file_checkpoint or last_device != device_test or last_ema != use_ema or tts_api is None: if last_checkpoint != file_checkpoint: last_checkpoint = file_checkpoint if last_device != device_test: last_device = device_test if last_ema != use_ema: last_ema = use_ema # Try to find vocab file vocab_file = None possible_vocab_paths = [ str(DATA_ROOT / project / "vocab.txt"), # legacy fallbacks for older layouts f"./data/{project}/vocab.txt", f"../../data/{project}/vocab.txt", "./data/Emilia_ZH_EN_pinyin/vocab.txt", "../../data/Emilia_ZH_EN_pinyin/vocab.txt", ] for path in possible_vocab_paths: if os.path.isfile(path): vocab_file = path break if vocab_file is None: return None, "Vocab file not found!", "" try: tts_api = F5TTS( model=exp_name, ckpt_file=file_checkpoint, vocab_file=vocab_file, device=device_test, use_ema=use_ema, frontend=frontend, use_prosody_encoder=use_prosody_encoder, prosody_cfg_path=str(CKPTS_ROOT / "prosody_encoder" / "pretssel_cfg.json"), prosody_ckpt_path=str(CKPTS_ROOT / "prosody_encoder" / "prosody_encoder_UnitY2.pt"), ) 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" ) else: gpu_stats = "No GPU available" return gpu_stats 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 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}" ) return cpu_stats 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 # Create Gradio interface with gr.Blocks(title="LEMAS-TTS Inference") as app: gr.Markdown( """ # Zero-Shot TTS 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=["F5TTS_v1_Base", "F5TTS_Base", "E2TTS_Base"], value="F5TTS_v1_Base", 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) 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) with gr.Row(): ch_use_ema = gr.Checkbox(label="Use EMA", value=True, scale=2, info="Turn off at early stage might offer better results") frontend = gr.Radio(label="Frontend", choices=["phone", "char", "bpe"], value="phone", scale=3) separate_langs = gr.Checkbox(label="Separate Languages", value=True, scale=2, info="separate language tokens") # 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=-1, 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", value=False, scale=1, info="No mel condition") use_acc_grl = gr.Checkbox(label="Use accent grl condition", value=False, scale=1, info="Use accent grl condition") use_prosody_encoder = gr.Checkbox(label="Use prosody encoder", value=False, scale=1, info="Use prosody encoder") seed = gr.Number(label="Random Seed", scale=1, value=5828684826493313192, 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.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 examples = gr.Examples( examples=[ [ "Ich glaub, mein Schwein pfeift.", str(DATA_ROOT / "test_examples" / "de.wav"), "我觉得我的猪在吹口哨。", ], [ "em, #1 I have a list of YouTubers, and I'm gonna be going to their houses and raiding them by.", str(DATA_ROOT / "test_examples" / "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.", str(DATA_ROOT / "test_examples" / "es.wav"), "我要给你一个从历史上最精明的商人之一约翰·洛克菲勒那里抄来的秘诀。", ], [ "Per l'amor di Dio #1 fai, #2 se pensi di non poterti fermare, fallo #1 e fallo.", str(DATA_ROOT / "test_examples" / "it.wav"), "看在上帝的份上，去做吧，如果你认为你无法停止，那就去做吧，继续做下去。", ], [ "Nova, #1 dia 25 desse mês vai rolar operação the last Frontier.", str(DATA_ROOT / "test_examples" / "pt.wav"), "新消息，本月二十五日，'最后的边疆行动'将启动。", ], # ["Good morning! #1 ", # "/mnt/code/lemas/F5-TTS/data/trueman/recognition_d0a02641c090813574a8ec398220339f_0.wav", # " #1" # ], # ["Good morning! #1 ", # "/mnt/code/lemas/F5-TTS/data/trueman/recognition_d0a02641c090813574a8ec398220339f_1.wav", # " #1", # ], # ["Good morning! #1 ", # "/mnt/code/lemas/F5-TTS/data/trueman/recognition_d0a02641c090813574a8ec398220339f_2.wav", # " #1", # ], # ["Oh, and in case I don't see ya, #1", # "/mnt/code/lemas/F5-TTS/data/trueman/recognition_d0a02641c090813574a8ec398220339f_3.wav", # " #1", # ], # ["Good afternoon, good evening, and good night. #1", # "/mnt/code/lemas/F5-TTS/data/trueman/recognition_d0a02641c090813574a8ec398220339f_4.wav", # " #1", # ], ], 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) def update_stats(): return get_combined_stats() denoise_btn.click(fn=denoise, inputs=[ref_audio], outputs=[denoise_audio]) cancel_btn.click(fn=cancel_denoise, inputs=[ref_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] ) 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(DATA_ROOT)], ) if __name__ == "__main__": main()