VoiceLDM: Text-to-Speech with Environmental Context
import os
import argparse

import torch
import torchaudio

from voiceldm import VoiceLDMPipeline


def parse_args(parser):

    parser.add_argument('--desc_prompt', '-d', type=str, help='description prompt')
    parser.add_argument('--cont_prompt', '-c', type=str, help='content prompt')
    parser.add_argument('--audio_prompt', '-a', type=str, help='path to audio file to be used as audio prompt')

    parser.add_argument("--model_config", type=str, default="m", help="configuration for VoiceLDM. 'm' for VoiceLDM-M and 's' for VoiceLDM-S")
    parser.add_argument("--ckpt_path", type=str, help="checkpoint file path for VoiceLDM")
    parser.add_argument("--output_dir", type=str, default="./outputs", help="directory to save generated audio")
    parser.add_argument("--file_name", type=str, help="filename for the generated audio")

    parser.add_argument('--num_inference_steps', type=int, default=50, help='number of inference steps for DDIM sampling')
    parser.add_argument('--audio_length_in_s', type=float, default=10, help='duration of the audio for generation')
    parser.add_argument('--guidance_scale', type=float, help='guidance weight for single classifier-free guidance')
    parser.add_argument('--desc_guidance_scale', type=float, default=5, required=False, help='desc guidance weight for dual classifier-free guidance')
    parser.add_argument('--cont_guidance_scale', type=float, default=7, required=False, help='cont guidance weight for dual classifier-free guidance')
    
    parser.add_argument("--device", type=str, default="auto", help="device to use for audio generation")
    parser.add_argument('--seed', type=int, help='random seed for generation')

    return parser.parse_args()


def main():
    parser = argparse.ArgumentParser()
    args = parse_args(parser)

    if args.device == "auto":
        if torch.cuda.is_available():
            args.device = torch.device("cuda:0")
        else:
            args.device = torch.device("cpu")
    elif args.device is not None:
        args.device = torch.device(device)
    else:
        args.device = torch.device("cpu")

    pipe = VoiceLDMPipeline(
        args.model_config,
        args.ckpt_path,
        args.device,
    )

    audio = pipe(**vars(args))

    file_name = args.file_name
    if file_name is None:
        file_name = args.desc_prompt or args.audio_prompt
        file_name = file_name + "-" + args.cont_prompt + ".wav"

    os.makedirs(args.output_dir, exist_ok=True)
    save_path = os.path.join(args.output_dir, file_name)
    torchaudio.save(save_path, src=audio, sample_rate=16000)


if __name__ == "__main__":
    main()
generate.py
glory20h/VoiceLDM
from .pipeline import VoiceLDMPipeline
voiceldm/__init__.py
glory20h/VoiceLDM
# Codes adapted from https://github.com/heatz123/naturalspeech/blob/main/models/models.py

import math
import torch
import torch.nn as nn
import torch.nn.functional as F


class LayerNorm(nn.Module):
    def __init__(self, channels, eps=1e-5):
        super().__init__()
        self.channels = channels
        self.eps = eps

        self.gamma = nn.Parameter(torch.ones(channels))
        self.beta = nn.Parameter(torch.zeros(channels))

    def forward(self, x):
        x = x.transpose(1, -1)
        x = F.layer_norm(x, (self.channels,), self.gamma, self.beta, self.eps)
        return x.transpose(1, -1)


class ConvBlock(nn.Module):
    def __init__(
        self, in_channels, out_channels, kernel_size, dropout, activation=nn.ReLU()
    ):
        super(ConvBlock, self).__init__()
        
        self.conv_layer = nn.Sequential(
            nn.Conv1d(
                in_channels,
                out_channels,
                kernel_size=kernel_size,
                stride=1,
                padding=int((kernel_size - 1) / 2),
                dilation=1,
            ),
            nn.BatchNorm1d(out_channels),
            activation,
        )
        self.dropout = dropout
        self.layer_norm = nn.LayerNorm(out_channels)
        
    def forward(self, x, mask=None):
        x = x.contiguous().transpose(1, 2)
        x = F.dropout(self.conv_layer(x), self.dropout, self.training)
        x = self.layer_norm(x.contiguous().transpose(1, 2))
        if mask is not None:
            x = x.masked_fill(mask.unsqueeze(-1), 0)
            
        return x
    
    
