upload files from https://github.com/facebookresearch/speech-resynthesis

models.py

@@ -0,0 +1,38 @@
+# adapted from https://github.com/jik876/hifi-gan
+
+from transformers.modeling_utils import PreTrainedModel
+
+from quantizer_config import QuantizerConfig
+from modules.jukebox import Encoder, Decoder
+from modules.vq import Bottleneck
+
+
+
+class Quantizer(PreTrainedModel):
+    config_class = QuantizerConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.config = config
+        self.encoder = Encoder(**config.f0_encoder_params)
+        self.vq = Bottleneck(**config.f0_vq_params)
+        self.decoder = Decoder(**config.f0_decoder_params)
+
+    def forward(self, **kwargs):
+        f0_h = self.encoder(kwargs['features'])
+        
+        zs, f0_h_q, f0_commit_losses, f0_metrics = self.vq(f0_h)
+        
+        f0 = self.decoder(f0_h_q)
+
+        return {
+            'f0': f0,                          
+            'commit_losses': f0_commit_losses, 
+            'metrics': f0_metrics,             
+            'codes': zs,                       
+            'hidden_states': f0_h_q            
+        }
+
+
+

modules/dist.py

@@ -0,0 +1,108 @@
+# Adapted from https://github.com/openai/jukebox
+
+from enum import Enum
+
+import torch.distributed as dist
+
+
+class ReduceOp(Enum):
+    SUM = 0,
+    PRODUCT = 1,
+    MIN = 2,
+    MAX = 3
+
+    def ToDistOp(self):
+        return {
+            self.SUM: dist.ReduceOp.SUM,
+            self.PRODUCT: dist.ReduceOp.PRODUCT,
+            self.MIN: dist.ReduceOp.MIN,
+            self.MAX: dist.ReduceOp.MAX
+        }[self]
+
+
+def is_available():
+    return dist.is_initialized()
+
+
+def get_rank():
+    if is_available():
+        return _get_rank()
+    else:
+        return 0
+
+
+def get_world_size():
+    if is_available():
+        return _get_world_size()
+    else:
+        return 1
+
+
+def barrier():
+    if is_available():
+        return _barrier()
+    # else: do nothing
+
+
+def all_gather(tensor_list, tensor):
+    if is_available():
+        return _all_gather(tensor_list, tensor)
+    else:
+        tensor_list[0] = tensor
+
+
+def all_reduce(tensor, op=ReduceOp.SUM):
+    if is_available():
+        return _all_reduce(tensor, op)
+    # else: do nothing
+
+
+def reduce(tensor, dst, op=ReduceOp.SUM):
+    if is_available():
+        return _reduce(tensor, dst, op)
+    # else: do nothing
+
+
+def broadcast(tensor, src):
+    if is_available():
+        return _broadcast(tensor, src)
+    # else: do nothing
+
+
+def init_process_group(backend, init_method):
+    if is_available():
+        return _init_process_group(backend, init_method)
+    # else: do nothing
+
+
+def _get_rank():
+    return dist.get_rank()
+
+
+def _barrier():
+    return dist.barrier()
+
+
+def _get_world_size():
+    return dist.get_world_size()
+
+
+def _all_gather(tensor_list, tensor):
+    return dist.all_gather(tensor_list, tensor)
+
+
+def _all_reduce(tensor, op):
+    return dist.all_reduce(tensor, op.ToDistOp())
+
+
+def _reduce(tensor, dst, op):
+    return dist.reduce(tensor, dst, op.ToDistOp())
+
+
+def _broadcast(tensor, src):
+    return dist.broadcast(tensor, src)
+
+
+def _init_process_group(backend, init_method):
+    return dist.init_process_group(backend, init_method)
+

modules/jukebox.py

