File size: 6,526 Bytes
fb0011a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 | # models_speaker.py
# ==============================================================================
# Model builder for Speaker-based codec (learnable speaker embeddings)
# ==============================================================================
import os
import os.path as osp
import copy
import math
import random
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.utils import spectral_norm
from torch.nn.utils.parametrizations import weight_norm
from collections import OrderedDict
from munch import Munch
import yaml
import json
import torchaudio
from typing import Optional, Tuple, List, Sequence
from Modules.text_encoder import TextEncoderTransformer
from Modules.codec_decoder_speaker import HybridTTSCodecVocoderSpeaker
from Modules.discriminators import MultiPeriodDiscriminator, MultiResSpecDiscriminator
def gaussian_upsample(
token_feats: torch.Tensor, # [B, D, N]
durations: torch.Tensor, # [B, N] (int or float) - can be fractional
T: torch.Tensor, # [B] frames per sample
sigma_sq: float = 10.0,
token_mask: Optional[torch.Tensor] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Numerically-stable gaussian upsampling.
Uses a softmax over negative squared distances (log-sum-exp stability).
"""
B, D, N = token_feats.shape
device = token_feats.device
T_list = [int(T[b].item()) for b in range(B)]
max_T = max(T_list) if len(T_list) > 0 else 0
out_list = []
weights_list = []
for b in range(B):
Nb = int(token_mask[b].sum().item()) if token_mask is not None else N
Tb = T_list[b]
if Nb <= 0 or Tb <= 0:
out_list.append(torch.zeros(D, max_T, device=device, dtype=token_feats.dtype))
weights_list.append(torch.zeros(N, max_T, device=device, dtype=token_feats.dtype))
continue
E = token_feats[b, :, :Nb] # [D, Nb]
d = durations[b, :Nb].float() # [Nb]
d_sum = d.sum().clamp(min=1e-8)
scale = Tb / d_sum
l = d * scale # [Nb] (fractional frames)
# centers
c = torch.cumsum(l, dim=0) - 0.5 * l # [Nb]
t = torch.arange(Tb, device=device).float().unsqueeze(0) # [1, Tb]
# squared distances
dist2 = (t - c.unsqueeze(1)) ** 2 # [Nb, Tb]
# stable softmax along token axis (dim=0)
logits = (-dist2 / float(sigma_sq)).to(torch.float32)
w = F.softmax(logits, dim=0).to(dist2.dtype) # [Nb, Tb]
# upsample
F_bt = E @ w # [D, Tb]
# pad to max_T
if Tb < max_T:
F_pad = torch.zeros(D, max_T, device=device, dtype=F_bt.dtype)
F_pad[:, :Tb] = F_bt
F_bt = F_pad
out_list.append(F_bt)
# pack back to [N, max_T] with padding zeros
w_full = torch.zeros(N, max_T, device=device, dtype=w.dtype)
w_full[:Nb, :Tb] = w
weights_list.append(w_full)
frame_feats = torch.stack(out_list, dim=0) # [B, D, max_T]
weights = torch.stack(weights_list, dim=0) # [B, N, max_T]
return frame_feats, weights
def build_model_speaker(args):
"""
Build model with speaker-based codec (learnable speaker embeddings).
Args:
args: Munch config with model parameters including:
- num_speakers: number of speakers (default 11 for IDs 0-10)
- speaker_dim: speaker embedding dimension (default 128)
- hidden_dim: text encoder hidden dim
- prosody_latent_dim: prosody latent dimension
- codec_hidden_dim: decoder hidden dimension
- fsq_levels: FSQ quantization levels
- decoder.*: decoder parameters
- n_token: vocabulary size
"""
hop_length = 441
fsq_levels = args.get('fsq_levels', [4] * 6)
codec_strides = [2]
codebook_size = np.prod(fsq_levels)
# Speaker-based codec with learnable embeddings
codec = HybridTTSCodecVocoderSpeaker(
# Speaker settings
num_speakers=args.get('num_speakers', 11), # 11 speakers: IDs 0-10
speaker_dim=args.get('speaker_dim', 128), # 128-dim learnable embeddings
# Input dimensions
text_dim=args.hidden_dim,
# Prosody latent settings
prosody_latent_dim=args.prosody_latent_dim,
hidden_dim=args.codec_hidden_dim,
# Compression
codec_strides=codec_strides,
# FSQ settings
codebook_size=codebook_size,
fsq_levels=fsq_levels,
# Decoder settings
upsample_rates=args.decoder.upsample_rates,
gen_istft_n_fft=args.decoder.gen_istft_n_fft,
gen_istft_hop_size=args.decoder.gen_istft_hop_size,
source_upsample_rate=hop_length,
language_dim=64,
)
text_encoder = TextEncoderTransformer(
channels=512,
language_count=7,
language_hidden_dim=64,
depth=6,
kernel_size=5,
n_symbols=args.n_token
)
nets = Munch(
codec=codec,
text_encoder=text_encoder,
mpd=MultiPeriodDiscriminator(),
msd=MultiResSpecDiscriminator(),
)
return nets
def load_checkpoint(model, optimizer, path, load_only_params=False, ignore_modules=[]):
"""Load checkpoint with optional module filtering."""
state = torch.load(path, map_location='cpu', weights_only=False)
params = state['net']
print('Loading checkpoint...')
for key in model:
if key in params and key not in ignore_modules:
try:
model[key].load_state_dict(params[key], strict=True)
except:
from collections import OrderedDict
state_dict = params[key]
new_state_dict = OrderedDict()
print(f'{key} key length: {len(model[key].state_dict().keys())}, '
f'state_dict key length: {len(state_dict.keys())}')
for (k_m, v_m), (k_c, v_c) in zip(model[key].state_dict().items(),
state_dict.items()):
new_state_dict[k_m] = v_c
model[key].load_state_dict(new_state_dict, strict=True)
print(f'{key} loaded')
if not load_only_params:
epoch = state["epoch"]
iters = state["iters"]
optimizer.load_state_dict(state["optimizer"])
else:
epoch = 0
iters = 0
return model, optimizer, epoch, iters
|