port do_tts to use the API

api.py

@@ -151,10 +151,10 @@ class TextToSpeech:
 
     def tts(self, text, voice_samples, k=1,
             # autoregressive generation parameters follow
-            num_autoregressive_samples=512, temperature=.9, length_penalty=1, repetition_penalty=1.0, top_k=50, top_p=.95,
+            num_autoregressive_samples=512, temperature=.5, length_penalty=2, repetition_penalty=2.0, top_p=.5,
             typical_sampling=False, typical_mass=.9,
             # diffusion generation parameters follow
-            diffusion_iterations=100, cond_free=True, cond_free_k=1, diffusion_temperature=1,):
+            diffusion_iterations=100, cond_free=True, cond_free_k=2, diffusion_temperature=.7,):
         text = torch.IntTensor(self.tokenizer.encode(text)).unsqueeze(0).cuda()
         text = F.pad(text, (0, 1))  # This may not be necessary.
 
@@ -181,7 +181,6 @@ class TextToSpeech:
             for b in tqdm(range(num_batches)):
                 codes = self.autoregressive.inference_speech(conds, text,
                                                              do_sample=True,
-                                                             top_k=top_k,
                                                              top_p=top_p,
                                                              temperature=temperature,
                                                              num_return_sequences=self.autoregressive_batch_size,
@@ -220,4 +219,28 @@ class TextToSpeech:
 
             if len(wav_candidates) > 1:
                 return wav_candidates
-            return wav_candidates[0]
+            return wav_candidates[0]
+
+    def refine_for_intellibility(self, wav_candidates, corresponding_codes, output_path):
+        """
+        Further refine the remaining candidates using a ASR model to pick out the ones that are the most understandable.
+        TODO: finish this function
+        :param wav_candidates:
+        :return:
+        """
+        transcriber = ocotillo.Transcriber(on_cuda=True)
+        transcriptions = transcriber.transcribe_batch(torch.cat(wav_candidates, dim=0).squeeze(1), 24000)
+        best = 99999999
+        for i, transcription in enumerate(transcriptions):
+            dist = lev_distance(transcription, args.text.lower())
+            if dist < best:
+                best = dist
+                best_codes = corresponding_codes[i].unsqueeze(0)
+                best_wav = wav_candidates[i]
+        del transcriber
+        torchaudio.save(os.path.join(output_path, f'{voice}_poor.wav'), best_wav.squeeze(0).cpu(), 24000)
+
+        # Perform diffusion again with the high-quality diffuser.
+        mel = do_spectrogram_diffusion(diffusion, final_diffuser, best_codes, cond_diffusion, mean=False)
+        wav = vocoder.inference(mel)
+        torchaudio.save(os.path.join(args.output_path, f'{voice}.wav'), wav.squeeze(0).cpu(), 24000)

do_tts.py

@@ -1,123 +1,13 @@
 import argparse
 import os
-import random
-from urllib import request
 
 import torch
 import torch.nn.functional as F
 import torchaudio
-import progressbar
-import ocotillo
-
-from models.diffusion_decoder import DiffusionTts
-from models.autoregressive import UnifiedVoice
-from tqdm import tqdm
-
-from models.arch_util import TorchMelSpectrogram
-from models.text_voice_clip import VoiceCLIP
-from models.vocoder import UnivNetGenerator
-from utils.audio import load_audio, wav_to_univnet_mel, denormalize_tacotron_mel
-from utils.diffusion import SpacedDiffusion, space_timesteps, get_named_beta_schedule
-from utils.tokenizer import VoiceBpeTokenizer, lev_distance
-
-pbar = None
-def download_models():
-    MODELS = {
-        'clip.pth': 'https://huggingface.co/jbetker/tortoise-tts-clip/resolve/main/pytorch-model.bin',
-        'diffusion.pth': 'https://huggingface.co/jbetker/tortoise-tts-diffusion-v1/resolve/main/pytorch-model.bin',
-        'autoregressive.pth': 'https://huggingface.co/jbetker/tortoise-tts-autoregressive/resolve/main/pytorch-model.bin'
-    }
-    os.makedirs('.models', exist_ok=True)
-    def show_progress(block_num, block_size, total_size):
-        global pbar
-        if pbar is None:
-            pbar = progressbar.ProgressBar(maxval=total_size)
-            pbar.start()
-
-        downloaded = block_num * block_size
-        if downloaded < total_size:
-            pbar.update(downloaded)
-        else:
-            pbar.finish()
-            pbar = None
-    for model_name, url in MODELS.items():
-        if os.path.exists(f'.models/{model_name}'):
-            continue
-        print(f'Downloading {model_name} from {url}...')
-        request.urlretrieve(url, f'.models/{model_name}', show_progress)
-        print('Done.')
-
-
-def load_discrete_vocoder_diffuser(trained_diffusion_steps=4000, desired_diffusion_steps=200, cond_free=True):
-    """
-    Helper function to load a GaussianDiffusion instance configured for use as a vocoder.
-    """
-    return SpacedDiffusion(use_timesteps=space_timesteps(trained_diffusion_steps, [desired_diffusion_steps]), model_mean_type='epsilon',
-                           model_var_type='learned_range', loss_type='mse', betas=get_named_beta_schedule('linear', trained_diffusion_steps),
-                           conditioning_free=cond_free, conditioning_free_k=1)
-
-
-def load_conditioning(path, sample_rate=22050, cond_length=132300):
-    rel_clip = load_audio(path, sample_rate)
-    gap = rel_clip.shape[-1] - cond_length
-    if gap < 0:
-        rel_clip = F.pad(rel_clip, pad=(0, abs(gap)))
-    elif gap > 0:
-        rand_start = random.randint(0, gap)
-        rel_clip = rel_clip[:, rand_start:rand_start + cond_length]
-    mel_clip = TorchMelSpectrogram()(rel_clip.unsqueeze(0)).squeeze(0)
-    return mel_clip.unsqueeze(0).cuda(), rel_clip.unsqueeze(0).cuda()
-
-
-def fix_autoregressive_output(codes, stop_token):
-    """
-    This function performs some padding on coded audio that fixes a mismatch issue between what the diffusion model was
-    trained on and what the autoregressive code generator creates (which has no padding or end).
-    This is highly specific to the DVAE being used, so this particular coding will not necessarily work if used with
-    a different DVAE. This can be inferred by feeding a audio clip padded with lots of zeros on the end through the DVAE
-    and copying out the last few codes.
-
-    Failing to do this padding will produce speech with a harsh end that sounds like "BLAH" or similar.
-    """
-    # Strip off the autoregressive stop token and add padding.
-    stop_token_indices = (codes == stop_token).nonzero()
-    if len(stop_token_indices) == 0:
-        print("No stop tokens found, enjoy that output of yours!")
-        return
-    else:
-        codes[stop_token_indices] = 83
-    stm = stop_token_indices.min().item()
-    codes[stm:] = 83
-    if stm - 3 < codes.shape[0]:
-        codes[-3] = 45
-        codes[-2] = 45
-        codes[-1] = 248
-
-    return codes
-
-
-def do_spectrogram_diffusion(diffusion_model, diffuser, mel_codes, conditioning_input, mean=False):
-    """
-    Uses the specified diffusion model and DVAE model to convert the provided MEL & conditioning inputs into an audio clip.
-    """
-    with torch.no_grad():
-        cond_mel = wav_to_univnet_mel(conditioning_input.squeeze(1), do_normalization=False)
-        # Pad MEL to multiples of 32
-        msl = mel_codes.shape[-1]
-        dsl = 32
-        gap = dsl - (msl % dsl)
-        if gap > 0:
-            mel = torch.nn.functional.pad(mel_codes, (0, gap))
-
-        output_shape = (mel.shape[0], 100, mel.shape[-1]*4)
-        precomputed_embeddings = diffusion_model.timestep_independent(mel_codes, cond_mel)
-        if mean:
-            mel = diffuser.p_sample_loop(diffusion_model, output_shape, noise=torch.zeros(output_shape, device=mel_codes.device),
-                                          model_kwargs={'precomputed_aligned_embeddings': precomputed_embeddings})
-        else:
-            mel = diffuser.p_sample_loop(diffusion_model, output_shape, model_kwargs={'precomputed_aligned_embeddings': precomputed_embeddings})
-        return denormalize_tacotron_mel(mel)[:,:,:msl*4]
 
+from api import TextToSpeech, load_conditioning
+from utils.audio import load_audio
+from utils.tokenizer import VoiceBpeTokenizer
 
 if __name__ == '__main__':
     # These are voices drawn randomly from the training set. You are free to substitute your own voices in, but testing
@@ -139,101 +29,23 @@ if __name__ == '__main__':
     parser.add_argument('-text', type=str, help='Text to speak.', default="I am a language model that has learned to speak.")
     parser.add_argument('-voice', type=str, help='Use a preset conditioning voice (defined above). Overrides cond_path.', default='dotrice,harris,lescault,otto,atkins,grace,kennard,mol')
     parser.add_argument('-num_samples', type=int, help='How many total outputs the autoregressive transformer should produce.', default=512)
-    parser.add_argument('-num_batches', type=int, help='How many batches those samples should be produced over.', default=16)
+    parser.add_argument('-batch_size', type=int, help='How many samples to process at once in the autoregressive model.', default=16)
     parser.add_argument('-num_diffusion_samples', type=int, help='Number of outputs that progress to the diffusion stage.', default=16)
     parser.add_argument('-output_path', type=str, help='Where to store outputs.', default='results/')
     args = parser.parse_args()
-
     os.makedirs(args.output_path, exist_ok=True)
-    download_models()
+
+    tts = TextToSpeech(autoregressive_batch_size=args.batch_size)
 
     for voice in args.voice.split(','):
-        print("Loading data..")
         tokenizer = VoiceBpeTokenizer()
         text = torch.IntTensor(tokenizer.encode(args.text)).unsqueeze(0).cuda()
         text = F.pad(text, (0,1))  # This may not be necessary.
         cond_paths = preselected_cond_voices[voice]
         conds = []
         for cond_path in cond_paths:
-            c, cond_wav = load_conditioning(cond_path)
+            c = load_audio(cond_path, 22050)
             conds.append(c)
-        conds = torch.stack(conds, dim=1)
-        cond_diffusion = cond_wav[:, :88200]  # The diffusion model expects <= 88200 conditioning samples.
-
-        print("Loading GPT TTS..")
-        autoregressive = UnifiedVoice(max_mel_tokens=300, max_text_tokens=200, max_conditioning_inputs=2, layers=30, model_dim=1024,
-                                      heads=16, number_text_tokens=256, start_text_token=255, checkpointing=False, train_solo_embeddings=False,
-                                      average_conditioning_embeddings=True).cuda().eval()
-        autoregressive.load_state_dict(torch.load('.models/autoregressive.pth'))
-        stop_mel_token = autoregressive.stop_mel_token
-
-        with torch.no_grad():
-            print("Performing autoregressive inference..")
-            samples = []
-            for b in tqdm(range(args.num_batches)):
-                codes = autoregressive.inference_speech(conds, text, num_beams=1, repetition_penalty=1.0, do_sample=True, top_k=50, top_p=.95,
-                                                        temperature=.9, num_return_sequences=args.num_samples//args.num_batches, length_penalty=1)
-                padding_needed = 250 - codes.shape[1]
-                codes = F.pad(codes, (0, padding_needed), value=stop_mel_token)
-                samples.append(codes)
-            del autoregressive
-
-            print("Loading CLIP..")
-            clip = VoiceCLIP(dim_text=512, dim_speech=512, dim_latent=512, num_text_tokens=256, text_enc_depth=12, text_seq_len=350, text_heads=8,
-                             num_speech_tokens=8192, speech_enc_depth=12, speech_heads=8, speech_seq_len=430, use_xformers=True).cuda().eval()
-            clip.load_state_dict(torch.load('.models/clip.pth'))
-            print("Performing CLIP filtering..")
-            clip_results = []
-            for batch in samples:
-                for i in range(batch.shape[0]):
-                    batch[i] = fix_autoregressive_output(batch[i], stop_mel_token)
-                clip_results.append(clip(text.repeat(batch.shape[0], 1), batch, return_loss=False))
-            clip_results = torch.cat(clip_results, dim=0)
-            samples = torch.cat(samples, dim=0)
-            best_results = samples[torch.topk(clip_results, k=args.num_diffusion_samples).indices]
-
-            # Delete the autoregressive and clip models to free up GPU memory
-            del samples, clip
-
-            print("Loading Diffusion Model..")
-            diffusion = DiffusionTts(model_channels=512, in_channels=100, out_channels=200, in_latent_channels=1024,
-                                     channel_mult=[1, 2, 3, 4], num_res_blocks=[3, 3, 3, 3], token_conditioning_resolutions=[1,4,8],
-                                     dropout=0, attention_resolutions=[4,8], num_heads=8, kernel_size=3, scale_factor=2,
-                                     time_embed_dim_multiplier=4, unconditioned_percentage=0, conditioning_dim_factor=2,
-                                     conditioning_expansion=1)
-            diffusion.load_state_dict(torch.load('.models/diffusion.pth'))
-            diffusion = diffusion.cuda().eval()
-            print("Loading vocoder..")
-            vocoder = UnivNetGenerator()
-            vocoder.load_state_dict(torch.load('.models/vocoder.pth')['model_g'])
-            vocoder = vocoder.cuda()
-            vocoder.eval(inference=True)
-            initial_diffuser = load_discrete_vocoder_diffuser(desired_diffusion_steps=40, cond_free=False)
-            final_diffuser = load_discrete_vocoder_diffuser(desired_diffusion_steps=500)
-
-            print("Performing vocoding..")
-            wav_candidates = []
-            for b in range(best_results.shape[0]):
-                code = best_results[b].unsqueeze(0)
-                mel = do_spectrogram_diffusion(diffusion, initial_diffuser, code, cond_diffusion, mean=False)
-                wav = vocoder.inference(mel)
-                wav_candidates.append(wav.cpu())
-
-            # Further refine the remaining candidates using a ASR model to pick out the ones that are the most understandable.
-            transcriber = ocotillo.Transcriber(on_cuda=True)
-            transcriptions = transcriber.transcribe_batch(torch.cat(wav_candidates, dim=0).squeeze(1), 24000)
-            best = 99999999
-            for i, transcription in enumerate(transcriptions):
-                dist = lev_distance(transcription, args.text.lower())
-                if dist < best:
-                    best = dist
-                    best_codes = best_results[i].unsqueeze(0)
-                    best_wav = wav_candidates[i]
-            del transcriber
-            torchaudio.save(os.path.join(args.output_path, f'{voice}_poor.wav'), best_wav.squeeze(0).cpu(), 24000)
-
-            # Perform diffusion again with the high-quality diffuser.
-            mel = do_spectrogram_diffusion(diffusion, final_diffuser, best_codes, cond_diffusion, mean=False)
-            wav = vocoder.inference(mel)
-            torchaudio.save(os.path.join(args.output_path, f'{voice}.wav'), wav.squeeze(0).cpu(), 24000)
+        gen = tts.tts(args.text, conds, num_autoregressive_samples=args.num_samples)
+        torchaudio.save(os.path.join(args.output_path, f'{voice}.wav'), gen.squeeze(0).cpu(), 24000)