Spaces:
Runtime error
Runtime error
| from dataclasses import dataclass | |
| from pathlib import Path | |
| import librosa | |
| import torch | |
| import torch.nn.functional as F | |
| from huggingface_hub import hf_hub_download | |
| from .models.t3 import T3 | |
| from .models.s3tokenizer import S3_SR, drop_invalid_tokens | |
| from .models.s3gen import S3GEN_SR, S3Gen | |
| from .models.tokenizers import EnTokenizer | |
| from .models.voice_encoder import VoiceEncoder | |
| from .models.t3.modules.cond_enc import T3Cond | |
| REPO_ID = "ResembleAI/Orator" | |
| class Conditionals: | |
| """ | |
| Conditionals for T3 and S3Gen | |
| - T3 conditionals: | |
| - speaker_emb | |
| - clap_emb | |
| - cond_prompt_speech_tokens | |
| - cond_prompt_speech_emb | |
| - emotion_adv | |
| - S3Gen conditionals: | |
| - prompt_token | |
| - prompt_token_len | |
| - prompt_feat | |
| - prompt_feat_len | |
| - embedding | |
| """ | |
| t3: T3Cond | |
| gen: dict | |
| def to(self, device): | |
| self.t3 = self.t3.to(device=device) | |
| for k, v in self.gen.items(): | |
| if torch.is_tensor(v): | |
| self.gen[k] = v.to(device=device) | |
| return self | |
| def save(self, fpath: Path): | |
| arg_dict = dict( | |
| t3=self.t3.__dict__, | |
| gen=self.gen | |
| ) | |
| torch.save(arg_dict, fpath) | |
| def load(cls, fpath, map_location="cpu"): | |
| kwargs = torch.load(fpath, map_location=map_location, weights_only=True) | |
| return cls(T3Cond(**kwargs['t3']), kwargs['gen']) | |
| class OratorTTS: | |
| ENC_COND_LEN = 6 * S3_SR | |
| DEC_COND_LEN = 10 * S3GEN_SR | |
| def __init__( | |
| self, | |
| t3: T3, | |
| s3gen: S3Gen, | |
| ve: VoiceEncoder, | |
| tokenizer: EnTokenizer, | |
| device: str, | |
| conds: Conditionals = None, | |
| ): | |
| self.sr = S3GEN_SR # sample rate of synthesized audio | |
| self.t3 = t3 | |
| self.s3gen = s3gen | |
| self.ve = ve | |
| self.tokenizer = tokenizer | |
| self.device = device | |
| self.conds = conds | |
| def from_local(cls, ckpt_dir, device) -> 'OratorTTS': | |
| ckpt_dir = Path(ckpt_dir) | |
| ve = VoiceEncoder() | |
| ve.load_state_dict( | |
| torch.load(ckpt_dir / "ve.pt") | |
| ) | |
| ve.to(device).eval() | |
| t3 = T3() | |
| t3.load_state_dict( | |
| torch.load(ckpt_dir / "t3.pt") | |
| ) | |
| t3.to(device).eval() | |
| s3gen = S3Gen() | |
| s3gen.load_state_dict( | |
| torch.load(ckpt_dir / "s3gen.pt") | |
| ) | |
| s3gen.to(device).eval() | |
| tokenizer = EnTokenizer( | |
| str(ckpt_dir / "tokenizer.json") | |
| ) | |
| conds = None | |
| if (builtin_voice := ckpt_dir / "conds.pt").exists(): | |
| conds = Conditionals.load(builtin_voice).to(device) | |
| return cls(t3, s3gen, ve, tokenizer, device, conds=conds) | |
| def from_pretrained(cls, device) -> 'OratorTTS': | |
| for fpath in ["ve.pt", "t3.pt", "s3gen.pt", "tokenizer.json", "conds.pt"]: | |
| local_path = hf_hub_download(repo_id=REPO_ID, filename=fpath) | |
| return cls.from_local(Path(local_path).parent, device) | |
| def prepare_conditionals(self, wav_fpath, emotion_adv=0.5): | |
| ## Load reference wav | |
| s3gen_ref_wav, _sr = librosa.load(wav_fpath, sr=S3GEN_SR) | |
| s3_ref_wav = librosa.resample(s3gen_ref_wav, orig_sr=S3GEN_SR, target_sr=S3_SR) | |
| s3gen_ref_wav = s3gen_ref_wav[:self.DEC_COND_LEN] | |
| s3gen_ref_dict = self.s3gen.embed_ref(s3gen_ref_wav, S3GEN_SR, device=self.device) | |
| # Speech cond prompt tokens | |
| if plen := self.t3.hp.speech_cond_prompt_len: | |
| s3_tokzr = self.s3gen.tokenizer | |
| t3_cond_prompt_tokens, _ = s3_tokzr.forward([s3_ref_wav[:self.ENC_COND_LEN]], max_len=plen) | |
| t3_cond_prompt_tokens = torch.atleast_2d(t3_cond_prompt_tokens).to(self.device) | |
| # # Voice-encoder speaker embedding | |
| ve_embed = torch.from_numpy(self.ve.embeds_from_wavs([s3_ref_wav], sample_rate=S3_SR)) | |
| ve_embed = ve_embed.mean(axis=0, keepdim=True).to(self.device) | |
| t3_cond = T3Cond( | |
| speaker_emb=ve_embed, | |
| cond_prompt_speech_tokens=t3_cond_prompt_tokens, | |
| emotion_adv=emotion_adv * torch.ones(1, 1, 1), | |
| ).to(device=self.device) | |
| self.conds = Conditionals(t3_cond, s3gen_ref_dict) | |
| def generate( | |
| self, | |
| text, | |
| audio_prompt_path=None, | |
| emotion_adv=0.5 | |
| ): | |
| if audio_prompt_path: | |
| self.prepare_conditionals(audio_prompt_path, emotion_adv=emotion_adv) | |
| else: | |
| assert self.conds is not None, "Please `prepare_conditionals` first or specify `audio_prompt_path`" | |
| # Update emotion_adv if needed | |
| if emotion_adv != self.conds.t3.emotion_adv[0, 0, 0]: | |
| _cond: T3Cond = self.conds.t3 | |
| self.conds.t3 = T3Cond( | |
| speaker_emb=_cond.speaker_emb, | |
| cond_prompt_speech_tokens=_cond.cond_prompt_speech_tokens, | |
| emotion_adv=emotion_adv * torch.ones(1, 1, 1), | |
| ).to(device=self.device) | |
| text_tokens = self.tokenizer.text_to_tokens(text).to(self.device) | |
| sot = self.t3.hp.start_text_token | |
| eot = self.t3.hp.stop_text_token | |
| text_tokens = F.pad(text_tokens, (1, 0), value=sot) | |
| text_tokens = F.pad(text_tokens, (0, 1), value=eot) | |
| with torch.inference_mode(): | |
| speech_tokens = self.t3.inference( | |
| t3_cond=self.conds.t3, | |
| text_tokens=text_tokens, | |
| max_new_tokens=1000, # TODO: use the value in config | |
| ) | |
| # TODO: output becomes 1D | |
| speech_tokens = drop_invalid_tokens(speech_tokens) | |
| speech_tokens = speech_tokens.to(self.device) | |
| wav, _ = self.s3gen.inference( | |
| speech_tokens=speech_tokens, | |
| ref_dict=self.conds.gen, | |
| ) | |
| return wav.detach().cpu() | |