class SwishBlock(nn.Module):
    def __init__(self, in_channels, hidden_dim, out_channels):
        super(SwishBlock, self).__init__()
        self.layer = nn.Sequential(
            nn.Linear(in_channels, hidden_dim, bias=True),
            nn.SiLU(),
            nn.Linear(hidden_dim, hidden_dim, bias=True),
            nn.SiLU(),
            nn.Linear(hidden_dim, out_channels, bias=True),
        )
        
    def forward(self, S, E, V):
        out = torch.cat(
            [
                S.unsqueeze(-1),
                E.unsqueeze(-1),
                V.unsqueeze(1).expand(-1, E.size(1), -1, -1),
            ],
            dim=-1,
        )
        out = self.layer(out)
        
        return out


class DurationPredictor(nn.Module):
    def __init__(
        self, in_channels, filter_channels, kernel_size, p_dropout, gin_channels=0
    ):
        super().__init__()

        self.in_channels = in_channels
        self.filter_channels = filter_channels
        self.kernel_size = kernel_size
        self.p_dropout = p_dropout
        self.gin_channels = gin_channels

        self.drop = nn.Dropout(p_dropout)
        self.conv_1 = nn.Conv1d(
            in_channels, filter_channels, kernel_size, padding=kernel_size // 2
        )
        self.norm_1 = LayerNorm(filter_channels)
        self.conv_2 = nn.Conv1d(
            filter_channels, filter_channels, kernel_size, padding=kernel_size // 2
        )
        self.norm_2 = LayerNorm(filter_channels)
        self.proj = nn.Conv1d(filter_channels, 1, 1)

        if gin_channels != 0:
            self.cond = nn.Conv1d(filter_channels, 1, 1)

    def forward(self, x, x_mask, g=None):
        x = torch.detach(x)
        if g is not None:
            g = torch.detach(g)
            x = x + self.cond(g)
        x = self.conv_1(x * x_mask)
        x = torch.relu(x)
        x = self.norm_1(x)
        x = self.drop(x)
        x = self.conv_2(x * x_mask)
        x = torch.relu(x)
        x = self.norm_2(x)
        x = self.drop(x)
        x = self.proj(x * x_mask)
        return x * x_mask


