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 torchaudio import soundfile as sf from pathlib import Path import spaces from cached_path import cached_path from lemas_tts.api import TTS, PRETRAINED_ROOT, CKPTS_ROOT device = ( "cuda" if torch.cuda.is_available() else "xpu" if torch.xpu.is_available() else "mps" if torch.backends.mps.is_available() else "cpu" ) REPO_ROOT = Path(__file__).resolve().parent # 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) class UVR5: """Small wrapper around the bundled uvr5 implementation for denoising.""" def __init__(self, model_dir): # Code directory is always the local `uvr5` folder in this repo self.code_dir = os.path.join(os.path.dirname(__file__), "uvr5") self.model_dir = model_dir self.model = None self.device = "cpu" def load_model(self, device="cpu"): import sys, json, os, torch if self.code_dir not in sys.path: sys.path.append(self.code_dir) # Reuse an already-loaded model if it matches the requested device. if self.model is not None and self.device == device: return self.model from multiprocess_cuda_infer import ModelData, Inference # In the minimal LEMAS-TTS layout, UVR5 weights live under: model_path = os.path.join(self.model_dir, "Kim_Vocal_1.onnx") config_path = os.path.join(self.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=self.model_dir, result_path=self.model_dir, device=device, process_method="MDX-Net", # Keep base_dir and model_dir the same so all UVR5 metadata # (model_data.json, model_name_mapper.json, etc.) are resolved # under `pretrained_models/uvr5`, matching LEMAS-TTS inference. base_dir=self.model_dir, **configs, ) uvr5_model = Inference(model_data, device) # On HF Spaces with stateless GPU, we must not initialize CUDA in the # main process. The heavy UVR5 loading happens lazily inside # @spaces.GPU functions; this guard is kept only for the CPU path to # avoid any accidental CUDA init. uvr5_model.load_model(model_path, 1) self.model = uvr5_model self.device = device return self.model def denoise(self, audio_info): # Prefer GPU if available; on Spaces this runs inside @spaces.GPU so # CUDA can be safely initialized here. model = self.load_model(device="cpu") input_audio = load_wav(audio_info, sr=44100, channel=2) output_audio = model.demix_base({0:input_audio.squeeze()}, is_match_mix=False, device="cpu") # transform = torchaudio.transforms.Resample(44100, 16000) # output_audio = transform(output_audio) return output_audio.squeeze().T.cpu().numpy(), 44100 denoise_model = UVR5( model_dir=Path(PRETRAINED_ROOT) / "uvr5", ) 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): # Return a numpy waveform tuple for direct playback in Gradio. denoised_audio, sr = denoise_model.denoise(audio_info) return (sr, denoised_audio) 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)] # 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", } 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()) 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) ] # 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 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_paths = [ str(Path(PRETRAINED_ROOT) / "data"), ] 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) 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(reverse=False) print("project_list:", project_list) return project_list @spaces.GPU @torch.no_grad() @torch.inference_mode() 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 tts_api # 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!", "" # Prepare reference audio: # - `ref_audio` from Gradio is a filepath (original reference) # - `denoise_audio` is an optional (sr, numpy_array) tuple from UVR5. # If provided, dump it to a temporary wav file and use that as ref_file. ref_audio_path = ref_audio tmp_ref_path = None if denoise_audio is not None: try: sr_d, wav_d = denoise_audio with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as f_ref: sf.write(f_ref.name, wav_d, int(sr_d), format="wav", subtype="PCM_24") tmp_ref_path = f_ref.name ref_audio_path = f_ref.name except Exception as e: traceback.print_exc() return None, f"Error preparing denoised reference audio: {str(e)}", "" # 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) # 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="cuda", 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() # Cleanup temp ref file if it was created if tmp_ref_path is not None and os.path.isfile(tmp_ref_path): os.remove(tmp_ref_path) return None, f"Error loading model: {str(e)}", "" print("Model loaded >>", 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_path, 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)}", "" finally: # Remove temporary reference file if created if tmp_ref_path is not None and os.path.isfile(tmp_ref_path): os.remove(tmp_ref_path) 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=["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) 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", 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") # 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.Accordion("Denoise audio (Optional / Recommend)", open=True): with gr.Row(): denoise_btn = gr.Button(value="Denoise") cancel_btn = gr.Button(value="Cancel Denoise") # Use numpy type here so we can reuse the waveform directly in Python. denoise_audio = gr.Audio( label="Denoised Audio", value=None, type="numpy", 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=[ ["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"), "我要给你一个从历史上最精明的商人之一，John Rockefeller那里抄来的秘诀。", ], ["Nova, #1 dia 25 desse mês vai rolar operação the last Frontier.", _resolve_example("pt.wav"), " Preparations are currently underway to ensure the operation proceeds as planned.", ], ], 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(Path(PRETRAINED_ROOT) / "data")], ) if __name__ == "__main__": main()