@@ -0,0 +1,178 @@
+# Adapted from https://github.com/openai/jukebox
+
+import numpy as np
+import torch.nn as nn
+from modules.resnet import Resnet1D
+
+
+def assert_shape(x, exp_shape):
+    assert x.shape == exp_shape, f"Expected {exp_shape} got {x.shape}"
+
+
+class EncoderConvBlock(nn.Module):
+    def __init__(self, input_emb_width, output_emb_width, down_t, stride_t, width, depth, m_conv,
+                 dilation_growth_rate=1, dilation_cycle=None, zero_out=False, res_scale=False):
+        super().__init__()
+        blocks = []
+        if type(stride_t) is tuple or type(stride_t) is list:
+            start = True
+            for s_t, d_t in zip(stride_t, down_t):
+                if s_t % 2 == 0:
+                    filter_t, pad_t = s_t * 2, s_t // 2
+                else:
+                    filter_t, pad_t = s_t * 2 + 1, s_t // 2 + 1
+                if d_t > 0:
+                    for i in range(d_t):
+                        block = nn.Sequential(
+                            nn.Conv1d(input_emb_width if i == 0 and start else width, width, filter_t, s_t, pad_t),
+                            Resnet1D(width, depth, m_conv, dilation_growth_rate, dilation_cycle, zero_out, res_scale), )
+                        blocks.append(block)
+                        start = False
+            block = nn.Conv1d(width, output_emb_width, 3, 1, 1)
+            blocks.append(block)
+        else:
+            filter_t, pad_t = stride_t * 2, stride_t // 2
+            if down_t > 0:
+                for i in range(down_t):
+                    block = nn.Sequential(
+                        nn.Conv1d(input_emb_width if i == 0 else width, width, filter_t, stride_t, pad_t),
+                        Resnet1D(width, depth, m_conv, dilation_growth_rate, dilation_cycle, zero_out, res_scale), )
+                    blocks.append(block)
+                block = nn.Conv1d(width, output_emb_width, 3, 1, 1)
+                blocks.append(block)
+        self.model = nn.Sequential(*blocks)
+
+    def forward(self, x):
+        return self.model(x)
+
+
+class DecoderConvBock(nn.Module):
+    def __init__(self, input_emb_width, output_emb_width, down_t, stride_t, width, depth, m_conv,
+                 dilation_growth_rate=1, dilation_cycle=None, zero_out=False, res_scale=False,
+                 reverse_decoder_dilation=False, checkpoint_res=False):
+        super().__init__()
+        blocks = []
+
+        if type(stride_t) is tuple or type(stride_t) is list:
+            block = nn.Conv1d(output_emb_width, width, 3, 1, 1)
+            blocks.append(block)
+            for k, (s_t, d_t) in enumerate(zip(stride_t, down_t)):
+                if d_t > 0:
+                    if s_t % 2 == 0:
+                        filter_t, pad_t = s_t * 2, s_t // 2
+                    else:
+                        filter_t, pad_t = s_t * 2 + 1, s_t // 2 + 1
+                    end = k == len(stride_t) - 1
+                    for i in range(d_t):
+                        block = nn.Sequential(
+                            Resnet1D(width, depth, m_conv, dilation_growth_rate, dilation_cycle, zero_out=zero_out,
+                                     res_scale=res_scale, reverse_dilation=reverse_decoder_dilation,
+                                     checkpoint_res=checkpoint_res),
+                            nn.ConvTranspose1d(width, input_emb_width if i == (d_t - 1) and end else width, filter_t,
+                                               s_t, pad_t))
+                        blocks.append(block)
+        else:
+            if down_t > 0:
+                filter_t, pad_t = stride_t * 2, stride_t // 2
+                block = nn.Conv1d(output_emb_width, width, 3, 1, 1)
+                blocks.append(block)
+                for i in range(down_t):
+                    block = nn.Sequential(
+                        Resnet1D(width, depth, m_conv, dilation_growth_rate, dilation_cycle, zero_out=zero_out,
+                                 res_scale=res_scale, reverse_dilation=reverse_decoder_dilation,
+                                 checkpoint_res=checkpoint_res),
+                        nn.ConvTranspose1d(width, input_emb_width if i == (down_t - 1) else width, filter_t, stride_t,
+                                           pad_t))
+                    blocks.append(block)
+        self.model = nn.Sequential(*blocks)
+
+    def forward(self, x):
+        return self.model(x)
+
+
+class Encoder(nn.Module):
+    def __init__(self, input_emb_width, output_emb_width, levels, downs_t, strides_t, **block_kwargs):
+        super().__init__()
+        self.input_emb_width = input_emb_width
+        self.output_emb_width = output_emb_width
+        self.levels = levels
+        self.downs_t = downs_t
+        self.strides_t = strides_t
+
+        block_kwargs_copy = dict(**block_kwargs)
+        if 'reverse_decoder_dilation' in block_kwargs_copy:
+            del block_kwargs_copy['reverse_decoder_dilation']
+        level_block = lambda level, down_t, stride_t: EncoderConvBlock(
+            input_emb_width if level == 0 else output_emb_width, output_emb_width, down_t, stride_t,
+            **block_kwargs_copy)
+        self.level_blocks = nn.ModuleList()
+        iterator = zip(list(range(self.levels)), downs_t, strides_t)
+        for level, down_t, stride_t in iterator:
+            self.level_blocks.append(level_block(level, down_t, stride_t))
+
+    def forward(self, x):
+        N, T = x.shape[0], x.shape[-1]
+        emb = self.input_emb_width
+        assert_shape(x, (N, emb, T))
+        xs = []
+
+        # 64, 32, ...
+        iterator = zip(list(range(self.levels)), self.downs_t, self.strides_t)
+        for level, down_t, stride_t in iterator:
+            level_block = self.level_blocks[level]
+            x = level_block(x)
+            if type(stride_t) is tuple or type(stride_t) is list:
+                emb, T = self.output_emb_width, T // np.prod([s ** d for s, d in zip(stride_t, down_t)])
+            else:
+                emb, T = self.output_emb_width, T // (stride_t ** down_t)
+            assert_shape(x, (N, emb, T))
+            xs.append(x)
+
+        return xs
+
+
+class Decoder(nn.Module):
+    def __init__(self, input_emb_width, output_emb_width, levels, downs_t, strides_t, **block_kwargs):
+        super().__init__()
+        self.input_emb_width = input_emb_width
+        self.output_emb_width = output_emb_width
+        self.levels = levels
+
+        self.downs_t = downs_t
+
+        self.strides_t = strides_t
+
+        level_block = lambda level, down_t, stride_t: DecoderConvBock(output_emb_width, output_emb_width, down_t,
+                                                                      stride_t, **block_kwargs)
+        self.level_blocks = nn.ModuleList()
+        iterator = zip(list(range(self.levels)), downs_t, strides_t)
+        for level, down_t, stride_t in iterator:
+            self.level_blocks.append(level_block(level, down_t, stride_t))
+
+        self.out = nn.Conv1d(output_emb_width, input_emb_width, 3, 1, 1)
+
+    def forward(self, xs, all_levels=True):
+        if all_levels:
+            assert len(xs) == self.levels
+        else:
+            assert len(xs) == 1
+        x = xs[-1]
+        N, T = x.shape[0], x.shape[-1]
+        emb = self.output_emb_width
+        assert_shape(x, (N, emb, T))
+
+        # 32, 64 ...
+        iterator = reversed(list(zip(list(range(self.levels)), self.downs_t, self.strides_t)))
+        for level, down_t, stride_t in iterator:
+            level_block = self.level_blocks[level]
+            x = level_block(x)
+            if type(stride_t) is tuple or type(stride_t) is list:
+                emb, T = self.output_emb_width, T * np.prod([s ** d for s, d in zip(stride_t, down_t)])
+            else:
+                emb, T = self.output_emb_width, T * (stride_t ** down_t)
+            assert_shape(x, (N, emb, T))
+            if level != 0 and all_levels:
+                x = x + xs[level - 1]
+
+        x = self.out(x)
+        return x