class LearnableUpsampler(nn.Module):
    def __init__(
        self,
        d_predictor=192,
        out_channels=192,
        kernel_size=3,
        dropout=0.0,
        conv_output_size=8,
        dim_w=4,
        dim_c=2,
        max_seq_len=1000,
    ):
        super(LearnableUpsampler, self).__init__()
        self.max_seq_len = max_seq_len

        # Attention (W)
        self.conv_w = ConvBlock(
            d_predictor,
            conv_output_size,
            kernel_size,
            dropout=dropout,
            activation=nn.SiLU(),
        )
        self.swish_w = SwishBlock(conv_output_size + 2, dim_w, dim_w)
        self.linear_w = nn.Linear(dim_w * d_predictor, d_predictor, bias=True)
        self.softmax_w = nn.Softmax(dim=2)

        # Auxiliary Attention Context (C)
        self.conv_c = ConvBlock(
            d_predictor,
            conv_output_size,
            kernel_size,
            dropout=dropout,
            activation=nn.SiLU(),
        )
        self.swish_c = SwishBlock(conv_output_size + 2, dim_c, dim_c)

        # Upsampled Representation (O)
        self.linear_einsum = nn.Linear(dim_c * dim_w, d_predictor)  # A
        self.layer_norm = nn.LayerNorm(d_predictor)

        self.proj_o = nn.Linear(d_predictor, out_channels)

    def forward(self, duration, V, src_len, src_mask, max_src_len):
        batch_size = duration.shape[0]

        # Duration Interpretation
        mel_len = torch.round(duration.sum(-1)).type(torch.LongTensor).to(V.device)
        mel_len = torch.clamp(mel_len, max=self.max_seq_len)
        max_mel_len = mel_len.max().item()
        mel_mask = self.get_mask_from_lengths(mel_len, max_mel_len)

        # Prepare Attention Mask
        src_mask_ = src_mask.unsqueeze(1).expand(
            -1, mel_mask.shape[1], -1
        )  # [B, tgt_len, src_len]
        mel_mask_ = mel_mask.unsqueeze(-1).expand(
            -1, -1, src_mask.shape[1]
        )  # [B, tgt_len, src_len]
        attn_mask = torch.zeros(
            (src_mask.shape[0], mel_mask.shape[1], src_mask.shape[1])
        ).to(V.device)
        attn_mask = attn_mask.masked_fill(src_mask_, 1.0)
        attn_mask = attn_mask.masked_fill(mel_mask_, 1.0)
        attn_mask = attn_mask.bool()

        # Token Boundary Grid
        e_k = torch.cumsum(duration, dim=1)
        s_k = e_k - duration
        e_k = e_k.unsqueeze(1).expand(batch_size, max_mel_len, -1)
        s_k = s_k.unsqueeze(1).expand(batch_size, max_mel_len, -1)
        t_arange = (
            torch.arange(1, max_mel_len + 1, device=V.device)
            .unsqueeze(0)
            .unsqueeze(-1)
            .expand(batch_size, -1, max_src_len)
        )

        S, E = (t_arange - s_k).masked_fill(attn_mask, 0), (e_k - t_arange).masked_fill(
            attn_mask, 0
        )

        # Attention (W)
        W = self.swish_w(S, E, self.conv_w(V))  # [B, T, K, dim_w]
        W = W.masked_fill(src_mask_.unsqueeze(-1), -float("Inf"))
        W = self.softmax_w(W)  # [B, T, K]
        W = W.masked_fill(mel_mask_.unsqueeze(-1), 0.0)
        W = W.permute(0, 3, 1, 2)

        # Auxiliary Attention Context (C)
        C = self.swish_c(S, E, self.conv_c(V))  # [B, T, K, dim_c]

        # Upsampled Representation (O)
        upsampled_rep = self.linear_w(
            torch.einsum("bqtk,bkh->bqth", W, V).permute(0, 2, 1, 3).flatten(2)
        ) + self.linear_einsum(
            torch.einsum("bqtk,btkp->bqtp", W, C).permute(0, 2, 1, 3).flatten(2)
        )  # [B, T, M]
        upsampled_rep = self.layer_norm(upsampled_rep)
        upsampled_rep = upsampled_rep.masked_fill(mel_mask.unsqueeze(-1), 0)
        upsampled_rep = self.proj_o(upsampled_rep)

        return upsampled_rep, mel_mask, mel_len, W

    def get_mask_from_lengths(self, lengths, max_len=None):
        batch_size = lengths.shape[0]
        if max_len is None:
            max_len = torch.max(lengths).item()

        ids = (
            torch.arange(0, max_len)
            .unsqueeze(0)
            .expand(batch_size, -1)
            .to(lengths.device)
        )
        mask = ids >= lengths.unsqueeze(1).expand(-1, max_len)

        return mask


