DL-Art-School/codes/models/gpt_voice/gpt_tts.py

import torch
import torch.nn as nn
import torch.nn.functional as F
from tqdm import tqdm

from models.arch_util import ConvGnSilu
from models.tacotron2.taco_utils import get_mask_from_lengths
from models.tacotron2.text import symbols
from models.gpt_voice.min_gpt import GPT, GPTConfig
from trainer.networks import register_model


# A Conv1d that masks out kernel elements ahead of the current location.
class CausalConv1d(nn.Conv1d):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.kernel_mask = torch.ones_like(self.weight)
        self.kernel_mask[:, :, -(self.kernel_size[0]//2):] = 0

    def forward(self, input):
        self.kernel_mask = self.kernel_mask.to(input.device)
        return self._conv_forward(input, self.weight * self.kernel_mask, self.bias)


class GptTts(nn.Module):
    def __init__(self):
        super().__init__()
        number_symbols = len(symbols)
        model_dim = 512
        max_symbols_per_phrase = 200
        max_mel_frames = 900
        mel_dim=80

        self.model_dim = model_dim
        self.max_mel_frames = max_mel_frames
        self.text_embedding = nn.Embedding(number_symbols, model_dim)
        # Whenever we process MEL frames, we need to be careful to use casually masked convolutions to avoid adding bias
        # into the model which we cannot provide in inference.
        self.mel_encoder = nn.Sequential(ConvGnSilu(mel_dim, model_dim//2, kernel_size=5, convnd=CausalConv1d),
                                         ConvGnSilu(model_dim//2, model_dim, kernel_size=5, stride=2, convnd=CausalConv1d))
        # *_tags are additively applied to
        self.text_tags = nn.Parameter(torch.randn(1, 1, model_dim)/256.0)
        self.separator = nn.Parameter(torch.randn(1, 1, model_dim))
        self.audio_tags = nn.Parameter(torch.randn(1, 1, model_dim)/256.0)
        self.gpt = GPT(GPTConfig(1+max_symbols_per_phrase+max_mel_frames//2, n_embd=model_dim, n_head=8))

        self.gate_head = nn.Sequential(ConvGnSilu(model_dim, model_dim, kernel_size=5, convnd=CausalConv1d),
                                       nn.Upsample(scale_factor=2, mode='nearest'),
                                       ConvGnSilu(model_dim, model_dim//2, kernel_size=5, convnd=CausalConv1d),
                                       # No need for causal convolutions when kernel_size=1
                                       nn.Conv1d(model_dim//2, 1, kernel_size=1))
        self.mel_head = nn.Sequential(ConvGnSilu(model_dim, model_dim, kernel_size=5, convnd=CausalConv1d),
                                      nn.Upsample(scale_factor=2, mode='nearest'),
                                      ConvGnSilu(model_dim, model_dim//2, kernel_size=5, convnd=CausalConv1d),
                                      ConvGnSilu(model_dim//2, model_dim//2, kernel_size=5, convnd=CausalConv1d),
                                      ConvGnSilu(model_dim//2, mel_dim, kernel_size=1, activation=False, norm=False, convnd=nn.Conv1d))

    def forward(self, text_inputs, mel_targets, output_lengths):
        # Pad mel_targets to be a multiple of 2
        padded = mel_targets.shape[-1] % 2 != 0
        if padded:
            mel_targets = F.pad(mel_targets, (0,1))

        text_emb = self.text_embedding(text_inputs)
        text_emb = text_emb + self.text_tags
        mel_emb = self.mel_encoder(mel_targets).permute(0,2,1)
        mel_emb = mel_emb + self.audio_tags
        emb = torch.cat([text_emb,
                         self.separator.repeat(text_emb.shape[0],1,1),
                         mel_emb], dim=1)
        enc = self.gpt(emb)
        mel_portion = enc[:, text_emb.shape[1]+1:].permute(0,2,1)
        gates = self.gate_head(mel_portion).squeeze(1)
        mel_pred = self.mel_head(mel_portion)

        # Mask portions of output which we don't need to predict.
        mask = ~get_mask_from_lengths(output_lengths, mel_pred.shape[-1])
        mask = mask.unsqueeze(1).repeat(1, mel_pred.shape[1], 1)
        mel_pred.data.masked_fill_(mask, 0)
        gates.data.masked_fill_(mask[:, 0, :], 1e3)

        if padded:
            mel_pred = mel_pred[:, :, :-1]
            gates = gates[:, :-1]
        return mel_pred, gates

    def test_guide(self, mel_guide, amount=50):
        mel_guide = mel_guide[:,:,:amount]
        mel_emb = self.mel_encoder(mel_guide).permute(0,2,1)
        mel_emb = mel_emb + self.audio_tags
        return mel_emb

    def inference(self, text_inputs, mel_guide):
        MEL_HEAD_EXPANSION = 2
        GATE_THRESHOLD = .95

        text_emb = self.text_embedding(text_inputs)
        text_emb = text_emb + self.text_tags
        b,s,c = text_emb.shape
        emb = torch.cat([text_emb,
                         self.separator.repeat(text_emb.shape[0],1,1)], dim=1)
                         #self.test_guide(mel_guide)], dim=1)
        completed = torch.zeros((b,), device=text_inputs.device, dtype=torch.bool)
        output = None
        for i in tqdm(range(self.max_mel_frames)):
            enc = self.gpt(emb)
            inferred = enc[:,s:,:].permute(0,2,1)
            # Create output frames.
            inferred_mel_frame = self.mel_head(inferred)[:,:,-MEL_HEAD_EXPANSION:]
            inferred_mel_frame = inferred_mel_frame * (~completed).float().view(b,1,1)
            if output is None:
                output = inferred_mel_frame
            else:
                output = torch.cat([output, inferred_mel_frame], dim=2)

            # Test termination condition
            gate = F.sigmoid(self.gate_head(inferred)).max(dim=-1).values  # TODO: accept single-frame terminations.
            completed = completed.logical_or((gate > GATE_THRESHOLD).squeeze(1))  # This comprises a latch - but that may not be wise.
            if torch.all(completed):
                break

            # Apply inferred mel_frames to emb for next pass.
            mel_emb = self.mel_encoder(output).permute(0,2,1)
            mel_emb = mel_emb + self.audio_tags
            emb = torch.cat([text_emb,
                             self.separator.repeat(text_emb.shape[0],1,1),
                             mel_emb], dim=1)
            if i == self.max_mel_frames//2:
                print("Warning! Inference hit mel frame cap without encountering a stop token.")
                break

        return output


@register_model
def register_gpt_tts(opt_net, opt):
    return GptTts()


if __name__ == '__main__':
    gpt = GptTts()
    m, g = gpt(torch.randint(high=24, size=(2,60)),
               torch.randn(2,80,747),
               torch.tensor([600,747]))
    print(m.shape)
    print(g.shape)

    o = gpt.infer(torch.randint(high=24, size=(2,60)))
    print(o.shape)