modules/resnet.py

@@ -0,0 +1,82 @@
+# Adapted from https://github.com/openai/jukebox
+
+import math
+import torch.nn as nn
+
+import modules.dist as dist
+
+
+class ResConvBlock(nn.Module):
+    def __init__(self, n_in, n_state):
+        super().__init__()
+        self.model = nn.Sequential(
+            nn.ReLU(),
+            nn.Conv2d(n_in, n_state, 3, 1, 1),
+            nn.ReLU(),
+            nn.Conv2d(n_state, n_in, 1, 1, 0),
+        )
+
+    def forward(self, x):
+        return x + self.model(x)
+
+
+class Resnet(nn.Module):
+    def __init__(self, n_in, n_depth, m_conv=1.0):
+        super().__init__()
+        self.model = nn.Sequential(*[ResConvBlock(n_in, int(m_conv * n_in)) for _ in range(n_depth)])
+
+    def forward(self, x):
+        return self.model(x)
+
+
+class ResConv1DBlock(nn.Module):
+    def __init__(self, n_in, n_state, dilation=1, zero_out=False, res_scale=1.0):
+        super().__init__()
+        padding = dilation
+        self.model = nn.Sequential(
+            nn.ReLU(),
+            nn.Conv1d(n_in, n_state, 3, 1, padding, dilation),
+            nn.ReLU(),
+            nn.Conv1d(n_state, n_in, 1, 1, 0),
+        )
+        if zero_out:
+            out = self.model[-1]
+            nn.init.zeros_(out.weight)
+            nn.init.zeros_(out.bias)
+        self.res_scale = res_scale
+
+    def forward(self, x):
+        return x + self.res_scale * self.model(x)
+
+
+class Resnet1D(nn.Module):
+    def __init__(self, n_in, n_depth, m_conv=1.0, dilation_growth_rate=1, dilation_cycle=None, zero_out=False,
+                 res_scale=False, reverse_dilation=False, checkpoint_res=False):
+        super().__init__()
+
+        def _get_depth(depth):
+            if dilation_cycle is None:
+                return depth
+            else:
+                return depth % dilation_cycle
+
+        blocks = [ResConv1DBlock(n_in, int(m_conv * n_in),
+                                 dilation=dilation_growth_rate ** _get_depth(depth),
+                                 zero_out=zero_out,
+                                 res_scale=1.0 if not res_scale else 1.0 / math.sqrt(n_depth))
+                  for depth in range(n_depth)]
+        if reverse_dilation:
+            blocks = blocks[::-1]
+        self.checkpoint_res = checkpoint_res
+        if self.checkpoint_res == 1:
+            if dist.get_rank() == 0:
+                print("Checkpointing convs")
+            self.blocks = nn.ModuleList(blocks)
+        else:
+            self.model = nn.Sequential(*blocks)
+
+    def forward(self, x):
+        if self.checkpoint_res == 1:
+            raise NotImplementedError("Checkpoint not implemented")
+        else:
+            return self.model(x)