class Durator(nn.Module):
    def __init__(
        self,
        in_channels=768,
        out_channels=768,
        kernel_size=3,
        dropout=0.5,
        conv_output_size=8,
        dim_w=4,
        dim_c=2,
        max_seq_len=1000,
    ):
        super().__init__()
        self.dp = DurationPredictor(
            in_channels=in_channels,
            filter_channels=in_channels,
            kernel_size=kernel_size,
            p_dropout=dropout,
        )
        
        self.learnable_upsampler = LearnableUpsampler(
            d_predictor=in_channels,
            out_channels=out_channels,
            kernel_size=kernel_size,
            dropout=dropout,
            conv_output_size=conv_output_size,
            dim_w=dim_w,
            dim_c=dim_c,
            max_seq_len=max_seq_len,
        )
        
    def forward(self, x, x_mask):
        logw = self.dp(x.transpose(1,2), x_mask)
        logw = torch.clamp(logw, max=3.5)
        w = torch.clamp(torch.exp(logw), min=1.0) * x_mask
        duration = w.squeeze(1)
        
        upsampled_rep, p_mask, _, W = self.learnable_upsampler(
            duration,
            x,
            x_mask.sum(-1),
            ~(x_mask.squeeze(1).bool()),
            x_mask.shape[-1],
        )
        
        return upsampled_rep, p_mask


class TextEncoder(nn.Module):
    def __init__(self, speecht5):
        super().__init__()
        self.encoder = speecht5.speecht5.encoder
        
        # Replace positional encoding with longer max_len
        max_len = 1000
        dim = self.encoder.config.hidden_size # 768
        pe = torch.zeros(max_len, dim)
        position = torch.arange(0, max_len).unsqueeze(1)
        div_term = torch.exp((torch.arange(0, dim, 2, dtype=torch.float) * -(math.log(10000.0) / dim)))
        pe[:, 0::2] = torch.sin(position.float() * div_term)
        pe[:, 1::2] = torch.cos(position.float() * div_term)
        pe = pe.unsqueeze(0)
        self.encoder.prenet.encode_positions.register_buffer("pe", pe)
        
    def forward(self, tokens):
        x = self.encoder(
            input_values=tokens.input_ids,
            attention_mask=tokens.attention_mask,
            return_dict=True,
        ).last_hidden_state
        return x


class UNetWrapper(nn.Module):
    def __init__(self, unet, text_encoder=None):
        super().__init__()
        self.unet = unet
        self.durator = Durator(
            in_channels=768,
            out_channels=768,
            kernel_size=3,
            dropout=0.5,
        )
        self.text_encoder = text_encoder
        
    def forward(self, sample, timestep, text_embed, text_embed_mask, cond):
        sample = self.unet(
            sample, 
            timestep, 
            encoder_hidden_states=text_embed,
            class_labels=cond,
        ).sample
    
        return sample
voiceldm/modules.py
glory20h/VoiceLDM
import torch
import torchaudio
import numpy as np
from tqdm.auto import tqdm

from transformers import (
    ClapModel,
    ClapProcessor,
    SpeechT5Processor,
    SpeechT5ForTextToSpeech,
    SpeechT5HifiGan,
)
from diffusers import (
    AutoencoderKL, 
    DDIMScheduler, 
    UNet2DConditionModel
)
from diffusers.utils import randn_tensor
from huggingface_hub import hf_hub_download

from .modules import UNetWrapper, TextEncoder


