| """Decoder half of the YuE2 48 kHz stereo waveform VAE.""" |
|
|
| import math |
|
|
| import torch |
| from torch import nn |
|
|
| SAMPLE_RATE = 48000 |
| HOP = 1920 |
|
|
|
|
| def WNConv1d(*args, **kwargs): |
| return nn.utils.parametrizations.weight_norm(nn.Conv1d(*args, **kwargs)) |
|
|
|
|
| def WNConvT1d(*args, **kwargs): |
| return nn.utils.parametrizations.weight_norm(nn.ConvTranspose1d(*args, **kwargs)) |
|
|
|
|
| class SnakeBeta(nn.Module): |
| def __init__(self, channels: int): |
| super().__init__() |
| self.alpha = nn.Parameter(torch.zeros(channels)) |
| self.beta = nn.Parameter(torch.zeros(channels)) |
|
|
| def forward(self, x): |
| alpha = self.alpha.unsqueeze(0).unsqueeze(-1).exp() |
| beta = self.beta.unsqueeze(0).unsqueeze(-1).exp() |
| return x + (1.0 / (beta + 1e-9)) * torch.sin(x * alpha).pow(2) |
|
|
|
|
| class ResUnit(nn.Module): |
| def __init__(self, channels: int, dilation: int): |
| super().__init__() |
| self.layers = nn.Sequential( |
| SnakeBeta(channels), |
| WNConv1d( |
| channels, |
| channels, |
| 7, |
| dilation=dilation, |
| padding=(dilation * 6) // 2, |
| ), |
| SnakeBeta(channels), |
| WNConv1d(channels, channels, 1), |
| ) |
|
|
| def forward(self, x): |
| return x + self.layers(x) |
|
|
|
|
| class DecBlock(nn.Module): |
| def __init__(self, in_channels: int, out_channels: int, stride: int): |
| super().__init__() |
| self.layers = nn.Sequential( |
| SnakeBeta(in_channels), |
| WNConvT1d( |
| in_channels, |
| out_channels, |
| 2 * stride, |
| stride=stride, |
| padding=math.ceil(stride / 2), |
| output_padding=stride % 2, |
| ), |
| ResUnit(out_channels, 1), |
| ResUnit(out_channels, 3), |
| ResUnit(out_channels, 9), |
| ) |
|
|
| def forward(self, x): |
| return self.layers(x) |
|
|
|
|
| class SequenceWrapper(nn.Module): |
| def __init__(self, *modules): |
| super().__init__() |
| self.layers = nn.Sequential(*modules) |
|
|
| def forward(self, x): |
| return self.layers(x) |
|
|
|
|
| def convert_weight_norm_keys(state_dict: dict) -> dict: |
| converted = {} |
| for key, value in state_dict.items(): |
| if key.endswith(".weight_g"): |
| key = key[: -len(".weight_g")] + ".parametrizations.weight.original0" |
| elif key.endswith(".weight_v"): |
| key = key[: -len(".weight_v")] + ".parametrizations.weight.original1" |
| converted[key] = value |
| return converted |
|
|
|
|
| class YuE2VAEDecoder(nn.Module): |
| def __init__( |
| self, |
| in_channels: int = 2, |
| channels: int = 64, |
| latent_dim: int = 64, |
| channel_multipliers=(1, 2, 4, 8, 16, 32), |
| strides=(2, 2, 4, 4, 5, 6), |
| ): |
| super().__init__() |
| multipliers = [1] + list(channel_multipliers) |
| decoder = [WNConv1d(latent_dim, multipliers[-1] * channels, 7, padding=3)] |
| for index in range(len(multipliers) - 1, 0, -1): |
| decoder.append( |
| DecBlock( |
| multipliers[index] * channels, |
| multipliers[index - 1] * channels, |
| strides[index - 1], |
| ) |
| ) |
| decoder += [ |
| SnakeBeta(multipliers[0] * channels), |
| WNConv1d(multipliers[0] * channels, in_channels, 7, padding=3, bias=False), |
| ] |
| self.decoder = SequenceWrapper(*decoder) |
|
|
| @classmethod |
| def load_from_state_dict(cls, state_dict: dict, dtype=torch.float32): |
| model = cls() |
| state_dict = convert_weight_norm_keys(state_dict) |
| model.load_state_dict({key: value.to(dtype) for key, value in state_dict.items()}, strict=True) |
| return model.to(dtype) |
|
|
| @property |
| def device(self): |
| return next(self.parameters()).device |
|
|
| @property |
| def dtype(self): |
| return next(self.parameters()).dtype |
|
|
| def decode(self, latents): |
| return self.decoder(latents) |
|
|
| def tiled_decode(self, latents, core_frames: int = 750, halo_frames: int = 16): |
| total = latents.shape[-1] |
| if total <= core_frames + 2 * halo_frames: |
| return self.decoder(latents) |
| pieces = [] |
| for start in range(0, total, core_frames): |
| end = min(start + core_frames, total) |
| low = max(0, start - halo_frames) |
| high = min(total, end + halo_frames) |
| audio = self.decoder(latents[..., low:high]) |
| pieces.append( |
| audio[..., (start - low) * HOP : (start - low + end - start) * HOP] |
| ) |
| return torch.cat(pieces, dim=-1) |
|
|