modules/vq.py

@@ -0,0 +1,249 @@
+# Adapted from https://github.com/openai/jukebox
+
+import numpy as np
+import torch as t
+import torch.nn as nn
+import torch.nn.functional as F
+
+import modules.dist as dist
+
+
+class BottleneckBlock(nn.Module):
+    def __init__(self, k_bins, emb_width, mu):
+        super().__init__()
+        self.k_bins = k_bins
+        self.emb_width = emb_width
+        self.mu = mu
+        self.reset_k()
+        self.threshold = 1.0
+
+    def reset_k(self):
+        self.init = False
+        self.k_sum = None
+        self.k_elem = None
+        self.register_buffer('k', t.zeros(self.k_bins, self.emb_width).cuda())
+
+    def _tile(self, x):
+        d, ew = x.shape
+        if d < self.k_bins:
+            n_repeats = (self.k_bins + d - 1) // d
+            std = 0.01 / np.sqrt(ew)
+            x = x.repeat(n_repeats, 1)
+            x = x + t.randn_like(x) * std
+        return x
+
+    def init_k(self, x):
+        mu, emb_width, k_bins = self.mu, self.emb_width, self.k_bins
+        self.init = True
+        # init k_w using random vectors from x
+        y = self._tile(x)
+        _k_rand = y[t.randperm(y.shape[0])][:k_bins]
+        dist.broadcast(_k_rand, 0)
+        self.k = _k_rand
+        assert self.k.shape == (k_bins, emb_width)
+        self.k_sum = self.k
+        self.k_elem = t.ones(k_bins, device=self.k.device)
+
+    def restore_k(self, num_tokens=None, threshold=1.0):
+        mu, emb_width, k_bins = self.mu, self.emb_width, self.k_bins
+        self.init = True
+        assert self.k.shape == (k_bins, emb_width)
+        self.k_sum = self.k.clone()
+        self.k_elem = t.ones(k_bins, device=self.k.device)
+        if num_tokens is not None:
+            expected_usage = num_tokens / k_bins
+            self.k_elem.data.mul_(expected_usage)
+            self.k_sum.data.mul_(expected_usage)
+        self.threshold = threshold
+
+    def update_k(self, x, x_l):
+        mu, emb_width, k_bins = self.mu, self.emb_width, self.k_bins
+        with t.no_grad():
+            # Calculate new centres
+            x_l_onehot = t.zeros(k_bins, x.shape[0], device=x.device)  # k_bins, N * L
+            x_l_onehot.scatter_(0, x_l.view(1, x.shape[0]), 1)
+
+            _k_sum = t.matmul(x_l_onehot, x)  # k_bins, w
+            _k_elem = x_l_onehot.sum(dim=-1)  # k_bins
+            y = self._tile(x)
+            _k_rand = y[t.randperm(y.shape[0])][:k_bins]
+
+            dist.broadcast(_k_rand, 0)
+            dist.all_reduce(_k_sum)
+            dist.all_reduce(_k_elem)
+
+            # Update centres
+            old_k = self.k
+            self.k_sum = mu * self.k_sum + (1. - mu) * _k_sum  # w, k_bins
+            self.k_elem = mu * self.k_elem + (1. - mu) * _k_elem  # k_bins
+            usage = (self.k_elem.view(k_bins, 1) >= self.threshold).float()
+            self.k = usage * (self.k_sum.view(k_bins, emb_width) / self.k_elem.view(k_bins, 1)) \
+                     + (1 - usage) * _k_rand
+            _k_prob = _k_elem / t.sum(_k_elem)  # x_l_onehot.mean(dim=-1)  # prob of each bin
+            entropy = -t.sum(_k_prob * t.log(_k_prob + 1e-8))  # entropy ie how diverse
+            used_curr = (_k_elem >= self.threshold).sum()
+            usage = t.sum(usage)
+            dk = t.norm(self.k - old_k) / np.sqrt(np.prod(old_k.shape))
+        return dict(entropy=entropy,
+                    used_curr=used_curr,
+                    usage=usage,
+                    dk=dk)
+
+    def preprocess(self, x):
+        # NCT -> NTC -> [NT, C]
+        x = x.permute(0, 2, 1).contiguous()
+        x = x.view(-1, x.shape[-1])  # x_en = (N * L, w), k_j = (w, k_bins)
+
+        if x.shape[-1] == self.emb_width:
+            prenorm = t.norm(x - t.mean(x)) / np.sqrt(np.prod(x.shape))
+        elif x.shape[-1] == 2 * self.emb_width:
+            x1, x2 = x[..., :self.emb_width], x[..., self.emb_width:]
+            prenorm = (t.norm(x1 - t.mean(x1)) / np.sqrt(np.prod(x1.shape))) + (
+                        t.norm(x2 - t.mean(x2)) / np.sqrt(np.prod(x2.shape)))
+
+            # Normalise
+            x = x1 + x2
+        else:
+            assert False, f"Expected {x.shape[-1]} to be (1 or 2) * {self.emb_width}"
+        return x, prenorm
+
+    def postprocess(self, x_l, x_d, x_shape):
+        # [NT, C] -> NTC -> NCT
+        N, T = x_shape
+        x_d = x_d.view(N, T, -1).permute(0, 2, 1).contiguous()
+        x_l = x_l.view(N, T)
+        return x_l, x_d
+
+    def quantise(self, x):
+        # Calculate latent code x_l
+        k_w = self.k.t()
+        distance = t.sum(x ** 2, dim=-1, keepdim=True) - 2 * t.matmul(x, k_w) + t.sum(k_w ** 2, dim=0,
+                                                                                      keepdim=True)  # (N * L, b)
+        min_distance, x_l = t.min(distance, dim=-1)
+        fit = t.mean(min_distance)
+        return x_l, fit
+
+    def dequantise(self, x_l):
+        x = F.embedding(x_l, self.k)
+        return x
+
+    def encode(self, x):
+        N, width, T = x.shape
+
+        # Preprocess.
+        x, prenorm = self.preprocess(x)
+
+        # Quantise
+        x_l, fit = self.quantise(x)
+
+        # Postprocess.
+        x_l = x_l.view(N, T)
+        return x_l
+
+    def decode(self, x_l):
+        N, T = x_l.shape
+        width = self.emb_width
+
+        # Dequantise
+        x_d = self.dequantise(x_l)
+
+        # Postprocess
+        x_d = x_d.view(N, T, width).permute(0, 2, 1).contiguous()
+        return x_d
+
+    def forward(self, x, update_k=True):
+        N, width, T = x.shape
+
+        # Preprocess
+        x, prenorm = self.preprocess(x)
+
+        # Init k if not inited
+        if update_k and not self.init:
+            self.init_k(x)
+
+        # Quantise and dequantise through bottleneck
+        x_l, fit = self.quantise(x)
+        x_d = self.dequantise(x_l)
+
+        # Update embeddings
+        if update_k and self.training:
+            update_metrics = self.update_k(x, x_l)
+        else:
+            update_metrics = {}
+
+        # Loss
+        commit_loss = t.norm(x_d.detach() - x) ** 2 / np.prod(x.shape)
+
+        # Passthrough
+        x_d = x + (x_d - x).detach()
+
+        # Postprocess
+        x_l, x_d = self.postprocess(x_l, x_d, (N, T))
+        return x_l, x_d, commit_loss, dict(fit=fit,
+                                           pn=prenorm,
+                                           **update_metrics)
+
+
+class Bottleneck(nn.Module):
+    def __init__(self, l_bins, emb_width, mu, levels):
+        super().__init__()
+        self.levels = levels
+        level_block = lambda level: BottleneckBlock(l_bins, emb_width, mu)
+        self.level_blocks = nn.ModuleList()
+        for level in range(self.levels):
+            self.level_blocks.append(level_block(level))
+
+    def encode(self, xs):
+        zs = [level_block.encode(x) for (level_block, x) in zip(self.level_blocks, xs)]
+        return zs
+
+    def decode(self, zs, start_level=0, end_level=None):
+        if end_level is None:
+            end_level = self.levels
+        xs_quantised = [level_block.decode(z) for (level_block, z) in zip(self.level_blocks[start_level:end_level], zs)]
+        return xs_quantised
+
+    def forward(self, xs):
+        zs, xs_quantised, commit_losses, metrics = [], [], [], []
+        for level in range(self.levels):
+            level_block = self.level_blocks[level]
+            x = xs[level]
+            z, x_quantised, commit_loss, metric = level_block(x, update_k=self.training)
+            zs.append(z)
+            if not self.training:
+                # Be extra paranoid and make sure the encoder weights can't
+                # change from straight-through estimator
+                x_quantised = x_quantised.detach()
+            xs_quantised.append(x_quantised)
+            commit_losses.append(commit_loss)
+            if self.training:
+                metrics.append(metric)
+        return zs, xs_quantised, commit_losses, metrics
+
+
+class NoBottleneckBlock(nn.Module):
+    def restore_k(self):
+        pass
+
+
+class NoBottleneck(nn.Module):
+    def __init__(self, levels):
+        super().__init__()
+        self.level_blocks = nn.ModuleList()
+        self.levels = levels
+        for level in range(levels):
+            self.level_blocks.append(NoBottleneckBlock())
+
+    def encode(self, xs):
+        return xs
+
+    def decode(self, zs, start_level=0, end_level=None):
+        if end_level is None:
+            end_level = self.levels
+        return zs
+
+    def forward(self, xs):
+        zero = t.zeros(()).cuda()
+        commit_losses = [zero for _ in range(self.levels)]
+        metrics = [dict(entropy=zero, usage=zero, used_curr=zero, pn=zero, dk=zero) for _ in range(self.levels)]
+        return xs, xs, commit_losses, metrics