class VoiceLDMPipeline():
    def __init__(
        self,
        model_config = None,
        ckpt_path = None,
        device = None,
    ):
        if model_config is None:
            model_config = "m"

        self.noise_scheduler = DDIMScheduler.from_pretrained("cvssp/audioldm-m-full", subfolder="scheduler")
        self.vae = AutoencoderKL.from_pretrained("cvssp/audioldm-m-full", subfolder="vae").eval()
        self.vocoder = SpeechT5HifiGan.from_pretrained("cvssp/audioldm-m-full", subfolder="vocoder").eval()
        self.clap_model = ClapModel.from_pretrained("laion/clap-htsat-unfused").eval()
        self.clap_processor = ClapProcessor.from_pretrained("laion/clap-htsat-unfused")
        self.text_encoder = TextEncoder(SpeechT5ForTextToSpeech.from_pretrained("microsoft/speecht5_tts")).eval()
        self.text_processor = SpeechT5Processor.from_pretrained("microsoft/speecht5_tts")

        unet = UNet2DConditionModel(
            sample_size = 128,
            in_channels = 8,
            out_channels = 8,
            down_block_types = (
                "DownBlock2D",
                "CrossAttnDownBlock2D",
                "CrossAttnDownBlock2D",
                "CrossAttnDownBlock2D",
            ),
            mid_block_type = "UNetMidBlock2DCrossAttn",
            up_block_types = (
                "CrossAttnUpBlock2D", 
                "CrossAttnUpBlock2D", 
                "CrossAttnUpBlock2D",
                "UpBlock2D",
            ),
            only_cross_attention = False,
            block_out_channels = [128, 256, 384, 640] if model_config == "s" else [192, 384, 576, 960],
            layers_per_block = 2,
            cross_attention_dim = 768,
            class_embed_type = 'simple_projection',
            projection_class_embeddings_input_dim = 512,
            class_embeddings_concat = True,
        )

        if ckpt_path is None:
            ckpt_path = hf_hub_download(
                repo_id=f"glory20h/voiceldm-{model_config}",
                filename=f"voiceldm-{model_config}.ckpt"
            )

        # TODO: Get checkpoints
        def load_ckpt(model, ckpt_path):
            ckpt = torch.load(ckpt_path, map_location="cpu")
            model.load_state_dict(ckpt)
            return model

        self.model = load_ckpt(UNetWrapper(unet, self.text_encoder), ckpt_path)

        self.device = device
        self.vae.to(device)
        self.vocoder.to(device)
        self.clap_model.to(device)
        self.text_encoder.to(device)
        self.model.to(device)

    def prepare_latents(self, batch_size, num_channels_latents, height, dtype, device, generator, latents=None):
        vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
        shape = (
            batch_size,
            num_channels_latents,
            height // vae_scale_factor,
            self.vocoder.config.model_in_dim // vae_scale_factor,
        )
        if isinstance(generator, list) and len(generator) != batch_size:
            raise ValueError(
                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
            )

        if latents is None:
            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
        else:
            latents = latents.to(device)

        latents = latents * self.noise_scheduler.init_noise_sigma
        return latents

    def decode_latents(self, latents):
        latents = 1 / self.vae.config.scaling_factor * latents
        mel_spectrogram = self.vae.decode(latents).sample
        return mel_spectrogram

    def mel_spectrogram_to_waveform(self, mel_spectrogram):
        if mel_spectrogram.dim() == 4:
            mel_spectrogram = mel_spectrogram.squeeze(1)

        waveform = self.vocoder(mel_spectrogram)
        waveform = waveform.cpu().float()
        return waveform

    def normalize_wav(self, waveform):
        waveform = waveform - torch.mean(waveform)
        waveform = waveform / (torch.max(torch.abs(waveform)) + 1e-8)
        return waveform

    def __call__(
        self,
        desc_prompt,
        cont_prompt,
        audio_prompt = None,
        batch_size = 1,
        num_inference_steps = 100,
        audio_length_in_s = 10,
        do_classifier_free_guidance = True,
        guidance_scale = None,
        desc_guidance_scale = None,
        cont_guidance_scale = None,
        device=None,
        seed=None,
        **kwargs,
    ):
        if guidance_scale is None and desc_guidance_scale is None and cont_guidance_scale is None:
            do_classifier_free_guidance = False

        guidance = None
        if guidance_scale is None:
            guidance = "dual"
        else:
            guidance = "single"

        # description condition
        if audio_prompt is None:
            if do_classifier_free_guidance:
                if guidance == "dual":
                    desc_prompt = [desc_prompt] * 2 + [""] * 2
                if guidance == "single":
                    desc_prompt = [desc_prompt] + [""]

            clap_inputs = self.clap_processor(
                text=desc_prompt, 
                return_tensors="pt", 
                padding=True
            ).to(self.device)
            embeds = self.clap_model.text_model(**clap_inputs).pooler_output
            c_desc = self.clap_model.text_projection(embeds)
        else:
            audio_sample, sr = torchaudio.load(audio_prompt)
            if sr != 48000:
                audio_sample = torchaudio.functional.resample(audio_sample, orig_freq=sr, new_freq=48000)
            audio_sample = audio_sample[0]

            clap_inputs = self.clap_processor(audios=audio_sample, sampling_rate=48000, return_tensors="pt", padding=True).to(self.device)

            embeds = self.clap_model.audio_model(**clap_inputs).pooler_output
            c_desc = self.clap_model.audio_projection(embeds)

            if do_classifier_free_guidance:
                clap_inputs = self.clap_processor(text=[""], return_tensors="pt", padding=True).to(self.device)
                uncond_embeds = self.clap_model.text_model(**clap_inputs).pooler_output
                uc_desc = self.clap_model.text_projection(uncond_embeds)
                
                if guidance == "dual":
                    c_desc = torch.cat((c_desc, c_desc, uc_desc, uc_desc))
                if guidance == "single":
                    c_desc = torch.cat((c_desc, uc_desc))
                
        # content condition
        if do_classifier_free_guidance:
            if guidance == "dual":
                cont_prompt = ([cont_prompt] + ["_"]) * 2
            if guidance == "single":
                cont_prompt = [cont_prompt] + ["_"]

        cont_tokens = self.text_processor(
            text=cont_prompt, 
            padding=True,
            truncation=True,
            max_length=1000,
            return_tensors="pt"
        ).to(self.device)
        cont_embed = self.text_encoder(cont_tokens)
        cont_embed_mask = cont_tokens.attention_mask.unsqueeze(1)

        with torch.no_grad():
            c_cont, cont_embed_mask = self.model.durator(cont_embed, cont_embed_mask)

        vocoder_upsample_factor = np.prod(self.vocoder.config.upsample_rates) / self.vocoder.config.sampling_rate
        height = int(audio_length_in_s * 1.024 / vocoder_upsample_factor)
        original_waveform_length = int(audio_length_in_s * self.vocoder.config.sampling_rate)

        # prepare timesteps
        self.noise_scheduler.set_timesteps(num_inference_steps)
        timesteps = self.noise_scheduler.timesteps

        # prepare latent variables
        num_channels_latents = self.model.unet.config.in_channels
        latents = self.prepare_latents(
            batch_size,
            num_channels_latents,
            height,
            c_desc.dtype,
            device=device,
            generator=torch.manual_seed(seed) if seed else None,
            latents=None,
        )

        # prepare extra step kwargs
        extra_step_kwargs = {
            'eta': 0.0,
            'generator': torch.manual_seed(seed) if seed else None,
        }

        # denoising loop
        with torch.no_grad():
            num_warmup_steps = len(timesteps) - num_inference_steps * self.noise_scheduler.order
            for i, t in enumerate(tqdm(timesteps)):
                if guidance == "dual":
                    latent_model_input = torch.cat([latents] * 4) if do_classifier_free_guidance else latents
                if guidance == "single":
                    latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
                
                latent_model_input = self.noise_scheduler.scale_model_input(latent_model_input, t)

                noise_pred = self.model(
                    latent_model_input, 
                    t, 
                    c_cont,
                    cont_embed_mask,
                    c_desc,
                )

                # perform guidance
                if do_classifier_free_guidance:
                    if guidance == "dual":
                        n1, n2, n3, n4 = noise_pred.chunk(4)
                        noise_pred = n1 + desc_guidance_scale * (n2 - n4) + cont_guidance_scale * (n3 - n4)
                    if guidance == "single":
                        noise_pred_text, noise_pred_uncond = noise_pred.chunk(2)
                        noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

                # compute the previous noise sample x_t -> x_t-1
                latents = self.noise_scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample

            mel_spectrogram = self.decode_latents(latents)
            audio = self.mel_spectrogram_to_waveform(mel_spectrogram)
            audio = audio[:, :original_waveform_length]
            audio = self.normalize_wav(audio)

        return audio
voiceldm/pipeline.py
glory20h/VoiceLDM