From 42cb1f36741aa3a24e7aab03e73b51becd182fa7 Mon Sep 17 00:00:00 2001 From: mrq Date: Wed, 15 Mar 2023 00:30:28 +0000 Subject: [PATCH] added args for tokenizer and diffusion model (so I don't have to add it later) --- setup.py | 2 +- tortoise/api.py | 62 +++++++++++++++++++++++++++++++++++++++++-------- 2 files changed, 53 insertions(+), 11 deletions(-) diff --git a/setup.py b/setup.py index 3a578f5..86a3938 100644 --- a/setup.py +++ b/setup.py @@ -6,7 +6,7 @@ with open("README.md", "r", encoding="utf-8") as fh: setuptools.setup( name="TorToiSe", packages=setuptools.find_packages(), - version="2.4.4", + version="2.4.5", author="James Betker", author_email="james@adamant.ai", description="A high quality multi-voice text-to-speech library", diff --git a/tortoise/api.py b/tortoise/api.py index f1a60f2..965e133 100755 --- a/tortoise/api.py +++ b/tortoise/api.py @@ -265,7 +265,11 @@ class TextToSpeech: Main entry point into Tortoise. """ - def __init__(self, autoregressive_batch_size=None, models_dir=MODELS_DIR, enable_redaction=True, device=None, minor_optimizations=True, input_sample_rate=22050, output_sample_rate=24000, autoregressive_model_path=None, vocoder_model=None): + def __init__(self, autoregressive_batch_size=None, models_dir=MODELS_DIR, enable_redaction=True, device=None, + minor_optimizations=True, + input_sample_rate=22050, output_sample_rate=24000, + autoregressive_model_path=None, diffusion_model_path=None, vocoder_model=None, tokenizer_json=None + ): """ Constructor :param autoregressive_batch_size: Specifies how many samples to generate per batch. Lower this if you are seeing @@ -300,23 +304,23 @@ class TextToSpeech: if self.enable_redaction: self.aligner = Wav2VecAlignment(device='cpu' if get_device_name() == "dml" else self.device) - self.tokenizer = VoiceBpeTokenizer() - + self.load_tokenizer_json(tokenizer_json) if os.path.exists(f'{models_dir}/autoregressive.ptt'): - # Assume this is a traced directory. self.autoregressive = torch.jit.load(f'{models_dir}/autoregressive.ptt') - self.diffusion = torch.jit.load(f'{models_dir}/diffusion_decoder.ptt') else: if not autoregressive_model_path or not os.path.exists(autoregressive_model_path): autoregressive_model_path = get_model_path('autoregressive.pth', models_dir) self.load_autoregressive_model(autoregressive_model_path) - self.diffusion = DiffusionTts(model_channels=1024, num_layers=10, in_channels=100, out_channels=200, - in_latent_channels=1024, in_tokens=8193, dropout=0, use_fp16=False, num_heads=16, - layer_drop=0, unconditioned_percentage=0).cpu().eval() - self.diffusion.load_state_dict(torch.load(get_model_path('diffusion_decoder.pth', models_dir))) + if os.path.exists(f'{models_dir}/diffusion_decoder.ptt'): + self.diffusion = torch.jit.load(f'{models_dir}/diffusion_decoder.ptt') + else: + if not diffusion_model_path or not os.path.exists(diffusion_model_path): + diffusion_model_path = get_model_path('diffusion_decoder.pth', models_dir) + + self.load_diffusion_model(diffusion_model_path) self.clvp = CLVP(dim_text=768, dim_speech=768, dim_latent=768, num_text_tokens=256, text_enc_depth=20, @@ -366,6 +370,28 @@ class TextToSpeech: self.loading = False print(f"Loaded autoregressive model") + def load_diffusion_model(self, diffusion_model_path): + if hasattr(self,"diffusion_model_path") and self.diffusion_model_path == diffusion_model_path: + return + + self.loading = True + + self.diffusion_model_path = diffusion_model_path if diffusion_model_path and os.path.exists(diffusion_model_path) else get_model_path('diffusion_decoder.pth', self.models_dir) + self.diffusion_model_hash = hash_file(self.diffusion_model_path) + + if hasattr(self, 'diffusion'): + del self.diffusion + + self.diffusion = DiffusionTts(model_channels=1024, num_layers=10, in_channels=100, out_channels=200, + in_latent_channels=1024, in_tokens=8193, dropout=0, use_fp16=False, num_heads=16, + layer_drop=0, unconditioned_percentage=0).cpu().eval() + self.diffusion.load_state_dict(torch.load(get_model_path('diffusion_decoder.pth', self.models_dir))) + if self.preloaded_tensors: + self.diffusion = migrate_to_device( self.diffusion, self.device ) + + self.loading = False + print(f"Loaded diffusion model") + def load_vocoder_model(self, vocoder_model): if hasattr(self,"vocoder_model_path") and self.vocoder_model_path == vocoder_model: return @@ -375,7 +401,7 @@ class TextToSpeech: if hasattr(self, 'vocoder'): del self.vocoder - print(vocoder_model) + print("Loading vocoder model:", vocoder_model) if vocoder_model is None: vocoder_model = 'bigvgan_24khz_100band' @@ -406,6 +432,22 @@ class TextToSpeech: self.loading = False print(f"Loaded vocoder model") + def load_tokenizer_json(self, tokenizer_json): + if hasattr(self,"tokenizer_json") and self.tokenizer_json == tokenizer_json: + return + + self.loading = True + self.tokenizer_json = tokenizer_json if tokenizer_json else os.path.join(os.path.dirname(os.path.realpath(__file__)), '../tortoise/data/tokenizer.json') + print("Loading tokenizer JSON:", self.tokenizer_json) + + if hasattr(self, 'tokenizer'): + del self.tokenizer + + self.tokenizer = VoiceBpeTokenizer(vocab_file=self.tokenizer_json) + + self.loading = False + print(f"Loaded tokenizer") + def load_cvvp(self): """Load CVVP model.""" self.cvvp = CVVP(model_dim=512, transformer_heads=8, dropout=0, mel_codes=8192, conditioning_enc_depth=8, cond_mask_percentage=0,