pipeline_utils.py

@@ -0,0 +1,120 @@
+
+import matplotlib.pyplot as plt
+from typing import List
+import numpy as np
+from dataclasses import dataclass
+
+@dataclass
+class SpeakerStats:
+    f0_mean: float
+    f0_std: float
+    intensity_mean: float
+    intensity_std: float
+    
+    @classmethod
+    def from_features(cls, f0_values: List[np.ndarray], intensity_values: List[np.ndarray]):
+
+        f0_arrays = [np.array(f0) for f0 in f0_values]
+        intensity_arrays = [np.array(i) for i in intensity_values]
+        
+        f0_concat = np.concatenate([f0[f0 != 0] for f0 in f0_arrays])
+        intensity_concat = np.concatenate(intensity_arrays)
+        
+        
+        return cls(
+            f0_mean=float(np.mean(f0_concat)),
+            f0_std=float(np.std(f0_concat)),
+            intensity_mean=float(np.mean(intensity_concat)),
+            intensity_std=float(np.std(intensity_concat))
+        )
+
+def compute_speaker_stats(dataset, speaker_column='speaker_id'):
+    """
+    Calculate speaker statistics from a preprocessed dataset.
+    
+    Args:
+        dataset: HuggingFace dataset containing f0 and intensity features
+        speaker_column: Name of the speaker ID column (default: 'speaker')
+    
+    Returns:
+        Dict[str, SpeakerStats]: Dictionary mapping speaker IDs to their statistics
+    """
+    speaker_features = {}
+    
+    # Group features by speaker
+    for item in dataset:
+        speaker_id = item[speaker_column]
+        if speaker_id not in speaker_features:
+            speaker_features[speaker_id] = {'f0': [], 'intensity': []}
+        
+        speaker_features[speaker_id]['f0'].append(item['f0'])
+        speaker_features[speaker_id]['intensity'].append(item['intensity'])
+    
+    # Calculate stats per speaker
+    speaker_stats = {
+        spk: SpeakerStats.from_features(
+            feats['f0'],
+            feats['intensity']
+        )
+        for spk, feats in speaker_features.items()
+    }
+    
+    return speaker_stats
+
+def plot_reconstruction(result, sample_idx):
+    # Get F0 data
+    input_f0 = result['input_features']['f0_orig']
+    output_f0 = np.array(result['f0_recon'])
+
+    length = len(input_f0)
+    truncated_length = (length // 16) * 16
+
+    input_f0 = np.array(input_f0[:truncated_length])
+    
+    # Get intensity data
+    input_intensity = np.array(result['input_features']['intensity_orig'][:truncated_length])
+    output_intensity = np.array(result['intensity_recon'])
+    
+    time = np.arange(len(input_f0))
+    
+    # Create figure with two subplots
+    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(15, 10))
+    
+    # Plot F0
+    ax1.plot(time, input_f0, label='Original F0', alpha=0.7)
+    ax1.plot(time, output_f0, label='Reconstructed F0', alpha=0.7)
+    
+    # Highlight large differences in F0 (>20% of original)
+    f0_diff_percent = np.abs(input_f0 - output_f0) / (input_f0 + 1e-8) * 100  # Add small epsilon to avoid division by zero
+    large_diff_mask = (f0_diff_percent > 20)
+    if np.any(large_diff_mask):
+        ax1.fill_between(time, input_f0, output_f0, 
+                       where=large_diff_mask,
+                       color='red', alpha=0.3,
+                       label='Diff > 20%')
+    
+    ax1.set_title(f'F0 Reconstruction (Sample {sample_idx})')
+    ax1.set_ylabel('Frequency (Hz)')
+    ax1.legend()
+    
+    # Plot Intensity
+    ax2.plot(time, input_intensity, label='Original Intensity', alpha=0.7)
+    ax2.plot(time, output_intensity, label='Reconstructed Intensity', alpha=0.7)
+    
+    # Highlight large differences in intensity (>20% of original)
+    intensity_diff_percent = np.abs(input_intensity - output_intensity) / (np.abs(input_intensity) + 1e-8) * 100
+    intensity_large_diff = intensity_diff_percent > 20
+    if np.any(intensity_large_diff):
+        ax2.fill_between(time, input_intensity, output_intensity,
+                       where=intensity_large_diff,
+                       color='red', alpha=0.3,
+                       label='Diff > 20%')
+    
+    ax2.set_title('Intensity Reconstruction')
+    ax2.set_ylabel('Intensity (dB)')
+    ax2.set_xlabel('Time (frames)')
+    ax2.legend()
+    
+    
+    plt.tight_layout()
+    return fig

