import gradio as gr import torch import safetensors from huggingface_hub import hf_hub_download import soundfile as sf import os import numpy as np import librosa from models.codec.kmeans.repcodec_model import RepCodec from models.tts.maskgct.maskgct_s2a import MaskGCT_S2A from models.tts.maskgct.maskgct_t2s import MaskGCT_T2S from models.codec.amphion_codec.codec import CodecEncoder, CodecDecoder from transformers import Wav2Vec2BertModel from utils.util import load_config from models.tts.maskgct.g2p.g2p_generation import g2p, chn_eng_g2p from transformers import SeamlessM4TFeatureExtractor processor = SeamlessM4TFeatureExtractor.from_pretrained("facebook/w2v-bert-2.0") device = torch.device("cuda" if torch.cuda.is_available() else "CPU") def g2p_(text, language): if language in ["zh", "en"]: return chn_eng_g2p(text) else: return g2p(text, sentence=None, language=language) def build_t2s_model(cfg, device): t2s_model = MaskGCT_T2S(cfg=cfg) t2s_model.eval() t2s_model.to(device) return t2s_model def build_s2a_model(cfg, device): soundstorm_model = MaskGCT_S2A(cfg=cfg) soundstorm_model.eval() soundstorm_model.to(device) return soundstorm_model def build_semantic_model(device): semantic_model = Wav2Vec2BertModel.from_pretrained("facebook/w2v-bert-2.0") semantic_model.eval() semantic_model.to(device) stat_mean_var = torch.load("./models/tts/maskgct/ckpt/wav2vec2bert_stats.pt") semantic_mean = stat_mean_var["mean"] semantic_std = torch.sqrt(stat_mean_var["var"]) semantic_mean = semantic_mean.to(device) semantic_std = semantic_std.to(device) return semantic_model, semantic_mean, semantic_std def build_semantic_codec(cfg, device): semantic_codec = RepCodec(cfg=cfg) semantic_codec.eval() semantic_codec.to(device) return semantic_codec def build_acoustic_codec(cfg, device): codec_encoder = CodecEncoder(cfg=cfg.encoder) codec_decoder = CodecDecoder(cfg=cfg.decoder) codec_encoder.eval() codec_decoder.eval() codec_encoder.to(device) codec_decoder.to(device) return codec_encoder, codec_decoder @torch.no_grad() def extract_features(speech, processor): inputs = processor(speech, sampling_rate=16000, return_tensors="pt") input_features = inputs["input_features"][0] attention_mask = inputs["attention_mask"][0] return input_features, attention_mask @torch.no_grad() def extract_semantic_code(semantic_mean, semantic_std, input_features, attention_mask): vq_emb = semantic_model( input_features=input_features, attention_mask=attention_mask, output_hidden_states=True, ) feat = vq_emb.hidden_states[17] # (B, T, C) feat = (feat - semantic_mean.to(feat)) / semantic_std.to(feat) semantic_code, rec_feat = semantic_codec.quantize(feat) # (B, T) return semantic_code, rec_feat @torch.no_grad() def extract_acoustic_code(speech): vq_emb = codec_encoder(speech.unsqueeze(1)) _, vq, _, _, _ = codec_decoder.quantizer(vq_emb) acoustic_code = vq.permute(1, 2, 0) return acoustic_code @torch.no_grad() def text2semantic( device, prompt_speech, prompt_text, prompt_language, target_text, target_language, target_len=None, n_timesteps=50, cfg=2.5, rescale_cfg=0.75, ): prompt_phone_id = g2p_(prompt_text, prompt_language)[1] target_phone_id = g2p_(target_text, target_language)[1] if target_len is None: target_len = int( (len(prompt_speech) * len(target_phone_id) / len(prompt_phone_id)) / 16000 * 50 ) else: target_len = int(target_len * 50) prompt_phone_id = torch.tensor(prompt_phone_id, dtype=torch.long).to(device) target_phone_id = torch.tensor(target_phone_id, dtype=torch.long).to(device) phone_id = torch.cat([prompt_phone_id, target_phone_id]) input_fetures, attention_mask = extract_features(prompt_speech, processor) input_fetures = input_fetures.unsqueeze(0).to(device) attention_mask = attention_mask.unsqueeze(0).to(device) semantic_code, _ = extract_semantic_code( semantic_mean, semantic_std, input_fetures, attention_mask ) predict_semantic = t2s_model.reverse_diffusion( semantic_code[:, :], target_len, phone_id.unsqueeze(0), n_timesteps=n_timesteps, cfg=cfg, rescale_cfg=rescale_cfg, ) combine_semantic_code = torch.cat([semantic_code[:, :], predict_semantic], dim=-1) prompt_semantic_code = semantic_code return combine_semantic_code, prompt_semantic_code @torch.no_grad() def semantic2acoustic( device, combine_semantic_code, acoustic_code, n_timesteps=[25, 10, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], cfg=2.5, rescale_cfg=0.75, ): semantic_code = combine_semantic_code cond = s2a_model_1layer.cond_emb(semantic_code) prompt = acoustic_code[:, :, :] predict_1layer = s2a_model_1layer.reverse_diffusion( cond=cond, prompt=prompt, temp=1.5, filter_thres=0.98, n_timesteps=n_timesteps[:1], cfg=cfg, rescale_cfg=rescale_cfg, ) cond = s2a_model_full.cond_emb(semantic_code) prompt = acoustic_code[:, :, :] predict_full = s2a_model_full.reverse_diffusion( cond=cond, prompt=prompt, temp=1.5, filter_thres=0.98, n_timesteps=n_timesteps, cfg=cfg, rescale_cfg=rescale_cfg, gt_code=predict_1layer, ) vq_emb = codec_decoder.vq2emb(predict_full.permute(2, 0, 1), n_quantizers=12) recovered_audio = codec_decoder(vq_emb) prompt_vq_emb = codec_decoder.vq2emb(prompt.permute(2, 0, 1), n_quantizers=12) recovered_prompt_audio = codec_decoder(prompt_vq_emb) recovered_prompt_audio = recovered_prompt_audio[0][0].cpu().numpy() recovered_audio = recovered_audio[0][0].cpu().numpy() combine_audio = np.concatenate([recovered_prompt_audio, recovered_audio]) return combine_audio, recovered_audio # Load the model and checkpoints def load_models(): cfg_path = "./models/tts/maskgct/config/maskgct.json" cfg = load_config(cfg_path) semantic_model, semantic_mean, semantic_std = build_semantic_model(device) semantic_codec = build_semantic_codec(cfg.model.semantic_codec, device) codec_encoder, codec_decoder = build_acoustic_codec( cfg.model.acoustic_codec, device ) t2s_model = build_t2s_model(cfg.model.t2s_model, device) s2a_model_1layer = build_s2a_model(cfg.model.s2a_model.s2a_1layer, device) s2a_model_full = build_s2a_model(cfg.model.s2a_model.s2a_full, device) # Download checkpoints semantic_code_ckpt = hf_hub_download( "amphion/MaskGCT", filename="semantic_codec/model.safetensors" ) codec_encoder_ckpt = hf_hub_download( "amphion/MaskGCT", filename="acoustic_codec/model.safetensors" ) codec_decoder_ckpt = hf_hub_download( "amphion/MaskGCT", filename="acoustic_codec/model_1.safetensors" ) t2s_model_ckpt = hf_hub_download( "amphion/MaskGCT", filename="t2s_model/model.safetensors" ) s2a_1layer_ckpt = hf_hub_download( "amphion/MaskGCT", filename="s2a_model/s2a_model_1layer/model.safetensors" ) s2a_full_ckpt = hf_hub_download( "amphion/MaskGCT", filename="s2a_model/s2a_model_full/model.safetensors" ) safetensors.torch.load_model(semantic_codec, semantic_code_ckpt) safetensors.torch.load_model(codec_encoder, codec_encoder_ckpt) safetensors.torch.load_model(codec_decoder, codec_decoder_ckpt) safetensors.torch.load_model(t2s_model, t2s_model_ckpt) safetensors.torch.load_model(s2a_model_1layer, s2a_1layer_ckpt) safetensors.torch.load_model(s2a_model_full, s2a_full_ckpt) return ( semantic_model, semantic_mean, semantic_std, semantic_codec, codec_encoder, codec_decoder, t2s_model, s2a_model_1layer, s2a_model_full, ) @torch.no_grad() def maskgct_inference( prompt_speech_path, prompt_text, target_text, language="en", target_language="en", target_len=None, n_timesteps=25, cfg=2.5, rescale_cfg=0.75, n_timesteps_s2a=[25, 10, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], cfg_s2a=2.5, rescale_cfg_s2a=0.75, device=torch.device("cuda:5"), ): speech_16k = librosa.load(prompt_speech_path, sr=16000)[0] speech = librosa.load(prompt_speech_path, sr=24000)[0] combine_semantic_code, _ = text2semantic( device, speech_16k, prompt_text, language, target_text, target_language, target_len, n_timesteps, cfg, rescale_cfg, ) acoustic_code = extract_acoustic_code(torch.tensor(speech).unsqueeze(0).to(device)) _, recovered_audio = semantic2acoustic( device, combine_semantic_code, acoustic_code, n_timesteps=n_timesteps_s2a, cfg=cfg_s2a, rescale_cfg=rescale_cfg_s2a, ) return recovered_audio @torch.no_grad() def inference( prompt_wav, prompt_text, target_text, target_len, n_timesteps, language, target_language, ): save_path = "./output/output.wav" os.makedirs("./output", exist_ok=True) recovered_audio = maskgct_inference( prompt_wav, prompt_text, target_text, language, target_language, target_len=target_len, n_timesteps=int(n_timesteps), device=device, ) sf.write(save_path, recovered_audio, 24000) return save_path # Language list language_list = ["en", "zh", "ja", "ko", "fr", "de"] # Gradio interface iface = gr.Interface( fn=inference, inputs=[ gr.Audio(label="Upload Prompt Wav", type="filepath"), gr.Textbox(label="Prompt Text"), gr.Textbox(label="Target Text"), gr.Number( label="Target Duration (in seconds)", value=None ), # Removed 'optional=True' gr.Slider( label="Number of Timesteps", minimum=15, maximum=100, value=25, step=1 ), gr.Dropdown(label="Language", choices=language_list, value="en"), gr.Dropdown(label="Target Language", choices=language_list, value="en"), ], outputs=gr.Audio(label="Generated Audio"), title="MaskGCT TTS Demo", description="Generate speech from text using the MaskGCT model.", ) # Launch the interface iface.launch(allowed_paths=["./output"])