FL2VA/audio_vae/dac_audio_vae.py
7.1 KB · 226 lines · python Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # DAC-lineage audio VAE: waveform encoder + BigVGAN decoder (inference-only bundle).
3 import math
4 from typing import List
5
6 import numpy as np
7 import torch
8 from torch import nn
9 from torch.nn.utils.parametrizations import weight_norm
10
11 from .dac_bigvgan import BigVGAN
12 from .dac_attn_proj import AttnProjection
13
14
15 class AttrDict(dict):
16 def __init__(self, *args, **kwargs):
17 super(AttrDict, self).__init__(*args, **kwargs)
18 self.__dict__ = self
19
20
21 def WNConv1d(*args, **kwargs):
22 return weight_norm(nn.Conv1d(*args, **kwargs))
23
24
25 @torch.jit.script
26 def snake(x, alpha):
27 shape = x.shape
28 x = x.reshape(shape[0], shape[1], -1)
29 x = x + (alpha + 1e-9).reciprocal() * torch.sin(alpha * x).pow(2)
30 x = x.reshape(shape)
31 return x
32
33
34 class Snake1d(nn.Module):
35 def __init__(self, channels):
36 super().__init__()
37 self.alpha = nn.Parameter(torch.ones(1, channels, 1))
38
39 def forward(self, x):
40 return snake(x, self.alpha)
41
42
43 def init_weights(m):
44 if isinstance(m, nn.Conv1d):
45 nn.init.trunc_normal_(m.weight, std=0.02)
46 if m.bias is not None:
47 nn.init.constant_(m.bias, 0)
48
49
50 class ResidualUnit(nn.Module):
51 def __init__(self, dim: int = 16, dilation: int = 1):
52 super().__init__()
53 pad = ((7 - 1) * dilation) // 2
54 self.block = nn.Sequential(
55 Snake1d(dim),
56 WNConv1d(dim, dim, kernel_size=7, dilation=dilation, padding=pad),
57 Snake1d(dim),
58 WNConv1d(dim, dim, kernel_size=1),
59 )
60
61 def forward(self, x):
62 y = self.block(x)
63 pad = (x.shape[-1] - y.shape[-1]) // 2
64 if pad > 0:
65 x = x[..., pad:-pad]
66 return x + y
67
68
69 class EncoderBlock(nn.Module):
70 def __init__(self, dim: int = 16, stride: int = 1):
71 super().__init__()
72 self.block = nn.Sequential(
73 ResidualUnit(dim // 2, dilation=1),
74 ResidualUnit(dim // 2, dilation=3),
75 ResidualUnit(dim // 2, dilation=9),
76 Snake1d(dim // 2),
77 WNConv1d(
78 dim // 2,
79 dim,
80 kernel_size=2 * stride,
81 stride=stride,
82 padding=math.ceil(stride / 2),
83 ),
84 )
85
86 def forward(self, x):
87 return self.block(x)
88
89
90 class Encoder(nn.Module):
91 def __init__(
92 self,
93 d_model: int = 64,
94 strides: list = [2, 4, 8, 8],
95 d_latent: int = 64,
96 ):
97 super().__init__()
98 # Create first convolution
99 self.block = [WNConv1d(1, d_model, kernel_size=7, padding=3)]
100
101 # Create EncoderBlocks that double channels as they downsample by `stride`
102 for stride in strides:
103 d_model *= 2
104 self.block += [EncoderBlock(d_model, stride=stride)]
105
106 # Create last convolution
107 self.block += [
108 Snake1d(d_model),
109 WNConv1d(d_model, d_latent, kernel_size=3, padding=1),
110 ]
111
112 # Wrap black into nn.Sequential
113 self.block = nn.Sequential(*self.block)
114 self.enc_dim = d_model
115
116 def forward(self, x):
117 return self.block(x)
118
119
120 class DacAudioVAE(nn.Module):
121 def __init__(
122 self,
123 encoder_dim: int = 64,
124 encoder_rates: List[int] = [2, 4, 8, 8],
125 latent_dim: int = None,
126 decoder_dim: int = 1536,
127 decoder_rates: List[int] = [8, 8, 4, 2],
128 sample_rate: int = 44100,
129 vae_latent_channels: int = 64,
130 attn_proj: bool = False,
131 decoder_type: str = "bigvgan",
132 ):
133 super().__init__()
134
135 self.encoder_dim = encoder_dim
136 self.encoder_rates = encoder_rates
137 self.decoder_dim = decoder_dim
138 self.decoder_rates = decoder_rates
139 self.sample_rate = sample_rate
140 self.attn_proj = attn_proj
141 self.decoder_type = decoder_type
142
143 if latent_dim is None:
144 latent_dim = encoder_dim * (2 ** len(encoder_rates))
145
146 self.latent_dim = latent_dim
147
148 self.hop_length = np.prod(encoder_rates)
149 self.encoder = Encoder(encoder_dim, encoder_rates, latent_dim)
150
151 if latent_dim % vae_latent_channels == 0:
152 self.attn_proj_dim = vae_latent_channels
153 else:
154 # smallest power of two >= vae_latent_channels
155 self.attn_proj_dim = 2 ** int(np.ceil(np.log2(vae_latent_channels)))
156
157 self.mean_proj = nn.Conv1d(self.attn_proj_dim, vae_latent_channels, 1)
158 self.logs_proj = nn.Conv1d(self.attn_proj_dim, vae_latent_channels, 1)
159
160 self.dec_in_proj = nn.Conv1d(vae_latent_channels, latent_dim, 1)
161
162 if self.decoder_type == "bigvgan":
163 if sample_rate == 16000:
164 bigvgan_conf = {"resblock": "1",
165 "num_mels": latent_dim,
166 "upsample_rates": [5,5,2,2,2,2],
167 "upsample_kernel_sizes": [9,9,4,4,4,4],
168 "upsample_initial_channel": decoder_dim,
169 "resblock_kernel_sizes": [3,7,11],
170 "resblock_dilation_sizes": [[1,3,5], [1,3,5], [1,3,5]],
171 "use_tanh_at_final": False,
172 "use_bias_at_final": False,
173 "activation": "snakebeta",
174 "snake_logscale": True}
175 elif sample_rate == 32000:
176 bigvgan_conf = {"resblock": "1",
177 "num_mels": latent_dim,
178 "upsample_rates": [5,5,2,2,2,2,2],
179 "upsample_kernel_sizes": [9,9,4,4,4,4,4],
180 "upsample_initial_channel": decoder_dim,
181 "resblock_kernel_sizes": [3,7,11],
182 "resblock_dilation_sizes": [[1,3,5], [1,3,5], [1,3,5]],
183 "use_tanh_at_final": False,
184 "use_bias_at_final": False,
185 "activation": "snakebeta",
186 "snake_logscale": True}
187 else:
188 raise ValueError(f"Invalid sample_rate: {sample_rate}")
189
190 h = AttrDict(**bigvgan_conf)
191 self.decoder = BigVGAN(h)
192 else:
193 raise ValueError(f"Invalid decoder type: {self.decoder_type}")
194
195 if self.attn_proj:
196 self.pre_block = AttnProjection(latent_dim, self.attn_proj_dim, num_heads=8)
197
198 self.sample_rate = sample_rate
199 self.apply(init_weights)
200
201 def preprocess(self, audio_data, sample_rate):
202 if sample_rate is None:
203 sample_rate = self.sample_rate
204
205 length = audio_data.shape[-1]
206 right_pad = math.ceil(length / self.hop_length) * self.hop_length - length
207 audio_data = nn.functional.pad(audio_data, (0, right_pad))
208
209 return audio_data
210
211 def decode(self, z: torch.Tensor):
212 """Decode given latent codes and return audio data
213
214 Parameters
215 ----------
216 z : Tensor[B x D x T]
217 Continuous latent representation
218
219 Returns
220 -------
221 Tensor[B x 1 x length]
222 Decoded audio data.
223 """
224 z = self.dec_in_proj(z)
225 return self.decoder(z)
226