quantizer_config.py

@@ -0,0 +1,167 @@
+
+from transformers import PretrainedConfig
+from typing import List, Optional
+
+class QuantizerConfig(PretrainedConfig):
+    model_type = "prosody_quantizer"
+    
+    def __init__(
+        self,
+        # VQ parameters
+        l_bins: int = 320,
+        emb_width: int = 64,
+        mu: float = 0.99,
+        levels: int = 1,
+        
+        # Encoder parameters
+        encoder_input_emb_width: int = 3,
+        encoder_output_emb_width: int = 64,
+        encoder_levels: int = 1,
+        encoder_downs_t: List[int] = [4],
+        encoder_strides_t: List[int] = [2],
+        encoder_width: int = 32,
+        encoder_depth: int = 4,
+        encoder_m_conv: float = 1.0,
+        encoder_dilation_growth_rate: int = 3,
+        
+        # Decoder parameters
+        decoder_input_emb_width: int = 3,
+        decoder_output_emb_width: int = 64,
+        decoder_levels: int = 1,
+        decoder_downs_t: List[int] = [4],
+        decoder_strides_t: List[int] = [2],
+        decoder_width: int = 32,
+        decoder_depth: int = 4,
+        decoder_m_conv: float = 1.0,
+        decoder_dilation_growth_rate: int = 3,
+        
+        # Training parameters
+        lambda_commit: float = 0.02,
+        f0_normalize: bool = True,
+        intensity_normalize: bool = True,
+        multispkr: str = "single",
+        f0_feats: bool = False,
+        f0_median: bool = False,
+        
+        # Optional training hyperparameters
+        learning_rate: float = 0.0002,
+        adam_b1: float = 0.8,
+        adam_b2: float = 0.99,
+        lr_decay: float = 0.999,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        
+        # VQ parameters
+        self.l_bins = l_bins
+        self.emb_width = emb_width
+        self.mu = mu
+        self.levels = levels
+        
+        # Encoder parameters
+        self.encoder_input_emb_width = encoder_input_emb_width
+        self.encoder_output_emb_width = encoder_output_emb_width
+        self.encoder_levels = encoder_levels
+        self.encoder_downs_t = encoder_downs_t
+        self.encoder_strides_t = encoder_strides_t
+        self.encoder_width = encoder_width
+        self.encoder_depth = encoder_depth
+        self.encoder_m_conv = encoder_m_conv
+        self.encoder_dilation_growth_rate = encoder_dilation_growth_rate
+        
+        # Decoder parameters
+        self.decoder_input_emb_width = decoder_input_emb_width
+        self.decoder_output_emb_width = decoder_output_emb_width
+        self.decoder_levels = decoder_levels
+        self.decoder_downs_t = decoder_downs_t
+        self.decoder_strides_t = decoder_strides_t
+        self.decoder_width = decoder_width
+        self.decoder_depth = decoder_depth
+        self.decoder_m_conv = decoder_m_conv
+        self.decoder_dilation_growth_rate = decoder_dilation_growth_rate
+        
+        # Training parameters
+        self.lambda_commit = lambda_commit
+        self.f0_normalize = f0_normalize
+        self.intensity_normalize = intensity_normalize
+        self.multispkr = multispkr
+        self.f0_feats = f0_feats
+        self.f0_median = f0_median
+        
+        # Training hyperparameters
+        self.learning_rate = learning_rate
+        self.adam_b1 = adam_b1
+        self.adam_b2 = adam_b2
+        self.lr_decay = lr_decay
+    
+    @property
+    def f0_vq_params(self):
+        return {
+            "l_bins": self.l_bins,
+            "emb_width": self.emb_width,
+            "mu": self.mu,
+            "levels": self.levels
+        }
+    
+    @property
+    def f0_encoder_params(self):
+        return {
+            "input_emb_width": self.encoder_input_emb_width,
+            "output_emb_width": self.encoder_output_emb_width,
+            "levels": self.encoder_levels,
+            "downs_t": self.encoder_downs_t,
+            "strides_t": self.encoder_strides_t,
+            "width": self.encoder_width,
+            "depth": self.encoder_depth,
+            "m_conv": self.encoder_m_conv,
+            "dilation_growth_rate": self.encoder_dilation_growth_rate
+        }
+    
+    @property
+    def f0_decoder_params(self):
+        return {
+            "input_emb_width": self.decoder_input_emb_width,
+            "output_emb_width": self.decoder_output_emb_width,
+            "levels": self.decoder_levels,
+            "downs_t": self.decoder_downs_t,
+            "strides_t": self.decoder_strides_t,
+            "width": self.decoder_width,
+            "depth": self.decoder_depth,
+            "m_conv": self.decoder_m_conv,
+            "dilation_growth_rate": self.decoder_dilation_growth_rate
+        }
+    
+    @classmethod
+    def from_yaml(cls, yaml_path: str):
+        """Load config from yaml file"""
+        import yaml
+        with open(yaml_path, 'r') as f:
+            config = yaml.safe_load(f)
+        
+        # Convert yaml config to kwargs
+        kwargs = {
+            # VQ params
+            **{k: v for k, v in config['f0_vq_params'].items()},
+            
+            # Encoder params
+            **{f"encoder_{k}": v for k, v in config['f0_encoder_params'].items()},
+            
+            # Decoder params  
+            **{f"decoder_{k}": v for k, v in config['f0_decoder_params'].items()},
+            
+            # Training params
+            "lambda_commit": config.get('lambda_commit', 0.02),
+            "f0_normalize": config.get('f0_normalize', True),
+            "intensity_normalize": config.get('intensity_normalize', True),
+            "multispkr": config.get('multispkr', "single"),
+            "f0_feats": config.get('f0_feats', False),
+            "f0_median": config.get('f0_median', False),
+            
+            # Training hyperparams
+            "learning_rate": config.get('learning_rate', 0.0002),
+            "adam_b1": config.get('adam_b1', 0.8), 
+            "adam_b2": config.get('adam_b2', 0.99),
+            "lr_decay": config.get('lr_decay', 0.999),
+        }
+        
+        return cls(**kwargs)

utils.py

@@ -0,0 +1,36 @@
+# # Copyright (c) Facebook, Inc. and its affiliates.
+# # All rights reserved.
+# #
+# # This source code is licensed under the license found in the
+# # LICENSE file in the root directory of this source tree.
+#
+# # Adapted from https://github.com/jik876/hifi-gan
+#
+import os
+import torch
+
+
+def init_weights(m, mean=0.0, std=0.01):
+    classname = m.__class__.__name__
+    if classname.find("Conv") != -1:
+        m.weight.data.normal_(mean, std)
+
+
+def get_padding(kernel_size, dilation=1):
+    return int((kernel_size*dilation - dilation)/2)
+
+
+def load_checkpoint(filepath, device):
+    assert os.path.isfile(filepath)
+    print("Loading '{}'".format(filepath))
+    checkpoint_dict = torch.load(filepath, map_location=device)
+    print("Complete.")
+    return checkpoint_dict
+
+
+class AttrDict(dict):
+    def __init__(self, *args, **kwargs):
+        super(AttrDict, self).__init__(*args, **kwargs)
+        self.__dict__ = self
+
+