chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,7 @@
|
||||
*.ds
|
||||
*.onnx
|
||||
*.npy
|
||||
*.wav
|
||||
temp/
|
||||
cache/
|
||||
assets/
|
||||
@@ -0,0 +1,32 @@
|
||||
import numpy as np
|
||||
import onnxruntime as ort
|
||||
import tqdm
|
||||
|
||||
n_tokens = 10
|
||||
n_frames = 100
|
||||
n_runs = 20
|
||||
speedup = 20
|
||||
provider = 'DmlExecutionProvider'
|
||||
|
||||
tokens = np.array([[1] * n_tokens], dtype=np.int64)
|
||||
durations = np.array([[n_frames // n_tokens] * n_tokens], dtype=np.int64)
|
||||
f0 = np.array([[440.] * n_frames], dtype=np.float32)
|
||||
speedup = np.array(speedup, dtype=np.int64)
|
||||
|
||||
session = ort.InferenceSession('model1.onnx', providers=[provider])
|
||||
for _ in tqdm.tqdm(range(n_runs)):
|
||||
session.run(['mel'], {
|
||||
'tokens': tokens,
|
||||
'durations': durations,
|
||||
'f0': f0,
|
||||
'speedup': speedup
|
||||
})
|
||||
|
||||
session = ort.InferenceSession('model2.onnx', providers=[provider])
|
||||
for _ in tqdm.tqdm(range(n_runs)):
|
||||
session.run(['mel'], {
|
||||
'tokens': tokens,
|
||||
'durations': durations,
|
||||
'f0': f0,
|
||||
'speedup': speedup
|
||||
})
|
||||
@@ -0,0 +1,16 @@
|
||||
import numpy as np
|
||||
import onnxruntime as ort
|
||||
import tqdm
|
||||
|
||||
n_frames = 1000
|
||||
n_runs = 20
|
||||
mel = np.random.randn(1, n_frames, 128).astype(np.float32)
|
||||
f0 = np.random.randn(1, n_frames).astype(np.float32) + 440.
|
||||
provider = 'DmlExecutionProvider'
|
||||
|
||||
session = ort.InferenceSession('nsf_hifigan.onnx', providers=[provider])
|
||||
for _ in tqdm.tqdm(range(n_runs)):
|
||||
session.run(['waveform'], {
|
||||
'mel': mel,
|
||||
'f0': f0
|
||||
})
|
||||
@@ -0,0 +1,3 @@
|
||||
from .acoustic_exporter import DiffSingerAcousticExporter
|
||||
from .variance_exporter import DiffSingerVarianceExporter
|
||||
from .nsf_hifigan_exporter import NSFHiFiGANExporter
|
||||
@@ -0,0 +1,424 @@
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Union, List, Tuple, Dict
|
||||
|
||||
import onnx
|
||||
import onnxsim
|
||||
import torch
|
||||
import yaml
|
||||
|
||||
from basics.base_exporter import BaseExporter
|
||||
from deployment.modules.toplevel import DiffSingerAcousticONNX
|
||||
from modules.fastspeech.param_adaptor import VARIANCE_CHECKLIST
|
||||
from utils import load_ckpt, onnx_helper, remove_suffix
|
||||
from utils.hparams import hparams
|
||||
from utils.phoneme_utils import load_phoneme_dictionary
|
||||
|
||||
|
||||
class DiffSingerAcousticExporter(BaseExporter):
|
||||
def __init__(
|
||||
self,
|
||||
device: Union[str, torch.device] = 'cpu',
|
||||
cache_dir: Path = None,
|
||||
ckpt_steps: int = None,
|
||||
freeze_gender: float = None,
|
||||
freeze_velocity: bool = False,
|
||||
export_spk: List[Tuple[str, Dict[str, float]]] = None,
|
||||
freeze_spk: Tuple[str, Dict[str, float]] = None
|
||||
):
|
||||
super().__init__(device=device, cache_dir=cache_dir)
|
||||
# Basic attributes
|
||||
self.model_name: str = hparams['exp_name']
|
||||
self.ckpt_steps: int = ckpt_steps
|
||||
self.spk_map: dict = self.build_spk_map()
|
||||
self.lang_map: dict = self.build_lang_map()
|
||||
self.phoneme_dictionary = load_phoneme_dictionary()
|
||||
self.use_lang_id = hparams.get('use_lang_id', False) and len(self.phoneme_dictionary.cross_lingual_phonemes) > 0
|
||||
self.model = self.build_model()
|
||||
self.fs2_aux_cache_path = self.cache_dir / (
|
||||
'fs2_aux.onnx' if self.model.use_shallow_diffusion else 'fs2.onnx'
|
||||
)
|
||||
self.diffusion_cache_path = self.cache_dir / 'diffusion.onnx'
|
||||
|
||||
# Attributes for logging
|
||||
self.model_class_name = remove_suffix(self.model.__class__.__name__, 'ONNX')
|
||||
fs2_aux_cls_logging = [remove_suffix(self.model.fs2.__class__.__name__, 'ONNX')]
|
||||
if self.model.use_shallow_diffusion:
|
||||
fs2_aux_cls_logging.append(remove_suffix(
|
||||
self.model.aux_decoder.decoder.__class__.__name__, 'ONNX'
|
||||
))
|
||||
self.fs2_aux_class_name = ', '.join(fs2_aux_cls_logging)
|
||||
self.aux_decoder_class_name = remove_suffix(
|
||||
self.model.aux_decoder.decoder.__class__.__name__, 'ONNX'
|
||||
) if self.model.use_shallow_diffusion else None
|
||||
self.backbone_class_name = remove_suffix(self.model.diffusion.backbone.__class__.__name__, 'ONNX')
|
||||
self.diffusion_class_name = remove_suffix(self.model.diffusion.__class__.__name__, 'ONNX')
|
||||
|
||||
# Attributes for exporting
|
||||
self.expose_gender = freeze_gender is None
|
||||
self.expose_velocity = not freeze_velocity
|
||||
self.freeze_spk: Tuple[str, Dict[str, float]] = freeze_spk \
|
||||
if hparams['use_spk_id'] else None
|
||||
self.export_spk: List[Tuple[str, Dict[str, float]]] = export_spk \
|
||||
if hparams['use_spk_id'] and export_spk is not None else []
|
||||
if hparams['use_key_shift_embed'] and not self.expose_gender:
|
||||
shift_min, shift_max = hparams['augmentation_args']['random_pitch_shifting']['range']
|
||||
key_shift = freeze_gender * shift_max if freeze_gender >= 0. else freeze_gender * abs(shift_min)
|
||||
key_shift = max(min(key_shift, shift_max), shift_min) # clip key shift
|
||||
self.model.fs2.register_buffer('frozen_key_shift', torch.FloatTensor([key_shift]).to(self.device))
|
||||
if hparams['use_spk_id']:
|
||||
if not self.export_spk and self.freeze_spk is None:
|
||||
# In case the user did not specify any speaker settings:
|
||||
if len(self.spk_map) == 1:
|
||||
# If there is only one speaker, freeze him/her.
|
||||
first_spk = next(iter(self.spk_map.keys()))
|
||||
self.freeze_spk = (first_spk, {first_spk: 1.0})
|
||||
else:
|
||||
# If there are multiple speakers, export them all.
|
||||
self.export_spk = [(name, {name: 1.0}) for name in self.spk_map.keys()]
|
||||
if self.freeze_spk is not None:
|
||||
self.model.fs2.register_buffer('frozen_spk_embed', self._perform_spk_mix(self.freeze_spk[1]))
|
||||
|
||||
def build_model(self) -> DiffSingerAcousticONNX:
|
||||
model = DiffSingerAcousticONNX(
|
||||
vocab_size=len(self.phoneme_dictionary),
|
||||
out_dims=hparams['audio_num_mel_bins'],
|
||||
cross_lingual_token_idx=sorted({
|
||||
self.phoneme_dictionary.encode_one(p)
|
||||
for p in self.phoneme_dictionary.cross_lingual_phonemes
|
||||
})
|
||||
).eval().to(self.device)
|
||||
load_ckpt(model, hparams['work_dir'], ckpt_steps=self.ckpt_steps,
|
||||
prefix_in_ckpt='model', strict=True, device=self.device)
|
||||
return model
|
||||
|
||||
def export(self, path: Path):
|
||||
path.mkdir(parents=True, exist_ok=True)
|
||||
model_name = self.model_name
|
||||
if self.freeze_spk is not None:
|
||||
model_name += '.' + self.freeze_spk[0]
|
||||
self.export_model(path / f'{model_name}.onnx')
|
||||
self.export_attachments(path)
|
||||
|
||||
def export_model(self, path: Path):
|
||||
self._torch_export_model()
|
||||
fs2_aux_onnx = self._optimize_fs2_aux_graph(onnx.load(self.fs2_aux_cache_path))
|
||||
diffusion_onnx = self._optimize_diffusion_graph(onnx.load(self.diffusion_cache_path))
|
||||
model_onnx = self._merge_fs2_aux_diffusion_graphs(fs2_aux_onnx, diffusion_onnx)
|
||||
onnx.save(model_onnx, path)
|
||||
self.fs2_aux_cache_path.unlink()
|
||||
self.diffusion_cache_path.unlink()
|
||||
print(f'| export model => {path}')
|
||||
|
||||
def export_attachments(self, path: Path):
|
||||
for spk in self.export_spk:
|
||||
self._export_spk_embed(
|
||||
path / f'{self.model_name}.{spk[0]}.emb',
|
||||
self._perform_spk_mix(spk[1])
|
||||
)
|
||||
self.export_dictionaries(path)
|
||||
self._export_phonemes(path)
|
||||
|
||||
model_name = self.model_name
|
||||
if self.freeze_spk is not None:
|
||||
model_name += '.' + self.freeze_spk[0]
|
||||
dsconfig = {
|
||||
# basic configs
|
||||
'phonemes': f'{self.model_name}.phonemes.json',
|
||||
'languages': f'{self.model_name}.languages.json',
|
||||
'use_lang_id': self.use_lang_id,
|
||||
'acoustic': f'{model_name}.onnx',
|
||||
'hidden_size': hparams['hidden_size'],
|
||||
'vocoder': 'pc_nsf_hifigan_44.1k_hop512_128bin_2025.02',
|
||||
}
|
||||
# multi-speaker
|
||||
if len(self.export_spk) > 0:
|
||||
dsconfig['speakers'] = [f'{self.model_name}.{spk[0]}' for spk in self.export_spk]
|
||||
# parameters
|
||||
if self.expose_gender:
|
||||
dsconfig['augmentation_args'] = {
|
||||
'random_pitch_shifting': {
|
||||
'range': hparams['augmentation_args']['random_pitch_shifting']['range']
|
||||
}
|
||||
}
|
||||
dsconfig['use_key_shift_embed'] = self.expose_gender
|
||||
dsconfig['use_speed_embed'] = self.expose_velocity
|
||||
for variance in VARIANCE_CHECKLIST:
|
||||
dsconfig[f'use_{variance}_embed'] = (variance in self.model.fs2.variance_embed_list)
|
||||
# sampling acceleration and shallow diffusion
|
||||
dsconfig['use_continuous_acceleration'] = True
|
||||
dsconfig['use_variable_depth'] = self.model.use_shallow_diffusion
|
||||
dsconfig['max_depth'] = 1 - self.model.diffusion.t_start
|
||||
# mel specification
|
||||
dsconfig['sample_rate'] = hparams['audio_sample_rate']
|
||||
dsconfig['hop_size'] = hparams['hop_size']
|
||||
dsconfig['win_size'] = hparams['win_size']
|
||||
dsconfig['fft_size'] = hparams['fft_size']
|
||||
dsconfig['num_mel_bins'] = hparams['audio_num_mel_bins']
|
||||
dsconfig['mel_fmin'] = hparams['fmin']
|
||||
dsconfig['mel_fmax'] = hparams['fmax'] if hparams['fmax'] is not None else hparams['audio_sample_rate'] / 2
|
||||
dsconfig['mel_base'] = 'e'
|
||||
dsconfig['mel_scale'] = 'slaney'
|
||||
config_path = path / 'dsconfig.yaml'
|
||||
with open(config_path, 'w', encoding='utf8') as fw:
|
||||
yaml.safe_dump(dsconfig, fw, sort_keys=False)
|
||||
print(f'| export configs => {config_path} **PLEASE EDIT BEFORE USE**')
|
||||
|
||||
@torch.no_grad()
|
||||
def _torch_export_model(self):
|
||||
# Prepare inputs for FastSpeech2 and aux decoder tracing
|
||||
n_frames = 10
|
||||
tokens = torch.LongTensor([[1]]).to(self.device)
|
||||
durations = torch.LongTensor([[n_frames]]).to(self.device)
|
||||
f0 = torch.FloatTensor([[440.] * n_frames]).to(self.device)
|
||||
variances = {
|
||||
v_name: torch.zeros(1, n_frames, dtype=torch.float32, device=self.device)
|
||||
for v_name in self.model.fs2.variance_embed_list
|
||||
}
|
||||
kwargs: Dict[str, torch.Tensor] = {}
|
||||
arguments = (tokens, durations, f0, variances, kwargs)
|
||||
input_names = ['tokens', 'durations', 'f0'] + self.model.fs2.variance_embed_list
|
||||
dynamic_axes = {
|
||||
'tokens': {
|
||||
1: 'n_tokens'
|
||||
},
|
||||
'durations': {
|
||||
1: 'n_tokens'
|
||||
},
|
||||
'f0': {
|
||||
1: 'n_frames'
|
||||
},
|
||||
**{
|
||||
v_name: {
|
||||
1: 'n_frames'
|
||||
}
|
||||
for v_name in self.model.fs2.variance_embed_list
|
||||
}
|
||||
}
|
||||
if hparams['use_key_shift_embed']:
|
||||
if self.expose_gender:
|
||||
kwargs['gender'] = torch.rand((1, n_frames), dtype=torch.float32, device=self.device)
|
||||
input_names.append('gender')
|
||||
dynamic_axes['gender'] = {
|
||||
1: 'n_frames'
|
||||
}
|
||||
if hparams['use_speed_embed']:
|
||||
if self.expose_velocity:
|
||||
kwargs['velocity'] = torch.rand((1, n_frames), dtype=torch.float32, device=self.device)
|
||||
input_names.append('velocity')
|
||||
dynamic_axes['velocity'] = {
|
||||
1: 'n_frames'
|
||||
}
|
||||
if hparams['use_spk_id'] and not self.freeze_spk:
|
||||
kwargs['spk_embed'] = torch.rand(
|
||||
(1, n_frames, hparams['hidden_size']),
|
||||
dtype=torch.float32, device=self.device
|
||||
)
|
||||
input_names.append('spk_embed')
|
||||
dynamic_axes['spk_embed'] = {
|
||||
1: 'n_frames'
|
||||
}
|
||||
if self.use_lang_id:
|
||||
kwargs['languages'] = torch.zeros_like(tokens)
|
||||
input_names.append('languages')
|
||||
dynamic_axes['languages'] = {
|
||||
1: 'n_tokens'
|
||||
}
|
||||
dynamic_axes['condition'] = {
|
||||
1: 'n_frames'
|
||||
}
|
||||
|
||||
# PyTorch ONNX export for FastSpeech2 and aux decoder
|
||||
output_names = ['condition']
|
||||
if self.model.use_shallow_diffusion:
|
||||
output_names.append('aux_mel')
|
||||
dynamic_axes['aux_mel'] = {
|
||||
1: 'n_frames'
|
||||
}
|
||||
print(f'Exporting {self.fs2_aux_class_name}...')
|
||||
torch.onnx.export(
|
||||
self.model.view_as_fs2_aux(),
|
||||
arguments,
|
||||
self.fs2_aux_cache_path,
|
||||
input_names=input_names,
|
||||
output_names=output_names,
|
||||
dynamic_axes=dynamic_axes,
|
||||
opset_version=17,
|
||||
**onnx_helper.TORCHSCRIPT_EXPORT_KWARGS
|
||||
)
|
||||
|
||||
condition = torch.rand((1, n_frames, hparams['hidden_size']), device=self.device)
|
||||
|
||||
# Prepare inputs for backbone tracing and GaussianDiffusion scripting
|
||||
shape = (1, 1, hparams['audio_num_mel_bins'], n_frames)
|
||||
noise = torch.randn(shape, device=self.device)
|
||||
x_aux = torch.randn((1, n_frames, hparams['audio_num_mel_bins']), device=self.device)
|
||||
dummy_time = (torch.rand((1,), device=self.device) * self.model.diffusion.time_scale_factor).float()
|
||||
dummy_depth = torch.tensor(0.1, device=self.device)
|
||||
dummy_steps = 5
|
||||
|
||||
print(f'Tracing {self.backbone_class_name} backbone...')
|
||||
if self.model.diffusion_type == 'ddpm':
|
||||
major_mel_decoder = self.model.view_as_diffusion()
|
||||
elif self.model.diffusion_type == 'reflow':
|
||||
major_mel_decoder = self.model.view_as_reflow()
|
||||
else:
|
||||
raise ValueError(f'Invalid diffusion type: {self.model.diffusion_type}')
|
||||
major_mel_decoder.diffusion.set_backbone(
|
||||
torch.jit.trace(
|
||||
major_mel_decoder.diffusion.backbone,
|
||||
(
|
||||
noise,
|
||||
dummy_time,
|
||||
condition.transpose(1, 2)
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
print(f'Scripting {self.diffusion_class_name}...')
|
||||
diffusion_inputs = [
|
||||
condition,
|
||||
*([x_aux, dummy_depth] if self.model.use_shallow_diffusion else [])
|
||||
]
|
||||
major_mel_decoder = torch.jit.script(
|
||||
major_mel_decoder,
|
||||
example_inputs=[
|
||||
(
|
||||
*diffusion_inputs,
|
||||
1 # p_sample branch
|
||||
),
|
||||
(
|
||||
*diffusion_inputs,
|
||||
dummy_steps # p_sample_plms branch
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
# PyTorch ONNX export for GaussianDiffusion
|
||||
print(f'Exporting {self.diffusion_class_name}...')
|
||||
torch.onnx.export(
|
||||
major_mel_decoder,
|
||||
(
|
||||
*diffusion_inputs,
|
||||
dummy_steps
|
||||
),
|
||||
self.diffusion_cache_path,
|
||||
input_names=[
|
||||
'condition',
|
||||
*(['x_aux', 'depth'] if self.model.use_shallow_diffusion else []),
|
||||
'steps'
|
||||
],
|
||||
output_names=[
|
||||
'mel'
|
||||
],
|
||||
dynamic_axes={
|
||||
'condition': {
|
||||
1: 'n_frames'
|
||||
},
|
||||
**({'x_aux': {1: 'n_frames'}} if self.model.use_shallow_diffusion else {}),
|
||||
'mel': {
|
||||
1: 'n_frames'
|
||||
}
|
||||
},
|
||||
opset_version=17,
|
||||
**onnx_helper.TORCHSCRIPT_EXPORT_KWARGS
|
||||
)
|
||||
|
||||
@torch.no_grad()
|
||||
def _perform_spk_mix(self, spk_mix: Dict[str, float]):
|
||||
spk_mix_ids = []
|
||||
spk_mix_values = []
|
||||
for name, value in spk_mix.items():
|
||||
spk_mix_ids.append(self.spk_map[name])
|
||||
assert value >= 0., f'Speaker mix checks failed.\n' \
|
||||
f'Proportion of speaker \'{name}\' is negative.'
|
||||
spk_mix_values.append(value)
|
||||
spk_mix_id_N = torch.LongTensor(spk_mix_ids).to(self.device)[None] # => [1, N]
|
||||
spk_mix_value_N = torch.FloatTensor(spk_mix_values).to(self.device)[None] # => [1, N]
|
||||
spk_mix_value_sum = spk_mix_value_N.sum()
|
||||
assert spk_mix_value_sum > 0., 'Speaker mix checks failed.\n' \
|
||||
'Proportions of speaker mix sum to zero.'
|
||||
spk_mix_value_N /= spk_mix_value_sum # normalize
|
||||
spk_mix_embed = torch.sum(
|
||||
self.model.fs2.spk_embed(spk_mix_id_N) * spk_mix_value_N.unsqueeze(2), # => [1, N, H]
|
||||
dim=1, keepdim=True
|
||||
) # => [1, 1, H]
|
||||
return spk_mix_embed
|
||||
|
||||
def _optimize_fs2_aux_graph(self, fs2: onnx.ModelProto) -> onnx.ModelProto:
|
||||
print(f'Running ONNX Simplifier on {self.fs2_aux_class_name}...')
|
||||
fs2, check = onnxsim.simplify(fs2, include_subgraph=True)
|
||||
assert check, 'Simplified ONNX model could not be validated'
|
||||
onnx_helper.model_reorder_io_list(
|
||||
fs2, 'input',
|
||||
target_name='languages', insert_after_name='tokens'
|
||||
)
|
||||
print(f'| optimize graph: {self.fs2_aux_class_name}')
|
||||
return fs2
|
||||
|
||||
def _optimize_diffusion_graph(self, diffusion: onnx.ModelProto) -> onnx.ModelProto:
|
||||
onnx_helper.model_override_io_shapes(diffusion, output_shapes={
|
||||
'mel': (1, 'n_frames', hparams['audio_num_mel_bins'])
|
||||
})
|
||||
print(f'Running ONNX Simplifier #1 on {self.diffusion_class_name}...')
|
||||
diffusion, check = onnxsim.simplify(diffusion, include_subgraph=True)
|
||||
assert check, 'Simplified ONNX model could not be validated'
|
||||
onnx_helper.graph_fold_back_to_squeeze(diffusion.graph)
|
||||
onnx_helper.graph_extract_conditioner_projections(
|
||||
graph=diffusion.graph, op_type='Conv',
|
||||
weight_pattern=r'diffusion\..*\.conditioner_projection\.weight',
|
||||
alias_prefix='/diffusion/backbone/cache'
|
||||
)
|
||||
onnx_helper.graph_remove_unused_values(diffusion.graph)
|
||||
print(f'Running ONNX Simplifier #2 on {self.diffusion_class_name}...')
|
||||
diffusion, check = onnxsim.simplify(
|
||||
diffusion,
|
||||
include_subgraph=True
|
||||
)
|
||||
assert check, 'Simplified ONNX model could not be validated'
|
||||
print(f'| optimize graph: {self.diffusion_class_name}')
|
||||
return diffusion
|
||||
|
||||
def _merge_fs2_aux_diffusion_graphs(self, fs2: onnx.ModelProto, diffusion: onnx.ModelProto) -> onnx.ModelProto:
|
||||
onnx_helper.model_add_prefixes(
|
||||
fs2, dim_prefix=('fs2aux.' if self.model.use_shallow_diffusion else 'fs2.'),
|
||||
ignored_pattern=r'(n_tokens)|(n_frames)'
|
||||
)
|
||||
onnx_helper.model_add_prefixes(diffusion, dim_prefix='diffusion.', ignored_pattern='n_frames')
|
||||
print(f'Merging {self.fs2_aux_class_name} and {self.diffusion_class_name} '
|
||||
f'back into {self.model_class_name}...')
|
||||
merged = onnx.compose.merge_models(
|
||||
fs2, diffusion, io_map=[
|
||||
('condition', 'condition'),
|
||||
*([('aux_mel', 'x_aux')] if self.model.use_shallow_diffusion else []),
|
||||
],
|
||||
prefix1='', prefix2='', doc_string='',
|
||||
producer_name=fs2.producer_name, producer_version=fs2.producer_version,
|
||||
domain=fs2.domain, model_version=fs2.model_version
|
||||
)
|
||||
merged.graph.name = fs2.graph.name
|
||||
|
||||
print(f'Running ONNX Simplifier on {self.model_class_name}...')
|
||||
merged, check = onnxsim.simplify(
|
||||
merged,
|
||||
include_subgraph=True
|
||||
)
|
||||
assert check, 'Simplified ONNX model could not be validated'
|
||||
print(f'| optimize graph: {self.model_class_name}')
|
||||
|
||||
return merged
|
||||
|
||||
# noinspection PyMethodMayBeStatic
|
||||
def _export_spk_embed(self, path: Path, spk_embed: torch.Tensor):
|
||||
with open(path, 'wb') as f:
|
||||
f.write(spk_embed.cpu().numpy().tobytes())
|
||||
print(f'| export spk embed => {path}')
|
||||
|
||||
def _export_phonemes(self, path: Path):
|
||||
ph_path = path / f'{self.model_name}.phonemes.json'
|
||||
self.phoneme_dictionary.dump(ph_path)
|
||||
print(f'| export phonemes => {ph_path}')
|
||||
lang_path = path / f'{self.model_name}.languages.json'
|
||||
with open(lang_path, 'w', encoding='utf8') as f:
|
||||
json.dump(self.lang_map, f, ensure_ascii=False, indent=2)
|
||||
print(f'| export languages => {lang_path}')
|
||||
@@ -0,0 +1,127 @@
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Union
|
||||
|
||||
import onnx
|
||||
import onnxsim
|
||||
import torch
|
||||
import yaml
|
||||
from torch import nn
|
||||
|
||||
from basics.base_exporter import BaseExporter
|
||||
from deployment.modules.nsf_hifigan import NSFHiFiGANONNX
|
||||
from utils import load_ckpt, onnx_helper, remove_suffix
|
||||
from utils.hparams import hparams
|
||||
|
||||
|
||||
class NSFHiFiGANExporter(BaseExporter):
|
||||
def __init__(
|
||||
self,
|
||||
device: Union[str, torch.device] = 'cpu',
|
||||
cache_dir: Path = None,
|
||||
model_path: Path = None,
|
||||
model_name: str = 'nsf_hifigan'
|
||||
):
|
||||
super().__init__(device=device, cache_dir=cache_dir)
|
||||
self.model_path = model_path
|
||||
self.model_name = model_name
|
||||
self.vocoder_pitch_controllable = False
|
||||
self.model = self.build_model()
|
||||
self.model_class_name = remove_suffix(self.model.__class__.__name__, 'ONNX')
|
||||
self.model_cache_path = (self.cache_dir / self.model_name).with_suffix('.onnx')
|
||||
|
||||
def build_model(self) -> nn.Module:
|
||||
config_path = self.model_path.with_name('config.json')
|
||||
with open(config_path, 'r', encoding='utf8') as f:
|
||||
config = json.load(f)
|
||||
assert hparams.get('mel_base') == 'e', (
|
||||
"Mel base must be set to \'e\' according to 2nd stage of the migration plan. "
|
||||
"See https://github.com/openvpi/DiffSinger/releases/tag/v2.3.0 for more details."
|
||||
)
|
||||
model = NSFHiFiGANONNX(config).eval().to(self.device)
|
||||
self.vocoder_pitch_controllable = config.get("pc_aug", False)
|
||||
load_ckpt(model.generator, str(self.model_path),
|
||||
prefix_in_ckpt=None, key_in_ckpt='generator',
|
||||
strict=True, device=self.device)
|
||||
model.generator.remove_weight_norm()
|
||||
return model
|
||||
|
||||
def export(self, path: Path):
|
||||
path.mkdir(parents=True, exist_ok=True)
|
||||
self.export_model(path / self.model_cache_path.name)
|
||||
self.export_attachments(path)
|
||||
|
||||
def export_model(self, path: Path):
|
||||
self._torch_export_model()
|
||||
model_onnx = self._optimize_model_graph(onnx.load(self.model_cache_path))
|
||||
onnx.save(model_onnx, path)
|
||||
self.model_cache_path.unlink()
|
||||
print(f'| export model => {path}')
|
||||
|
||||
def export_attachments(self, path: Path):
|
||||
config_path = path / 'vocoder.yaml'
|
||||
with open(config_path, 'w', encoding='utf8') as fw:
|
||||
yaml.safe_dump({
|
||||
# basic configs
|
||||
'name': self.model_name,
|
||||
'model': self.model_cache_path.name,
|
||||
# mel specifications
|
||||
'sample_rate': hparams['audio_sample_rate'],
|
||||
'hop_size': hparams['hop_size'],
|
||||
'win_size': hparams['win_size'],
|
||||
'fft_size': hparams['fft_size'],
|
||||
'num_mel_bins': hparams['audio_num_mel_bins'],
|
||||
'mel_fmin': hparams['fmin'],
|
||||
'mel_fmax': hparams['fmax'] if hparams['fmax'] is not None else hparams['audio_sample_rate'] / 2,
|
||||
'mel_base': 'e',
|
||||
'mel_scale': 'slaney',
|
||||
'pitch_controllable': self.vocoder_pitch_controllable,
|
||||
# Some old vocoder versions may have severe performance issues on CUDA;
|
||||
# the issues were fixed in newer versions, and this flag is to distinguish them
|
||||
'force_on_cpu': False,
|
||||
}, fw, sort_keys=False)
|
||||
print(f'| export configs => {config_path} **PLEASE EDIT BEFORE USE**')
|
||||
|
||||
@torch.no_grad()
|
||||
def _torch_export_model(self):
|
||||
# Prepare inputs for NSFHiFiGAN
|
||||
n_frames = 10
|
||||
mel = torch.randn((1, n_frames, hparams['audio_num_mel_bins']), dtype=torch.float32, device=self.device)
|
||||
f0 = torch.randn((1, n_frames), dtype=torch.float32, device=self.device) + 440.
|
||||
|
||||
# PyTorch ONNX export for NSFHiFiGAN
|
||||
print(f'Exporting {self.model_class_name}...')
|
||||
torch.onnx.export(
|
||||
self.model,
|
||||
(
|
||||
mel,
|
||||
f0
|
||||
),
|
||||
self.model_cache_path,
|
||||
input_names=[
|
||||
'mel',
|
||||
'f0'
|
||||
],
|
||||
output_names=[
|
||||
'waveform'
|
||||
],
|
||||
dynamic_axes={
|
||||
'mel': {
|
||||
1: 'n_frames'
|
||||
},
|
||||
'f0': {
|
||||
1: 'n_frames'
|
||||
},
|
||||
'waveform': {
|
||||
1: 'n_samples'
|
||||
}
|
||||
},
|
||||
opset_version=17,
|
||||
**onnx_helper.TORCHSCRIPT_EXPORT_KWARGS
|
||||
)
|
||||
|
||||
def _optimize_model_graph(self, model: onnx.ModelProto) -> onnx.ModelProto:
|
||||
print(f'Running ONNX simplifier for {self.model_class_name}...')
|
||||
model, check = onnxsim.simplify(model, include_subgraph=True)
|
||||
assert check, 'Simplified ONNX model could not be validated'
|
||||
return model
|
||||
@@ -0,0 +1,808 @@
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Union, List, Tuple, Dict
|
||||
|
||||
import onnx
|
||||
import onnxsim
|
||||
import torch
|
||||
import yaml
|
||||
|
||||
from basics.base_exporter import BaseExporter
|
||||
from deployment.modules.toplevel import DiffSingerVarianceONNX
|
||||
from modules.fastspeech.param_adaptor import VARIANCE_CHECKLIST
|
||||
from utils import load_ckpt, onnx_helper, remove_suffix
|
||||
from utils.hparams import hparams
|
||||
from utils.phoneme_utils import load_phoneme_dictionary
|
||||
|
||||
|
||||
class DiffSingerVarianceExporter(BaseExporter):
|
||||
def __init__(
|
||||
self,
|
||||
device: Union[str, torch.device] = 'cpu',
|
||||
cache_dir: Path = None,
|
||||
ckpt_steps: int = None,
|
||||
freeze_glide: bool = False,
|
||||
freeze_expr: bool = False,
|
||||
export_spk: List[Tuple[str, Dict[str, float]]] = None,
|
||||
freeze_spk: Tuple[str, Dict[str, float]] = None
|
||||
):
|
||||
super().__init__(device=device, cache_dir=cache_dir)
|
||||
# Basic attributes
|
||||
self.model_name: str = hparams['exp_name']
|
||||
self.ckpt_steps: int = ckpt_steps
|
||||
self.spk_map: dict = self.build_spk_map()
|
||||
self.lang_map: dict = self.build_lang_map()
|
||||
self.phoneme_dictionary = load_phoneme_dictionary()
|
||||
self.use_lang_id = hparams.get('use_lang_id', False) and len(self.phoneme_dictionary.cross_lingual_phonemes) > 0
|
||||
self.model = self.build_model()
|
||||
self.linguistic_encoder_cache_path = self.cache_dir / 'linguistic.onnx'
|
||||
self.dur_predictor_cache_path = self.cache_dir / 'dur.onnx'
|
||||
self.pitch_preprocess_cache_path = self.cache_dir / 'pitch_pre.onnx'
|
||||
self.pitch_predictor_cache_path = self.cache_dir / 'pitch.onnx'
|
||||
self.pitch_postprocess_cache_path = self.cache_dir / 'pitch_post.onnx'
|
||||
self.variance_preprocess_cache_path = self.cache_dir / 'variance_pre.onnx'
|
||||
self.multi_var_predictor_cache_path = self.cache_dir / 'variance.onnx'
|
||||
self.variance_postprocess_cache_path = self.cache_dir / 'variance_post.onnx'
|
||||
|
||||
# Attributes for logging
|
||||
self.fs2_class_name = remove_suffix(self.model.fs2.__class__.__name__, 'ONNX')
|
||||
self.dur_predictor_class_name = \
|
||||
remove_suffix(self.model.fs2.dur_predictor.__class__.__name__, 'ONNX') \
|
||||
if self.model.predict_dur else None
|
||||
self.pitch_backbone_class_name = \
|
||||
remove_suffix(self.model.pitch_predictor.backbone.__class__.__name__, 'ONNX') \
|
||||
if self.model.predict_pitch else None
|
||||
self.pitch_predictor_class_name = \
|
||||
remove_suffix(self.model.pitch_predictor.__class__.__name__, 'ONNX') \
|
||||
if self.model.predict_pitch else None
|
||||
self.variance_backbone_class_name = \
|
||||
remove_suffix(self.model.variance_predictor.backbone.__class__.__name__, 'ONNX') \
|
||||
if self.model.predict_variances else None
|
||||
self.multi_var_predictor_class_name = \
|
||||
remove_suffix(self.model.variance_predictor.__class__.__name__, 'ONNX') \
|
||||
if self.model.predict_variances else None
|
||||
|
||||
# Attributes for exporting
|
||||
self.expose_expr = not freeze_expr
|
||||
self.freeze_glide = freeze_glide
|
||||
self.freeze_spk: Tuple[str, Dict[str, float]] = freeze_spk \
|
||||
if hparams['use_spk_id'] else None
|
||||
self.export_spk: List[Tuple[str, Dict[str, float]]] = export_spk \
|
||||
if hparams['use_spk_id'] and export_spk is not None else []
|
||||
if hparams['use_spk_id']:
|
||||
if not self.export_spk and self.freeze_spk is None:
|
||||
# In case the user did not specify any speaker settings:
|
||||
if len(self.spk_map) == 1:
|
||||
# If there is only one speaker, freeze him/her.
|
||||
first_spk = next(iter(self.spk_map.keys()))
|
||||
self.freeze_spk = (first_spk, {first_spk: 1.0})
|
||||
else:
|
||||
# If there are multiple speakers, export them all.
|
||||
self.export_spk = [(name, {name: 1.0}) for name in self.spk_map.keys()]
|
||||
if self.freeze_spk is not None:
|
||||
self.model.register_buffer('frozen_spk_embed', self._perform_spk_mix(self.freeze_spk[1]))
|
||||
|
||||
def build_model(self) -> DiffSingerVarianceONNX:
|
||||
model = DiffSingerVarianceONNX(
|
||||
vocab_size=len(self.phoneme_dictionary),
|
||||
cross_lingual_token_idx=sorted({
|
||||
self.phoneme_dictionary.encode_one(p)
|
||||
for p in self.phoneme_dictionary.cross_lingual_phonemes
|
||||
})
|
||||
).eval().to(self.device)
|
||||
load_ckpt(model, hparams['work_dir'], ckpt_steps=self.ckpt_steps,
|
||||
prefix_in_ckpt='model', strict=True, device=self.device)
|
||||
model.build_smooth_op(self.device)
|
||||
return model
|
||||
|
||||
def export(self, path: Path):
|
||||
path.mkdir(parents=True, exist_ok=True)
|
||||
model_name = self.model_name
|
||||
if self.freeze_spk is not None:
|
||||
model_name += '.' + self.freeze_spk[0]
|
||||
self.export_model(path, model_name)
|
||||
self.export_attachments(path)
|
||||
|
||||
def export_model(self, path: Path, model_name: str = None):
|
||||
self._torch_export_model()
|
||||
linguistic_onnx = self._optimize_linguistic_graph(onnx.load(self.linguistic_encoder_cache_path))
|
||||
linguistic_path = path / f'{model_name}.linguistic.onnx'
|
||||
onnx.save(linguistic_onnx, linguistic_path)
|
||||
print(f'| export linguistic encoder => {linguistic_path}')
|
||||
self.linguistic_encoder_cache_path.unlink()
|
||||
if self.model.predict_dur:
|
||||
dur_predictor_onnx = self._optimize_dur_predictor_graph(onnx.load(self.dur_predictor_cache_path))
|
||||
dur_predictor_path = path / f'{model_name}.dur.onnx'
|
||||
onnx.save(dur_predictor_onnx, dur_predictor_path)
|
||||
self.dur_predictor_cache_path.unlink()
|
||||
print(f'| export dur predictor => {dur_predictor_path}')
|
||||
if self.model.predict_pitch:
|
||||
pitch_predictor_onnx = self._optimize_merge_pitch_predictor_graph(
|
||||
onnx.load(self.pitch_preprocess_cache_path),
|
||||
onnx.load(self.pitch_predictor_cache_path),
|
||||
onnx.load(self.pitch_postprocess_cache_path)
|
||||
)
|
||||
pitch_predictor_path = path / f'{model_name}.pitch.onnx'
|
||||
onnx.save(pitch_predictor_onnx, pitch_predictor_path)
|
||||
self.pitch_preprocess_cache_path.unlink()
|
||||
self.pitch_predictor_cache_path.unlink()
|
||||
self.pitch_postprocess_cache_path.unlink()
|
||||
print(f'| export pitch predictor => {pitch_predictor_path}')
|
||||
if self.model.predict_variances:
|
||||
variance_predictor_onnx = self._optimize_merge_variance_predictor_graph(
|
||||
onnx.load(self.variance_preprocess_cache_path),
|
||||
onnx.load(self.multi_var_predictor_cache_path),
|
||||
onnx.load(self.variance_postprocess_cache_path)
|
||||
)
|
||||
variance_predictor_path = path / f'{model_name}.variance.onnx'
|
||||
onnx.save(variance_predictor_onnx, variance_predictor_path)
|
||||
self.variance_preprocess_cache_path.unlink()
|
||||
self.multi_var_predictor_cache_path.unlink()
|
||||
self.variance_postprocess_cache_path.unlink()
|
||||
print(f'| export variance predictor => {variance_predictor_path}')
|
||||
|
||||
def export_attachments(self, path: Path):
|
||||
for spk in self.export_spk:
|
||||
self._export_spk_embed(
|
||||
path / f'{self.model_name}.{spk[0]}.emb',
|
||||
self._perform_spk_mix(spk[1])
|
||||
)
|
||||
self.export_dictionaries(path)
|
||||
self._export_phonemes(path)
|
||||
|
||||
model_name = self.model_name
|
||||
if self.freeze_spk is not None:
|
||||
model_name += '.' + self.freeze_spk[0]
|
||||
dsconfig = {
|
||||
# basic configs
|
||||
'phonemes': f'{self.model_name}.phonemes.json',
|
||||
'languages': f'{self.model_name}.languages.json',
|
||||
'use_lang_id': self.use_lang_id,
|
||||
'linguistic': f'{model_name}.linguistic.onnx',
|
||||
'hidden_size': self.model.hidden_size,
|
||||
'predict_dur': self.model.predict_dur,
|
||||
}
|
||||
# multi-speaker
|
||||
if len(self.export_spk) > 0:
|
||||
dsconfig['speakers'] = [f'{self.model_name}.{spk[0]}' for spk in self.export_spk]
|
||||
# functionalities
|
||||
if self.model.predict_dur:
|
||||
dsconfig['dur'] = f'{model_name}.dur.onnx'
|
||||
if self.model.predict_pitch:
|
||||
dsconfig['pitch'] = f'{model_name}.pitch.onnx'
|
||||
dsconfig['use_expr'] = self.expose_expr
|
||||
dsconfig['use_note_rest'] = self.model.use_melody_encoder
|
||||
if self.model.predict_variances:
|
||||
dsconfig['variance'] = f'{model_name}.variance.onnx'
|
||||
for variance in VARIANCE_CHECKLIST:
|
||||
dsconfig[f'predict_{variance}'] = (variance in self.model.variance_prediction_list)
|
||||
# sampling acceleration
|
||||
dsconfig['use_continuous_acceleration'] = True
|
||||
# frame specifications
|
||||
dsconfig['sample_rate'] = hparams['audio_sample_rate']
|
||||
dsconfig['hop_size'] = hparams['hop_size']
|
||||
config_path = path / 'dsconfig.yaml'
|
||||
with open(config_path, 'w', encoding='utf8') as fw:
|
||||
yaml.safe_dump(dsconfig, fw, sort_keys=False)
|
||||
print(f'| export configs => {config_path} **PLEASE EDIT BEFORE USE**')
|
||||
|
||||
@torch.no_grad()
|
||||
def _torch_export_model(self):
|
||||
# Prepare inputs for FastSpeech2 and dur predictor tracing
|
||||
tokens = torch.LongTensor([[1] * 5]).to(self.device)
|
||||
ph_dur = torch.LongTensor([[3, 5, 2, 1, 4]]).to(self.device)
|
||||
word_div = torch.LongTensor([[2, 2, 1]]).to(self.device)
|
||||
word_dur = torch.LongTensor([[8, 3, 4]]).to(self.device)
|
||||
languages = torch.LongTensor([[0] * 5]).to(self.device)
|
||||
encoder_out = torch.rand(1, 5, hparams['hidden_size'], dtype=torch.float32, device=self.device)
|
||||
x_masks = tokens == 0
|
||||
ph_midi = torch.LongTensor([[60] * 5]).to(self.device)
|
||||
encoder_output_names = ['encoder_out', 'x_masks']
|
||||
encoder_common_axes = {
|
||||
'encoder_out': {
|
||||
1: 'n_tokens'
|
||||
},
|
||||
'x_masks': {
|
||||
1: 'n_tokens'
|
||||
}
|
||||
}
|
||||
input_lang_id = self.use_lang_id
|
||||
input_spk_embed = hparams['use_spk_id'] and not self.freeze_spk
|
||||
|
||||
print(f'Exporting {self.fs2_class_name}...')
|
||||
if self.model.predict_dur:
|
||||
torch.onnx.export(
|
||||
self.model.view_as_linguistic_encoder(),
|
||||
(
|
||||
tokens,
|
||||
word_div,
|
||||
word_dur,
|
||||
*([languages] if input_lang_id else [])
|
||||
),
|
||||
self.linguistic_encoder_cache_path,
|
||||
input_names=[
|
||||
'tokens',
|
||||
'word_div',
|
||||
'word_dur',
|
||||
*(['languages'] if input_lang_id else [])
|
||||
],
|
||||
output_names=encoder_output_names,
|
||||
dynamic_axes={
|
||||
'tokens': {
|
||||
1: 'n_tokens'
|
||||
},
|
||||
'word_div': {
|
||||
1: 'n_words'
|
||||
},
|
||||
'word_dur': {
|
||||
1: 'n_words'
|
||||
},
|
||||
**encoder_common_axes,
|
||||
**({'languages': {1: 'n_tokens'}} if input_lang_id else {})
|
||||
},
|
||||
opset_version=17,
|
||||
**onnx_helper.TORCHSCRIPT_EXPORT_KWARGS
|
||||
)
|
||||
|
||||
print(f'Exporting {self.dur_predictor_class_name}...')
|
||||
torch.onnx.export(
|
||||
self.model.view_as_dur_predictor(),
|
||||
(
|
||||
encoder_out,
|
||||
x_masks,
|
||||
ph_midi,
|
||||
*([torch.rand(
|
||||
1, 5, hparams['hidden_size'],
|
||||
dtype=torch.float32, device=self.device
|
||||
)] if input_spk_embed else [])
|
||||
),
|
||||
self.dur_predictor_cache_path,
|
||||
input_names=[
|
||||
'encoder_out',
|
||||
'x_masks',
|
||||
'ph_midi',
|
||||
*(['spk_embed'] if input_spk_embed else [])
|
||||
],
|
||||
output_names=[
|
||||
'ph_dur_pred'
|
||||
],
|
||||
dynamic_axes={
|
||||
'ph_midi': {
|
||||
1: 'n_tokens'
|
||||
},
|
||||
'ph_dur_pred': {
|
||||
1: 'n_tokens'
|
||||
},
|
||||
**({'spk_embed': {1: 'n_tokens'}} if input_spk_embed else {}),
|
||||
**encoder_common_axes
|
||||
},
|
||||
opset_version=17,
|
||||
**onnx_helper.TORCHSCRIPT_EXPORT_KWARGS
|
||||
)
|
||||
else:
|
||||
torch.onnx.export(
|
||||
self.model.view_as_linguistic_encoder(),
|
||||
(
|
||||
tokens,
|
||||
ph_dur,
|
||||
*([languages] if input_lang_id else [])
|
||||
),
|
||||
self.linguistic_encoder_cache_path,
|
||||
input_names=[
|
||||
'tokens',
|
||||
'ph_dur',
|
||||
*(['languages'] if input_lang_id else [])
|
||||
],
|
||||
output_names=encoder_output_names,
|
||||
dynamic_axes={
|
||||
'tokens': {
|
||||
1: 'n_tokens'
|
||||
},
|
||||
'ph_dur': {
|
||||
1: 'n_tokens'
|
||||
},
|
||||
**encoder_common_axes,
|
||||
**({'languages': {1: 'n_tokens'}} if input_lang_id else {})
|
||||
},
|
||||
opset_version=17,
|
||||
**onnx_helper.TORCHSCRIPT_EXPORT_KWARGS
|
||||
)
|
||||
|
||||
# Common dummy inputs
|
||||
dummy_time = (torch.rand((1,), device=self.device) * hparams.get('time_scale_factor', 1.0)).float()
|
||||
dummy_steps = 5
|
||||
|
||||
if self.model.predict_pitch:
|
||||
use_melody_encoder = hparams.get('use_melody_encoder', False)
|
||||
use_glide_embed = use_melody_encoder and hparams['use_glide_embed'] and not self.freeze_glide
|
||||
# Prepare inputs for preprocessor of the pitch predictor
|
||||
note_midi = torch.FloatTensor([[60.] * 4]).to(self.device)
|
||||
note_dur = torch.LongTensor([[2, 6, 3, 4]]).to(self.device)
|
||||
pitch = torch.FloatTensor([[60.] * 15]).to(self.device)
|
||||
retake = torch.ones_like(pitch, dtype=torch.bool)
|
||||
pitch_input_args = (
|
||||
encoder_out,
|
||||
ph_dur,
|
||||
{
|
||||
'note_midi': note_midi,
|
||||
**({'note_rest': note_midi >= 0} if use_melody_encoder else {}),
|
||||
'note_dur': note_dur,
|
||||
**({'note_glide': torch.zeros_like(note_midi, dtype=torch.long)} if use_glide_embed else {}),
|
||||
'pitch': pitch,
|
||||
**({'expr': torch.ones_like(pitch)} if self.expose_expr else {}),
|
||||
'retake': retake,
|
||||
**({'spk_embed': torch.rand(
|
||||
1, 15, hparams['hidden_size'], dtype=torch.float32, device=self.device
|
||||
)} if input_spk_embed else {})
|
||||
}
|
||||
)
|
||||
torch.onnx.export(
|
||||
self.model.view_as_pitch_preprocess(),
|
||||
pitch_input_args,
|
||||
self.pitch_preprocess_cache_path,
|
||||
input_names=[
|
||||
'encoder_out', 'ph_dur', 'note_midi',
|
||||
*(['note_rest'] if use_melody_encoder else []),
|
||||
'note_dur',
|
||||
*(['note_glide'] if use_glide_embed else []),
|
||||
'pitch',
|
||||
*(['expr'] if self.expose_expr else []),
|
||||
'retake',
|
||||
*(['spk_embed'] if input_spk_embed else [])
|
||||
],
|
||||
output_names=[
|
||||
'pitch_cond', 'base_pitch'
|
||||
],
|
||||
dynamic_axes={
|
||||
'encoder_out': {
|
||||
1: 'n_tokens'
|
||||
},
|
||||
'ph_dur': {
|
||||
1: 'n_tokens'
|
||||
},
|
||||
'note_midi': {
|
||||
1: 'n_notes'
|
||||
},
|
||||
**({'note_rest': {1: 'n_notes'}} if use_melody_encoder else {}),
|
||||
'note_dur': {
|
||||
1: 'n_notes'
|
||||
},
|
||||
**({'note_glide': {1: 'n_notes'}} if use_glide_embed else {}),
|
||||
'pitch': {
|
||||
1: 'n_frames'
|
||||
},
|
||||
**({'expr': {1: 'n_frames'}} if self.expose_expr else {}),
|
||||
'retake': {
|
||||
1: 'n_frames'
|
||||
},
|
||||
'pitch_cond': {
|
||||
1: 'n_frames'
|
||||
},
|
||||
'base_pitch': {
|
||||
1: 'n_frames'
|
||||
},
|
||||
**({'spk_embed': {1: 'n_frames'}} if input_spk_embed else {})
|
||||
},
|
||||
opset_version=17,
|
||||
**onnx_helper.TORCHSCRIPT_EXPORT_KWARGS
|
||||
)
|
||||
|
||||
# Prepare inputs for backbone tracing and pitch predictor scripting
|
||||
shape = (1, 1, hparams['pitch_prediction_args']['repeat_bins'], 15)
|
||||
noise = torch.randn(shape, device=self.device)
|
||||
condition = torch.rand((1, hparams['hidden_size'], 15), device=self.device)
|
||||
|
||||
print(f'Tracing {self.pitch_backbone_class_name} backbone...')
|
||||
pitch_predictor = self.model.view_as_pitch_predictor()
|
||||
pitch_predictor.pitch_predictor.set_backbone(
|
||||
torch.jit.trace(
|
||||
pitch_predictor.pitch_predictor.backbone,
|
||||
(
|
||||
noise,
|
||||
dummy_time,
|
||||
condition
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
print(f'Scripting {self.pitch_predictor_class_name}...')
|
||||
pitch_predictor = torch.jit.script(
|
||||
pitch_predictor,
|
||||
example_inputs=[
|
||||
(
|
||||
condition.transpose(1, 2),
|
||||
1 # p_sample branch
|
||||
),
|
||||
(
|
||||
condition.transpose(1, 2),
|
||||
dummy_steps # p_sample_plms branch
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
print(f'Exporting {self.pitch_predictor_class_name}...')
|
||||
torch.onnx.export(
|
||||
pitch_predictor,
|
||||
(
|
||||
condition.transpose(1, 2),
|
||||
dummy_steps
|
||||
),
|
||||
self.pitch_predictor_cache_path,
|
||||
input_names=[
|
||||
'pitch_cond',
|
||||
'steps'
|
||||
],
|
||||
output_names=[
|
||||
'x_pred'
|
||||
],
|
||||
dynamic_axes={
|
||||
'pitch_cond': {
|
||||
1: 'n_frames'
|
||||
},
|
||||
'x_pred': {
|
||||
1: 'n_frames'
|
||||
}
|
||||
},
|
||||
opset_version=17,
|
||||
**onnx_helper.TORCHSCRIPT_EXPORT_KWARGS
|
||||
)
|
||||
|
||||
# Prepare inputs for postprocessor of the multi-variance predictor
|
||||
torch.onnx.export(
|
||||
self.model.view_as_pitch_postprocess(),
|
||||
(
|
||||
pitch,
|
||||
pitch
|
||||
),
|
||||
self.pitch_postprocess_cache_path,
|
||||
input_names=[
|
||||
'x_pred',
|
||||
'base_pitch'
|
||||
],
|
||||
output_names=[
|
||||
'pitch_pred'
|
||||
],
|
||||
dynamic_axes={
|
||||
'x_pred': {
|
||||
1: 'n_frames'
|
||||
},
|
||||
'base_pitch': {
|
||||
1: 'n_frames'
|
||||
},
|
||||
'pitch_pred': {
|
||||
1: 'n_frames'
|
||||
}
|
||||
},
|
||||
opset_version=17,
|
||||
**onnx_helper.TORCHSCRIPT_EXPORT_KWARGS
|
||||
)
|
||||
|
||||
if self.model.predict_variances:
|
||||
total_repeat_bins = hparams['variances_prediction_args']['total_repeat_bins']
|
||||
repeat_bins = total_repeat_bins // len(self.model.variance_prediction_list)
|
||||
|
||||
# Prepare inputs for preprocessor of the multi-variance predictor
|
||||
pitch = torch.FloatTensor([[60.] * 15]).to(self.device)
|
||||
variances = {
|
||||
v_name: torch.FloatTensor([[0.] * 15]).to(self.device)
|
||||
for v_name in self.model.variance_prediction_list
|
||||
}
|
||||
retake = torch.ones_like(pitch, dtype=torch.bool)[..., None].tile(len(self.model.variance_prediction_list))
|
||||
torch.onnx.export(
|
||||
self.model.view_as_variance_preprocess(),
|
||||
(
|
||||
encoder_out,
|
||||
ph_dur,
|
||||
pitch,
|
||||
variances,
|
||||
retake,
|
||||
*([torch.rand(
|
||||
1, 15, hparams['hidden_size'],
|
||||
dtype=torch.float32, device=self.device
|
||||
)] if input_spk_embed else [])
|
||||
),
|
||||
self.variance_preprocess_cache_path,
|
||||
input_names=[
|
||||
'encoder_out', 'ph_dur', 'pitch',
|
||||
*self.model.variance_prediction_list,
|
||||
'retake',
|
||||
*(['spk_embed'] if input_spk_embed else [])
|
||||
],
|
||||
output_names=[
|
||||
'variance_cond'
|
||||
],
|
||||
dynamic_axes={
|
||||
'encoder_out': {
|
||||
1: 'n_tokens'
|
||||
},
|
||||
'ph_dur': {
|
||||
1: 'n_tokens'
|
||||
},
|
||||
'pitch': {
|
||||
1: 'n_frames'
|
||||
},
|
||||
**{
|
||||
v_name: {
|
||||
1: 'n_frames'
|
||||
}
|
||||
for v_name in self.model.variance_prediction_list
|
||||
},
|
||||
'retake': {
|
||||
1: 'n_frames'
|
||||
},
|
||||
**({'spk_embed': {1: 'n_frames'}} if input_spk_embed else {})
|
||||
},
|
||||
opset_version=17,
|
||||
**onnx_helper.TORCHSCRIPT_EXPORT_KWARGS
|
||||
)
|
||||
|
||||
# Prepare inputs for backbone tracing and multi-variance predictor scripting
|
||||
shape = (1, len(self.model.variance_prediction_list), repeat_bins, 15)
|
||||
noise = torch.randn(shape, device=self.device)
|
||||
condition = torch.rand((1, hparams['hidden_size'], 15), device=self.device)
|
||||
step = (torch.rand((1,), device=self.device) * hparams.get('time_scale_factor', hparams['K_step']))
|
||||
|
||||
print(f'Tracing {self.variance_backbone_class_name} backbone...')
|
||||
multi_var_predictor = self.model.view_as_variance_predictor()
|
||||
multi_var_predictor.variance_predictor.set_backbone(
|
||||
torch.jit.trace(
|
||||
multi_var_predictor.variance_predictor.backbone,
|
||||
(
|
||||
noise,
|
||||
step,
|
||||
condition
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
print(f'Scripting {self.multi_var_predictor_class_name}...')
|
||||
multi_var_predictor = torch.jit.script(
|
||||
multi_var_predictor,
|
||||
example_inputs=[
|
||||
(
|
||||
condition.transpose(1, 2),
|
||||
1 # p_sample branch
|
||||
),
|
||||
(
|
||||
condition.transpose(1, 2),
|
||||
dummy_steps # p_sample_plms branch
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
print(f'Exporting {self.multi_var_predictor_class_name}...')
|
||||
torch.onnx.export(
|
||||
multi_var_predictor,
|
||||
(
|
||||
condition.transpose(1, 2),
|
||||
dummy_steps
|
||||
),
|
||||
self.multi_var_predictor_cache_path,
|
||||
input_names=[
|
||||
'variance_cond',
|
||||
'steps'
|
||||
],
|
||||
output_names=[
|
||||
'xs_pred'
|
||||
],
|
||||
dynamic_axes={
|
||||
'variance_cond': {
|
||||
1: 'n_frames'
|
||||
},
|
||||
'xs_pred': {
|
||||
(1 if len(self.model.variance_prediction_list) == 1 else 2): 'n_frames'
|
||||
}
|
||||
},
|
||||
opset_version=17,
|
||||
**onnx_helper.TORCHSCRIPT_EXPORT_KWARGS
|
||||
)
|
||||
|
||||
# Prepare inputs for postprocessor of the multi-variance predictor
|
||||
xs_shape = (1, 15) \
|
||||
if len(self.model.variance_prediction_list) == 1 \
|
||||
else (1, len(self.model.variance_prediction_list), 15)
|
||||
xs_pred = torch.randn(xs_shape, dtype=torch.float32, device=self.device)
|
||||
torch.onnx.export(
|
||||
self.model.view_as_variance_postprocess(),
|
||||
(
|
||||
xs_pred
|
||||
),
|
||||
self.variance_postprocess_cache_path,
|
||||
input_names=[
|
||||
'xs_pred'
|
||||
],
|
||||
output_names=[
|
||||
f'{v_name}_pred'
|
||||
for v_name in self.model.variance_prediction_list
|
||||
],
|
||||
dynamic_axes={
|
||||
'xs_pred': {
|
||||
(1 if len(self.model.variance_prediction_list) == 1 else 2): 'n_frames'
|
||||
},
|
||||
**{
|
||||
f'{v_name}_pred': {
|
||||
1: 'n_frames'
|
||||
}
|
||||
for v_name in self.model.variance_prediction_list
|
||||
}
|
||||
},
|
||||
opset_version=17,
|
||||
**onnx_helper.TORCHSCRIPT_EXPORT_KWARGS
|
||||
)
|
||||
|
||||
@torch.no_grad()
|
||||
def _perform_spk_mix(self, spk_mix: Dict[str, float]):
|
||||
spk_mix_ids = []
|
||||
spk_mix_values = []
|
||||
for name, value in spk_mix.items():
|
||||
spk_mix_ids.append(self.spk_map[name])
|
||||
assert value >= 0., f'Speaker mix checks failed.\n' \
|
||||
f'Proportion of speaker \'{name}\' is negative.'
|
||||
spk_mix_values.append(value)
|
||||
spk_mix_id_N = torch.LongTensor(spk_mix_ids).to(self.device)[None] # => [1, N]
|
||||
spk_mix_value_N = torch.FloatTensor(spk_mix_values).to(self.device)[None] # => [1, N]
|
||||
spk_mix_value_sum = spk_mix_value_N.sum()
|
||||
assert spk_mix_value_sum > 0., 'Speaker mix checks failed.\n' \
|
||||
'Proportions of speaker mix sum to zero.'
|
||||
spk_mix_value_N /= spk_mix_value_sum # normalize
|
||||
spk_mix_embed = torch.sum(
|
||||
self.model.spk_embed(spk_mix_id_N) * spk_mix_value_N.unsqueeze(2), # => [1, N, H]
|
||||
dim=1, keepdim=True
|
||||
) # => [1, 1, H]
|
||||
return spk_mix_embed
|
||||
|
||||
def _optimize_linguistic_graph(self, linguistic: onnx.ModelProto) -> onnx.ModelProto:
|
||||
onnx_helper.model_override_io_shapes(
|
||||
linguistic,
|
||||
output_shapes={
|
||||
'encoder_out': (1, 'n_tokens', hparams['hidden_size'])
|
||||
}
|
||||
)
|
||||
print(f'Running ONNX Simplifier on {self.fs2_class_name}...')
|
||||
linguistic, check = onnxsim.simplify(linguistic, include_subgraph=True)
|
||||
assert check, 'Simplified ONNX model could not be validated'
|
||||
onnx_helper.model_reorder_io_list(
|
||||
linguistic, 'input',
|
||||
target_name='languages', insert_after_name='tokens'
|
||||
)
|
||||
print(f'| optimize graph: {self.fs2_class_name}')
|
||||
return linguistic
|
||||
|
||||
def _optimize_dur_predictor_graph(self, dur_predictor: onnx.ModelProto) -> onnx.ModelProto:
|
||||
onnx_helper.model_override_io_shapes(
|
||||
dur_predictor,
|
||||
output_shapes={
|
||||
'ph_dur_pred': (1, 'n_tokens')
|
||||
}
|
||||
)
|
||||
print(f'Running ONNX Simplifier on {self.dur_predictor_class_name}...')
|
||||
dur_predictor, check = onnxsim.simplify(dur_predictor, include_subgraph=True)
|
||||
assert check, 'Simplified ONNX model could not be validated'
|
||||
print(f'| optimize graph: {self.dur_predictor_class_name}')
|
||||
return dur_predictor
|
||||
|
||||
def _optimize_merge_pitch_predictor_graph(
|
||||
self, pitch_pre: onnx.ModelProto, pitch_predictor: onnx.ModelProto, pitch_post: onnx.ModelProto
|
||||
) -> onnx.ModelProto:
|
||||
onnx_helper.model_override_io_shapes(
|
||||
pitch_pre, output_shapes={'pitch_cond': (1, 'n_frames', hparams['hidden_size'])}
|
||||
)
|
||||
pitch_pre, check = onnxsim.simplify(pitch_pre, include_subgraph=True)
|
||||
assert check, 'Simplified ONNX model could not be validated'
|
||||
|
||||
onnx_helper.model_override_io_shapes(
|
||||
pitch_predictor, output_shapes={'pitch_pred': (1, 'n_frames')}
|
||||
)
|
||||
print(f'Running ONNX Simplifier #1 on {self.pitch_predictor_class_name}...')
|
||||
pitch_predictor, check = onnxsim.simplify(pitch_predictor, include_subgraph=True)
|
||||
assert check, 'Simplified ONNX model could not be validated'
|
||||
onnx_helper.graph_fold_back_to_squeeze(pitch_predictor.graph)
|
||||
onnx_helper.graph_extract_conditioner_projections(
|
||||
graph=pitch_predictor.graph, op_type='Conv',
|
||||
weight_pattern=r'pitch_predictor\..*\.conditioner_projection\.weight',
|
||||
alias_prefix='/pitch_predictor/backbone/cache'
|
||||
)
|
||||
onnx_helper.graph_remove_unused_values(pitch_predictor.graph)
|
||||
print(f'Running ONNX Simplifier #2 on {self.pitch_predictor_class_name}...')
|
||||
pitch_predictor, check = onnxsim.simplify(pitch_predictor, include_subgraph=True)
|
||||
assert check, 'Simplified ONNX model could not be validated'
|
||||
|
||||
onnx_helper.model_add_prefixes(pitch_pre, node_prefix='/pre', ignored_pattern=r'.*embed.*')
|
||||
onnx_helper.model_add_prefixes(pitch_pre, dim_prefix='pre.', ignored_pattern='(n_tokens)|(n_notes)|(n_frames)')
|
||||
onnx_helper.model_add_prefixes(pitch_post, node_prefix='/post', ignored_pattern=None)
|
||||
onnx_helper.model_add_prefixes(pitch_post, dim_prefix='post.', ignored_pattern='n_frames')
|
||||
pitch_pre_diffusion = onnx.compose.merge_models(
|
||||
pitch_pre, pitch_predictor, io_map=[('pitch_cond', 'pitch_cond')],
|
||||
prefix1='', prefix2='', doc_string='',
|
||||
producer_name=pitch_pre.producer_name, producer_version=pitch_pre.producer_version,
|
||||
domain=pitch_pre.domain, model_version=pitch_pre.model_version
|
||||
)
|
||||
pitch_pre_diffusion.graph.name = pitch_pre.graph.name
|
||||
pitch_predictor = onnx.compose.merge_models(
|
||||
pitch_pre_diffusion, pitch_post, io_map=[
|
||||
('x_pred', 'x_pred'), ('base_pitch', 'base_pitch')
|
||||
], prefix1='', prefix2='', doc_string='',
|
||||
producer_name=pitch_pre.producer_name, producer_version=pitch_pre.producer_version,
|
||||
domain=pitch_pre.domain, model_version=pitch_pre.model_version
|
||||
)
|
||||
pitch_predictor.graph.name = pitch_pre.graph.name
|
||||
|
||||
print(f'| optimize graph: {self.pitch_predictor_class_name}')
|
||||
return pitch_predictor
|
||||
|
||||
def _optimize_merge_variance_predictor_graph(
|
||||
self, var_pre: onnx.ModelProto, var_diffusion: onnx.ModelProto, var_post: onnx.ModelProto
|
||||
):
|
||||
onnx_helper.model_override_io_shapes(
|
||||
var_pre, output_shapes={'variance_cond': (1, 'n_frames', hparams['hidden_size'])}
|
||||
)
|
||||
var_pre, check = onnxsim.simplify(var_pre, include_subgraph=True)
|
||||
assert check, 'Simplified ONNX model could not be validated'
|
||||
|
||||
onnx_helper.model_override_io_shapes(
|
||||
var_diffusion, output_shapes={
|
||||
'xs_pred': (1, 'n_frames')
|
||||
if len(self.model.variance_prediction_list) == 1
|
||||
else (1, len(self.model.variance_prediction_list), 'n_frames')
|
||||
}
|
||||
)
|
||||
print(f'Running ONNX Simplifier #1 on {self.multi_var_predictor_class_name}...')
|
||||
var_diffusion, check = onnxsim.simplify(var_diffusion, include_subgraph=True)
|
||||
assert check, 'Simplified ONNX model could not be validated'
|
||||
onnx_helper.graph_fold_back_to_squeeze(var_diffusion.graph)
|
||||
onnx_helper.graph_extract_conditioner_projections(
|
||||
graph=var_diffusion.graph, op_type='Conv',
|
||||
weight_pattern=r'variance_predictor\..*\.conditioner_projection\.weight',
|
||||
alias_prefix='/variance_predictor/backbone/cache'
|
||||
)
|
||||
onnx_helper.graph_remove_unused_values(var_diffusion.graph)
|
||||
print(f'Running ONNX Simplifier #2 on {self.multi_var_predictor_class_name}...')
|
||||
var_diffusion, check = onnxsim.simplify(var_diffusion, include_subgraph=True)
|
||||
assert check, 'Simplified ONNX model could not be validated'
|
||||
|
||||
var_post, check = onnxsim.simplify(var_post, include_subgraph=True)
|
||||
assert check, 'Simplified ONNX model could not be validated'
|
||||
|
||||
ignored_variance_names = '|'.join([f'({v_name})' for v_name in self.model.variance_prediction_list])
|
||||
onnx_helper.model_add_prefixes(
|
||||
var_pre, node_prefix='/pre', value_info_prefix='/pre', initializer_prefix='/pre',
|
||||
ignored_pattern=fr'.*((embed)|{ignored_variance_names}).*'
|
||||
)
|
||||
onnx_helper.model_add_prefixes(var_pre, dim_prefix='pre.', ignored_pattern='(n_tokens)|(n_frames)')
|
||||
onnx_helper.model_add_prefixes(
|
||||
var_post, node_prefix='/post', value_info_prefix='/post', initializer_prefix='/post',
|
||||
ignored_pattern=None
|
||||
)
|
||||
onnx_helper.model_add_prefixes(var_post, dim_prefix='post.', ignored_pattern='n_frames')
|
||||
|
||||
print(f'Merging {self.multi_var_predictor_class_name} subroutines...')
|
||||
var_pre_diffusion = onnx.compose.merge_models(
|
||||
var_pre, var_diffusion, io_map=[('variance_cond', 'variance_cond')],
|
||||
prefix1='', prefix2='', doc_string='',
|
||||
producer_name=var_pre.producer_name, producer_version=var_pre.producer_version,
|
||||
domain=var_pre.domain, model_version=var_pre.model_version
|
||||
)
|
||||
var_pre_diffusion.graph.name = var_pre.graph.name
|
||||
var_predictor = onnx.compose.merge_models(
|
||||
var_pre_diffusion, var_post, io_map=[('xs_pred', 'xs_pred')],
|
||||
prefix1='', prefix2='', doc_string='',
|
||||
producer_name=var_pre.producer_name, producer_version=var_pre.producer_version,
|
||||
domain=var_pre.domain, model_version=var_pre.model_version
|
||||
)
|
||||
var_predictor.graph.name = var_pre.graph.name
|
||||
return var_predictor
|
||||
|
||||
# noinspection PyMethodMayBeStatic
|
||||
def _export_spk_embed(self, path: Path, spk_embed: torch.Tensor):
|
||||
with open(path, 'wb') as f:
|
||||
f.write(spk_embed.cpu().numpy().tobytes())
|
||||
print(f'| export spk embed => {path}')
|
||||
|
||||
def _export_phonemes(self, path: Path):
|
||||
ph_path = path / f'{self.model_name}.phonemes.json'
|
||||
self.phoneme_dictionary.dump(ph_path)
|
||||
print(f'| export phonemes => {ph_path}')
|
||||
lang_path = path / f'{self.model_name}.languages.json'
|
||||
with open(lang_path, 'w', encoding='utf8') as fw:
|
||||
json.dump(self.lang_map, fw, ensure_ascii=False, indent=2)
|
||||
print(f'| export languages => {lang_path}')
|
||||
@@ -0,0 +1,220 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import List, Tuple
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
|
||||
from modules.core import (
|
||||
GaussianDiffusion, PitchDiffusion, MultiVarianceDiffusion
|
||||
)
|
||||
|
||||
|
||||
def extract(a, t):
|
||||
return a[t].reshape((1, 1, 1, 1))
|
||||
|
||||
|
||||
# noinspection PyMethodOverriding
|
||||
class GaussianDiffusionONNX(GaussianDiffusion):
|
||||
@property
|
||||
def backbone(self):
|
||||
return self.denoise_fn
|
||||
|
||||
# We give up the setter for the property `backbone` because this will cause TorchScript to fail
|
||||
# @backbone.setter
|
||||
@torch.jit.unused
|
||||
def set_backbone(self, value):
|
||||
self.denoise_fn = value
|
||||
|
||||
def q_sample(self, x_start, t, noise):
|
||||
return (
|
||||
extract(self.sqrt_alphas_cumprod, t) * x_start +
|
||||
extract(self.sqrt_one_minus_alphas_cumprod, t) * noise
|
||||
)
|
||||
|
||||
def p_sample(self, x, t, cond):
|
||||
x_pred = self.denoise_fn(x, t, cond)
|
||||
x_recon = (
|
||||
extract(self.sqrt_recip_alphas_cumprod, t) * x -
|
||||
extract(self.sqrt_recipm1_alphas_cumprod, t) * x_pred
|
||||
)
|
||||
# This is previously inherited from original DiffSinger repository
|
||||
# and disabled due to some loudness issues when speedup = 1.
|
||||
# x_recon = torch.clamp(x_recon, min=-1., max=1.)
|
||||
|
||||
model_mean = (
|
||||
extract(self.posterior_mean_coef1, t) * x_recon +
|
||||
extract(self.posterior_mean_coef2, t) * x
|
||||
)
|
||||
model_log_variance = extract(self.posterior_log_variance_clipped, t)
|
||||
noise = torch.randn_like(x)
|
||||
# no noise when t == 0
|
||||
nonzero_mask = ((t > 0).float()).reshape(1, 1, 1, 1)
|
||||
return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise
|
||||
|
||||
def p_sample_ddim(self, x, t, interval: int, cond):
|
||||
a_t = extract(self.alphas_cumprod, t)
|
||||
t_prev = t - interval
|
||||
a_prev = extract(self.alphas_cumprod, t_prev * (t_prev > 0))
|
||||
|
||||
noise_pred = self.denoise_fn(x, t, cond=cond)
|
||||
x_prev = a_prev.sqrt() * (
|
||||
x / a_t.sqrt() + (((1 - a_prev) / a_prev).sqrt() - ((1 - a_t) / a_t).sqrt()) * noise_pred
|
||||
)
|
||||
return x_prev
|
||||
|
||||
def plms_get_x_pred(self, x, noise_t, t, t_prev):
|
||||
a_t = extract(self.alphas_cumprod, t)
|
||||
a_prev = extract(self.alphas_cumprod, t_prev)
|
||||
a_t_sq, a_prev_sq = a_t.sqrt(), a_prev.sqrt()
|
||||
|
||||
x_delta = (a_prev - a_t) * ((1 / (a_t_sq * (a_t_sq + a_prev_sq))) * x - 1 / (
|
||||
a_t_sq * (((1 - a_prev) * a_t).sqrt() + ((1 - a_t) * a_prev).sqrt())) * noise_t)
|
||||
x_pred = x + x_delta
|
||||
|
||||
return x_pred
|
||||
|
||||
def p_sample_plms(self, x_prev, t, interval: int, cond, noise_list: List[Tensor], stage: int):
|
||||
noise_pred = self.denoise_fn(x_prev, t, cond)
|
||||
t_prev = t - interval
|
||||
t_prev = t_prev * (t_prev > 0)
|
||||
if stage == 0:
|
||||
x_pred = self.plms_get_x_pred(x_prev, noise_pred, t, t_prev)
|
||||
noise_pred_prev = self.denoise_fn(x_pred, t_prev, cond)
|
||||
noise_pred_prime = (noise_pred + noise_pred_prev) / 2.
|
||||
elif stage == 1:
|
||||
noise_pred_prime = (3. * noise_pred - noise_list[-1]) / 2.
|
||||
elif stage == 2:
|
||||
noise_pred_prime = (23. * noise_pred - 16. * noise_list[-1] + 5. * noise_list[-2]) / 12.
|
||||
else:
|
||||
noise_pred_prime = (55. * noise_pred - 59. * noise_list[-1] + 37.
|
||||
* noise_list[-2] - 9. * noise_list[-3]) / 24.
|
||||
x_prev = self.plms_get_x_pred(x_prev, noise_pred_prime, t, t_prev)
|
||||
return noise_pred, x_prev
|
||||
|
||||
def norm_spec(self, x):
|
||||
k = (self.spec_max - self.spec_min) / 2.
|
||||
b = (self.spec_max + self.spec_min) / 2.
|
||||
return (x - b) / k
|
||||
|
||||
def denorm_spec(self, x):
|
||||
k = (self.spec_max - self.spec_min) / 2.
|
||||
b = (self.spec_max + self.spec_min) / 2.
|
||||
return x * k + b
|
||||
|
||||
def forward(self, condition, x_start=None, depth=None, steps: int = 10):
|
||||
condition = condition.transpose(1, 2) # [1, T, H] => [1, H, T]
|
||||
device = condition.device
|
||||
n_frames = condition.shape[2]
|
||||
|
||||
noise = torch.randn((1, self.num_feats, self.out_dims, n_frames), device=device)
|
||||
if x_start is None:
|
||||
speedup = max(1, self.timesteps // steps)
|
||||
speedup = self.timestep_factors[torch.sum(self.timestep_factors <= speedup) - 1]
|
||||
step_range = torch.arange(0, self.k_step, speedup, dtype=torch.long, device=device).flip(0)[:, None]
|
||||
x = noise
|
||||
else:
|
||||
depth_int64 = min(torch.round(depth * self.timesteps).long(), self.k_step)
|
||||
speedup = max(1, depth_int64 // steps)
|
||||
depth_int64 = depth_int64 // speedup * speedup # make depth_int64 a multiple of speedup
|
||||
step_range = torch.arange(0, depth_int64, speedup, dtype=torch.long, device=device).flip(0)[:, None]
|
||||
x_start = self.norm_spec(x_start).transpose(-2, -1)
|
||||
if self.num_feats == 1:
|
||||
x_start = x_start[:, None, :, :]
|
||||
if depth_int64 >= self.timesteps:
|
||||
x = noise
|
||||
elif depth_int64 > 0:
|
||||
x = self.q_sample(
|
||||
x_start, torch.full((1,), depth_int64 - 1, device=device, dtype=torch.long), noise
|
||||
)
|
||||
else:
|
||||
x = x_start
|
||||
|
||||
if speedup > 1:
|
||||
for t in step_range:
|
||||
x = self.p_sample_ddim(x, t, interval=speedup, cond=condition)
|
||||
# plms_noise_stage: int = 0
|
||||
# noise_list: List[Tensor] = []
|
||||
# for t in step_range:
|
||||
# noise_pred, x = self.p_sample_plms(
|
||||
# x, t, interval=speedup, cond=condition,
|
||||
# noise_list=noise_list, stage=plms_noise_stage
|
||||
# )
|
||||
# if plms_noise_stage == 0:
|
||||
# noise_list = [noise_pred]
|
||||
# plms_noise_stage = plms_noise_stage + 1
|
||||
# else:
|
||||
# if plms_noise_stage >= 3:
|
||||
# noise_list.pop(0)
|
||||
# else:
|
||||
# plms_noise_stage = plms_noise_stage + 1
|
||||
# noise_list.append(noise_pred)
|
||||
else:
|
||||
for t in step_range:
|
||||
x = self.p_sample(x, t, cond=condition)
|
||||
|
||||
if self.num_feats == 1:
|
||||
x = x.squeeze(1).permute(0, 2, 1) # [B, 1, M, T] => [B, T, M]
|
||||
else:
|
||||
x = x.permute(0, 1, 3, 2) # [B, F, M, T] => [B, F, T, M]
|
||||
x = self.denorm_spec(x)
|
||||
return x
|
||||
|
||||
|
||||
class PitchDiffusionONNX(GaussianDiffusionONNX, PitchDiffusion):
|
||||
def __init__(self, vmin: float, vmax: float,
|
||||
cmin: float, cmax: float, repeat_bins,
|
||||
timesteps=1000, k_step=1000,
|
||||
backbone_type=None, backbone_args=None,
|
||||
betas=None):
|
||||
self.vmin = vmin
|
||||
self.vmax = vmax
|
||||
self.cmin = cmin
|
||||
self.cmax = cmax
|
||||
super(PitchDiffusion, self).__init__(
|
||||
vmin=vmin, vmax=vmax, repeat_bins=repeat_bins,
|
||||
timesteps=timesteps, k_step=k_step,
|
||||
backbone_type=backbone_type, backbone_args=backbone_args,
|
||||
betas=betas
|
||||
)
|
||||
|
||||
def clamp_spec(self, x):
|
||||
return x.clamp(min=self.cmin, max=self.cmax)
|
||||
|
||||
def denorm_spec(self, x):
|
||||
d = (self.spec_max - self.spec_min) / 2.
|
||||
m = (self.spec_max + self.spec_min) / 2.
|
||||
x = x * d + m
|
||||
x = x.mean(dim=-1)
|
||||
return x
|
||||
|
||||
|
||||
class MultiVarianceDiffusionONNX(GaussianDiffusionONNX, MultiVarianceDiffusion):
|
||||
def __init__(
|
||||
self, ranges: List[Tuple[float, float]],
|
||||
clamps: List[Tuple[float | None, float | None] | None],
|
||||
repeat_bins, timesteps=1000, k_step=1000,
|
||||
backbone_type=None, backbone_args=None,
|
||||
betas=None
|
||||
):
|
||||
assert len(ranges) == len(clamps)
|
||||
self.clamps = clamps
|
||||
vmin = [r[0] for r in ranges]
|
||||
vmax = [r[1] for r in ranges]
|
||||
if len(vmin) == 1:
|
||||
vmin = vmin[0]
|
||||
if len(vmax) == 1:
|
||||
vmax = vmax[0]
|
||||
super(MultiVarianceDiffusion, self).__init__(
|
||||
vmin=vmin, vmax=vmax, repeat_bins=repeat_bins,
|
||||
timesteps=timesteps, k_step=k_step,
|
||||
backbone_type=backbone_type, backbone_args=backbone_args,
|
||||
betas=betas
|
||||
)
|
||||
|
||||
def denorm_spec(self, x):
|
||||
d = (self.spec_max - self.spec_min) / 2.
|
||||
m = (self.spec_max + self.spec_min) / 2.
|
||||
x = x * d + m
|
||||
x = x.mean(dim=-1)
|
||||
return x
|
||||
@@ -0,0 +1,235 @@
|
||||
import copy
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from modules.commons.common_layers import NormalInitEmbedding as Embedding
|
||||
from modules.fastspeech.acoustic_encoder import FastSpeech2Acoustic
|
||||
from modules.fastspeech.variance_encoder import FastSpeech2Variance
|
||||
from utils.hparams import hparams
|
||||
from utils.phoneme_utils import PAD_INDEX
|
||||
|
||||
f0_bin = 256
|
||||
f0_max = 1100.0
|
||||
f0_min = 50.0
|
||||
f0_mel_min = 1127 * np.log(1 + f0_min / 700)
|
||||
f0_mel_max = 1127 * np.log(1 + f0_max / 700)
|
||||
|
||||
|
||||
def uniform_attention_pooling(spk_embed, durations):
|
||||
_, T_mel, _ = spk_embed.shape
|
||||
ph_starts = torch.cumsum(torch.cat([torch.zeros_like(durations[:, :1]), durations[:, :-1]], dim=1), dim=1)
|
||||
ph_ends = ph_starts + durations
|
||||
mel_indices = torch.arange(T_mel, device=spk_embed.device).view(1, 1, T_mel)
|
||||
phoneme_to_mel_mask = (mel_indices >= ph_starts.unsqueeze(-1)) & (mel_indices < ph_ends.unsqueeze(-1))
|
||||
uniform_scores = phoneme_to_mel_mask.float()
|
||||
sum_scores = uniform_scores.sum(dim=2, keepdim=True)
|
||||
attn_weights = uniform_scores / (sum_scores + (sum_scores == 0).float()) # [B, T_ph, T_mel]
|
||||
ph_spk_embed = torch.bmm(attn_weights, spk_embed)
|
||||
|
||||
return ph_spk_embed
|
||||
|
||||
|
||||
def f0_to_coarse(f0):
|
||||
f0_mel = 1127 * (1 + f0 / 700).log()
|
||||
a = (f0_bin - 2) / (f0_mel_max - f0_mel_min)
|
||||
b = f0_mel_min * a - 1.
|
||||
f0_mel = torch.where(f0_mel > 0, f0_mel * a - b, f0_mel)
|
||||
torch.clip_(f0_mel, min=1., max=float(f0_bin - 1))
|
||||
f0_coarse = torch.round(f0_mel).long()
|
||||
return f0_coarse
|
||||
|
||||
|
||||
class LengthRegulator(nn.Module):
|
||||
# noinspection PyMethodMayBeStatic
|
||||
def forward(self, dur):
|
||||
token_idx = torch.arange(1, dur.shape[1] + 1, device=dur.device)[None, :, None]
|
||||
dur_cumsum = torch.cumsum(dur, dim=1)
|
||||
dur_cumsum_prev = F.pad(dur_cumsum, (1, -1), mode='constant', value=0)
|
||||
pos_idx = torch.arange(dur.sum(dim=1).max(), device=dur.device)[None, None]
|
||||
token_mask = (pos_idx >= dur_cumsum_prev[:, :, None]) & (pos_idx < dur_cumsum[:, :, None])
|
||||
mel2ph = (token_idx * token_mask).sum(dim=1)
|
||||
return mel2ph
|
||||
|
||||
|
||||
class FastSpeech2AcousticONNX(FastSpeech2Acoustic):
|
||||
def __init__(self, vocab_size, cross_lingual_token_idx=None):
|
||||
super().__init__(vocab_size=vocab_size)
|
||||
self.register_buffer(
|
||||
'cross_lingual_token_idx',
|
||||
torch.LongTensor(cross_lingual_token_idx),
|
||||
persistent=False
|
||||
) # [N,]
|
||||
if len(cross_lingual_token_idx) == 0:
|
||||
self.use_lang_id = False
|
||||
|
||||
# for temporary compatibility; will be completely removed in the future
|
||||
self.f0_embed_type = hparams.get('f0_embed_type', 'continuous')
|
||||
if self.f0_embed_type == 'discrete':
|
||||
self.pitch_embed = Embedding(300, hparams['hidden_size'], PAD_INDEX)
|
||||
|
||||
self.lr = LengthRegulator()
|
||||
if hparams['use_key_shift_embed']:
|
||||
self.shift_min, self.shift_max = hparams['augmentation_args']['random_pitch_shifting']['range']
|
||||
if hparams['use_speed_embed']:
|
||||
self.speed_min, self.speed_max = hparams['augmentation_args']['random_time_stretching']['range']
|
||||
|
||||
# noinspection PyMethodOverriding
|
||||
def forward(
|
||||
self, tokens, durations,
|
||||
f0, variances: dict,
|
||||
gender=None, velocity=None,
|
||||
spk_embed=None,
|
||||
languages=None
|
||||
):
|
||||
txt_embed = self.txt_embed(tokens)
|
||||
durations = durations * (tokens > 0)
|
||||
mel2ph = self.lr(durations)
|
||||
_mel2ph = mel2ph
|
||||
f0 = f0 * (mel2ph > 0)
|
||||
mel2ph = mel2ph[..., None].repeat((1, 1, hparams['hidden_size']))
|
||||
if self.use_variance_scaling:
|
||||
dur_embed = self.dur_embed(torch.log(1 + durations.float())[:, :, None])
|
||||
else:
|
||||
dur_embed = self.dur_embed(durations.float()[:, :, None])
|
||||
if self.use_lang_id:
|
||||
lang_mask = torch.any(
|
||||
tokens[..., None] == self.cross_lingual_token_idx[None, None],
|
||||
dim=-1
|
||||
)
|
||||
lang_embed = self.lang_embed(languages * lang_mask)
|
||||
extra_embed = dur_embed + lang_embed
|
||||
else:
|
||||
extra_embed = dur_embed
|
||||
if hparams.get('use_mix_ln', False):
|
||||
if hasattr(self, 'frozen_spk_embed'):
|
||||
ph_spk_embed = self.frozen_spk_embed.repeat(1, tokens.shape[1], 1)
|
||||
else:
|
||||
ph_spk_embed = uniform_attention_pooling(spk_embed, durations)
|
||||
else:
|
||||
ph_spk_embed = None
|
||||
encoded = self.encoder(txt_embed, extra_embed, tokens == PAD_INDEX, spk_embed=ph_spk_embed)
|
||||
encoded = F.pad(encoded, (0, 0, 1, 0))
|
||||
condition = torch.gather(encoded, 1, mel2ph)
|
||||
|
||||
if self.use_stretch_embed:
|
||||
stretch = torch.round(1000 * self.sr(_mel2ph, durations))
|
||||
table = self.stretch_embed(torch.arange(0, 1001, device=stretch.device))
|
||||
stretch_embed = torch.index_select(table, 0, stretch.view(-1).long()).view_as(condition)
|
||||
condition += stretch_embed
|
||||
stretch_embed_rnn_out, _ = self.stretch_embed_rnn(condition)
|
||||
condition += stretch_embed_rnn_out
|
||||
|
||||
if self.f0_embed_type == 'discrete':
|
||||
pitch = f0_to_coarse(f0)
|
||||
pitch_embed = self.pitch_embed(pitch)
|
||||
else:
|
||||
f0_mel = (1 + f0 / 700).log()
|
||||
pitch_embed = self.pitch_embed(f0_mel[:, :, None])
|
||||
condition += pitch_embed
|
||||
|
||||
if self.use_variance_embeds:
|
||||
variance_embeds = torch.stack([
|
||||
self.variance_embeds[v_name](variances[v_name][:, :, None] * self.variance_scaling_factor[v_name])
|
||||
for v_name in self.variance_embed_list
|
||||
], dim=-1).sum(-1)
|
||||
condition += variance_embeds
|
||||
|
||||
if hparams['use_key_shift_embed']:
|
||||
if hasattr(self, 'frozen_key_shift'):
|
||||
key_shift_embed = self.key_shift_embed(self.frozen_key_shift[:, None, None] * self.variance_scaling_factor['key_shift'])
|
||||
else:
|
||||
gender = torch.clip(gender, min=-1., max=1.)
|
||||
gender_mask = (gender < 0.).float()
|
||||
key_shift = gender * ((1. - gender_mask) * self.shift_max + gender_mask * abs(self.shift_min))
|
||||
key_shift_embed = self.key_shift_embed(key_shift[:, :, None] * self.variance_scaling_factor['key_shift'])
|
||||
condition += key_shift_embed
|
||||
|
||||
if hparams['use_speed_embed']:
|
||||
if velocity is not None:
|
||||
velocity = torch.clip(velocity, min=self.speed_min, max=self.speed_max)
|
||||
speed_embed = self.speed_embed(velocity[:, :, None] * self.variance_scaling_factor['speed'])
|
||||
else:
|
||||
speed_embed = self.speed_embed(torch.FloatTensor([1.]).to(condition.device)[:, None, None] * self.variance_scaling_factor['speed'])
|
||||
condition += speed_embed
|
||||
|
||||
if hparams['use_spk_id']:
|
||||
if hasattr(self, 'frozen_spk_embed'):
|
||||
condition += self.frozen_spk_embed
|
||||
else:
|
||||
condition += spk_embed
|
||||
return condition
|
||||
|
||||
|
||||
class FastSpeech2VarianceONNX(FastSpeech2Variance):
|
||||
def __init__(self, vocab_size, cross_lingual_token_idx=None):
|
||||
super().__init__(vocab_size=vocab_size)
|
||||
self.register_buffer(
|
||||
'cross_lingual_token_idx',
|
||||
torch.LongTensor(cross_lingual_token_idx),
|
||||
persistent=False
|
||||
)
|
||||
if len(cross_lingual_token_idx) == 0:
|
||||
self.use_lang_id = False
|
||||
self.lr = LengthRegulator()
|
||||
|
||||
def forward_encoder_word(self, tokens, word_div, word_dur, languages=None):
|
||||
txt_embed = self.txt_embed(tokens)
|
||||
ph2word = self.lr(word_div)
|
||||
onset = ph2word > F.pad(ph2word, [1, -1])
|
||||
onset_embed = self.onset_embed(onset.long())
|
||||
ph_word_dur = torch.gather(F.pad(word_dur, [1, 0]), 1, ph2word)
|
||||
word_dur_embed = self.word_dur_embed(ph_word_dur.float()[:, :, None])
|
||||
extra_embed = onset_embed + word_dur_embed
|
||||
if self.use_lang_id:
|
||||
lang_mask = torch.any(
|
||||
tokens[..., None] == self.cross_lingual_token_idx[None, None],
|
||||
dim=-1
|
||||
)
|
||||
lang_embed = self.lang_embed(languages * lang_mask)
|
||||
extra_embed += lang_embed
|
||||
x_masks = tokens == PAD_INDEX
|
||||
return self.encoder(txt_embed, extra_embed, x_masks), x_masks
|
||||
|
||||
def forward_encoder_phoneme(self, tokens, ph_dur, languages=None):
|
||||
txt_embed = self.txt_embed(tokens)
|
||||
if self.use_variance_scaling:
|
||||
ph_dur_embed = self.ph_dur_embed(torch.log(1 + ph_dur.float())[:, :, None])
|
||||
else:
|
||||
ph_dur_embed = self.ph_dur_embed(ph_dur.float()[:, :, None])
|
||||
if self.use_lang_id:
|
||||
lang_mask = torch.any(
|
||||
tokens[..., None] == self.cross_lingual_token_idx[None, None],
|
||||
dim=-1
|
||||
)
|
||||
lang_embed = self.lang_embed(languages * lang_mask)
|
||||
extra_embed = ph_dur_embed + lang_embed
|
||||
else:
|
||||
extra_embed = ph_dur_embed
|
||||
x_masks = tokens == PAD_INDEX
|
||||
return self.encoder(txt_embed, extra_embed, x_masks), x_masks
|
||||
|
||||
def forward_dur_predictor(self, encoder_out, x_masks, ph_midi, spk_embed=None):
|
||||
midi_embed = self.midi_embed(ph_midi)
|
||||
dur_cond = encoder_out + midi_embed
|
||||
if hparams['use_spk_id'] and spk_embed is not None:
|
||||
dur_cond += spk_embed
|
||||
ph_dur = self.dur_predictor(dur_cond, x_masks=x_masks)
|
||||
return ph_dur
|
||||
|
||||
def view_as_encoder(self):
|
||||
model = copy.deepcopy(self)
|
||||
if self.predict_dur:
|
||||
del model.dur_predictor
|
||||
model.forward = model.forward_encoder_word
|
||||
else:
|
||||
model.forward = model.forward_encoder_phoneme
|
||||
return model
|
||||
|
||||
def view_as_dur_predictor(self):
|
||||
model = copy.deepcopy(self)
|
||||
del model.encoder
|
||||
model.forward = model.forward_dur_predictor
|
||||
return model
|
||||
@@ -0,0 +1,16 @@
|
||||
import torch
|
||||
|
||||
from modules.nsf_hifigan.env import AttrDict
|
||||
from modules.nsf_hifigan.models import Generator
|
||||
|
||||
|
||||
# noinspection SpellCheckingInspection
|
||||
class NSFHiFiGANONNX(torch.nn.Module):
|
||||
def __init__(self, attrs: dict):
|
||||
super().__init__()
|
||||
self.generator = Generator(AttrDict(attrs))
|
||||
|
||||
def forward(self, mel: torch.Tensor, f0: torch.Tensor):
|
||||
mel = mel.transpose(1, 2)
|
||||
wav = self.generator(mel, f0)
|
||||
return wav.squeeze(1)
|
||||
@@ -0,0 +1,123 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import List, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
from modules.core import (
|
||||
RectifiedFlow, PitchRectifiedFlow, MultiVarianceRectifiedFlow
|
||||
)
|
||||
|
||||
|
||||
class RectifiedFlowONNX(RectifiedFlow):
|
||||
@property
|
||||
def backbone(self):
|
||||
return self.velocity_fn
|
||||
|
||||
# We give up the setter for the property `backbone` because this will cause TorchScript to fail
|
||||
# @backbone.setter
|
||||
@torch.jit.unused
|
||||
def set_backbone(self, value):
|
||||
self.velocity_fn = value
|
||||
|
||||
def sample_euler(self, x, t, dt: float, cond):
|
||||
x += self.velocity_fn(x, t * self.time_scale_factor, cond) * dt
|
||||
return x
|
||||
|
||||
def norm_spec(self, x):
|
||||
k = (self.spec_max - self.spec_min) / 2.
|
||||
b = (self.spec_max + self.spec_min) / 2.
|
||||
return (x - b) / k
|
||||
|
||||
def denorm_spec(self, x):
|
||||
k = (self.spec_max - self.spec_min) / 2.
|
||||
b = (self.spec_max + self.spec_min) / 2.
|
||||
return x * k + b
|
||||
|
||||
def forward(self, condition, x_end=None, depth=None, steps: int = 10):
|
||||
condition = condition.transpose(1, 2) # [1, T, H] => [1, H, T]
|
||||
device = condition.device
|
||||
n_frames = condition.shape[2]
|
||||
noise = torch.randn((1, self.num_feats, self.out_dims, n_frames), device=device)
|
||||
if x_end is None:
|
||||
t_start = 0.
|
||||
x = noise
|
||||
else:
|
||||
t_start = torch.max(1 - depth, torch.tensor(self.t_start, dtype=torch.float32, device=device))
|
||||
x_end = self.norm_spec(x_end).transpose(-2, -1)
|
||||
if self.num_feats == 1:
|
||||
x_end = x_end[:, None, :, :]
|
||||
if t_start <= 0.:
|
||||
x = noise
|
||||
elif t_start >= 1.:
|
||||
x = x_end
|
||||
else:
|
||||
x = t_start * x_end + (1 - t_start) * noise
|
||||
|
||||
t_width = 1. - t_start
|
||||
if t_width >= 0.:
|
||||
dt = t_width / max(1, steps)
|
||||
for t in torch.arange(steps, dtype=torch.long, device=device)[:, None].float() * dt + t_start:
|
||||
x = self.sample_euler(x, t, dt, condition)
|
||||
|
||||
if self.num_feats == 1:
|
||||
x = x.squeeze(1).permute(0, 2, 1) # [B, 1, M, T] => [B, T, M]
|
||||
else:
|
||||
x = x.permute(0, 1, 3, 2) # [B, F, M, T] => [B, F, T, M]
|
||||
x = self.denorm_spec(x)
|
||||
return x
|
||||
|
||||
|
||||
class PitchRectifiedFlowONNX(RectifiedFlowONNX, PitchRectifiedFlow):
|
||||
def __init__(self, vmin: float, vmax: float,
|
||||
cmin: float, cmax: float, repeat_bins,
|
||||
time_scale_factor=1000,
|
||||
backbone_type=None, backbone_args=None):
|
||||
self.vmin = vmin
|
||||
self.vmax = vmax
|
||||
self.cmin = cmin
|
||||
self.cmax = cmax
|
||||
super(PitchRectifiedFlow, self).__init__(
|
||||
vmin=vmin, vmax=vmax, repeat_bins=repeat_bins,
|
||||
time_scale_factor=time_scale_factor,
|
||||
backbone_type=backbone_type, backbone_args=backbone_args
|
||||
)
|
||||
|
||||
def clamp_spec(self, x):
|
||||
return x.clamp(min=self.cmin, max=self.cmax)
|
||||
|
||||
def denorm_spec(self, x):
|
||||
d = (self.spec_max - self.spec_min) / 2.
|
||||
m = (self.spec_max + self.spec_min) / 2.
|
||||
x = x * d + m
|
||||
x = x.mean(dim=-1)
|
||||
return x
|
||||
|
||||
|
||||
class MultiVarianceRectifiedFlowONNX(RectifiedFlowONNX, MultiVarianceRectifiedFlow):
|
||||
def __init__(
|
||||
self, ranges: List[Tuple[float, float]],
|
||||
clamps: List[Tuple[float | None, float | None] | None],
|
||||
repeat_bins, time_scale_factor=1000,
|
||||
backbone_type=None, backbone_args=None
|
||||
):
|
||||
assert len(ranges) == len(clamps)
|
||||
self.clamps = clamps
|
||||
vmin = [r[0] for r in ranges]
|
||||
vmax = [r[1] for r in ranges]
|
||||
if len(vmin) == 1:
|
||||
vmin = vmin[0]
|
||||
if len(vmax) == 1:
|
||||
vmax = vmax[0]
|
||||
super(MultiVarianceRectifiedFlow, self).__init__(
|
||||
vmin=vmin, vmax=vmax, repeat_bins=repeat_bins,
|
||||
time_scale_factor=time_scale_factor,
|
||||
backbone_type=backbone_type, backbone_args=backbone_args
|
||||
)
|
||||
|
||||
def denorm_spec(self, x):
|
||||
d = (self.spec_max - self.spec_min) / 2.
|
||||
m = (self.spec_max + self.spec_min) / 2.
|
||||
x = x * d + m
|
||||
x = x.mean(dim=-1)
|
||||
return x
|
||||
@@ -0,0 +1,413 @@
|
||||
import copy
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from torch import Tensor
|
||||
|
||||
from deployment.modules.diffusion import (
|
||||
GaussianDiffusionONNX, PitchDiffusionONNX, MultiVarianceDiffusionONNX
|
||||
)
|
||||
from deployment.modules.rectified_flow import (
|
||||
RectifiedFlowONNX, PitchRectifiedFlowONNX, MultiVarianceRectifiedFlowONNX
|
||||
)
|
||||
from deployment.modules.fastspeech2 import FastSpeech2AcousticONNX, FastSpeech2VarianceONNX
|
||||
from modules.toplevel import DiffSingerAcoustic, DiffSingerVariance
|
||||
from utils.hparams import hparams
|
||||
|
||||
|
||||
class DiffSingerAcousticONNX(DiffSingerAcoustic):
|
||||
def __init__(self, vocab_size, out_dims, cross_lingual_token_idx=None):
|
||||
super().__init__(vocab_size, out_dims)
|
||||
del self.fs2
|
||||
del self.diffusion
|
||||
self.fs2 = FastSpeech2AcousticONNX(
|
||||
vocab_size=vocab_size,
|
||||
cross_lingual_token_idx=cross_lingual_token_idx
|
||||
)
|
||||
if self.diffusion_type == 'ddpm':
|
||||
self.diffusion = GaussianDiffusionONNX(
|
||||
out_dims=out_dims,
|
||||
num_feats=1,
|
||||
timesteps=hparams['timesteps'],
|
||||
k_step=hparams['K_step'],
|
||||
backbone_type=self.backbone_type,
|
||||
backbone_args=self.backbone_args,
|
||||
spec_min=hparams['spec_min'],
|
||||
spec_max=hparams['spec_max']
|
||||
)
|
||||
elif self.diffusion_type == 'reflow':
|
||||
self.diffusion = RectifiedFlowONNX(
|
||||
out_dims=out_dims,
|
||||
num_feats=1,
|
||||
t_start=hparams['T_start'],
|
||||
time_scale_factor=hparams['time_scale_factor'],
|
||||
backbone_type=self.backbone_type,
|
||||
backbone_args=self.backbone_args,
|
||||
spec_min=hparams['spec_min'],
|
||||
spec_max=hparams['spec_max']
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Invalid diffusion type: {self.diffusion_type}")
|
||||
self.mel_base = hparams.get('mel_base', '10')
|
||||
|
||||
def ensure_mel_base(self, mel):
|
||||
if self.mel_base != 'e':
|
||||
# log10 mel to log mel
|
||||
mel = mel * 2.30259
|
||||
return mel
|
||||
|
||||
def forward_fs2_aux(
|
||||
self,
|
||||
tokens: Tensor,
|
||||
durations: Tensor,
|
||||
f0: Tensor,
|
||||
variances: dict,
|
||||
gender: Tensor = None,
|
||||
velocity: Tensor = None,
|
||||
spk_embed: Tensor = None,
|
||||
languages: Tensor = None
|
||||
):
|
||||
condition = self.fs2(
|
||||
tokens, durations, f0, variances=variances,
|
||||
gender=gender, velocity=velocity, spk_embed=spk_embed,
|
||||
languages=languages
|
||||
)
|
||||
if self.use_shallow_diffusion:
|
||||
aux_mel_pred = self.aux_decoder(condition, infer=True)
|
||||
return condition, aux_mel_pred
|
||||
else:
|
||||
return condition
|
||||
|
||||
def forward_shallow_diffusion(
|
||||
self, condition: Tensor, x_start: Tensor,
|
||||
depth, steps: int
|
||||
) -> Tensor:
|
||||
mel_pred = self.diffusion(condition, x_start=x_start, depth=depth, steps=steps)
|
||||
return self.ensure_mel_base(mel_pred)
|
||||
|
||||
def forward_diffusion(self, condition: Tensor, steps: int):
|
||||
mel_pred = self.diffusion(condition, steps=steps)
|
||||
return self.ensure_mel_base(mel_pred)
|
||||
|
||||
def forward_shallow_reflow(
|
||||
self, condition: Tensor, x_end: Tensor,
|
||||
depth, steps: int
|
||||
):
|
||||
mel_pred = self.diffusion(condition, x_end=x_end, depth=depth, steps=steps)
|
||||
return self.ensure_mel_base(mel_pred)
|
||||
|
||||
def forward_reflow(self, condition: Tensor, steps: int):
|
||||
mel_pred = self.diffusion(condition, steps=steps)
|
||||
return self.ensure_mel_base(mel_pred)
|
||||
|
||||
def view_as_fs2_aux(self) -> nn.Module:
|
||||
model = copy.deepcopy(self)
|
||||
del model.diffusion
|
||||
model.forward = model.forward_fs2_aux
|
||||
return model
|
||||
|
||||
def view_as_diffusion(self) -> nn.Module:
|
||||
model = copy.deepcopy(self)
|
||||
del model.fs2
|
||||
if self.use_shallow_diffusion:
|
||||
del model.aux_decoder
|
||||
model.forward = model.forward_shallow_diffusion
|
||||
else:
|
||||
model.forward = model.forward_diffusion
|
||||
return model
|
||||
|
||||
def view_as_reflow(self) -> nn.Module:
|
||||
model = copy.deepcopy(self)
|
||||
del model.fs2
|
||||
if self.use_shallow_diffusion:
|
||||
del model.aux_decoder
|
||||
model.forward = model.forward_shallow_reflow
|
||||
else:
|
||||
model.forward = model.forward_reflow
|
||||
return model
|
||||
|
||||
|
||||
class DiffSingerVarianceONNX(DiffSingerVariance):
|
||||
def __init__(self, vocab_size, cross_lingual_token_idx=None):
|
||||
super().__init__(vocab_size=vocab_size)
|
||||
del self.fs2
|
||||
self.fs2 = FastSpeech2VarianceONNX(
|
||||
vocab_size=vocab_size,
|
||||
cross_lingual_token_idx=cross_lingual_token_idx
|
||||
)
|
||||
self.hidden_size = hparams['hidden_size']
|
||||
if self.predict_pitch:
|
||||
del self.pitch_predictor
|
||||
self.smooth: nn.Conv1d = None
|
||||
pitch_hparams = hparams['pitch_prediction_args']
|
||||
if self.diffusion_type == 'ddpm':
|
||||
self.pitch_predictor = PitchDiffusionONNX(
|
||||
vmin=pitch_hparams['pitd_norm_min'],
|
||||
vmax=pitch_hparams['pitd_norm_max'],
|
||||
cmin=pitch_hparams['pitd_clip_min'],
|
||||
cmax=pitch_hparams['pitd_clip_max'],
|
||||
repeat_bins=pitch_hparams['repeat_bins'],
|
||||
timesteps=hparams['timesteps'],
|
||||
k_step=hparams['K_step'],
|
||||
backbone_type=self.pitch_backbone_type,
|
||||
backbone_args=self.pitch_backbone_args
|
||||
)
|
||||
elif self.diffusion_type == 'reflow':
|
||||
self.pitch_predictor = PitchRectifiedFlowONNX(
|
||||
vmin=pitch_hparams['pitd_norm_min'],
|
||||
vmax=pitch_hparams['pitd_norm_max'],
|
||||
cmin=pitch_hparams['pitd_clip_min'],
|
||||
cmax=pitch_hparams['pitd_clip_max'],
|
||||
repeat_bins=pitch_hparams['repeat_bins'],
|
||||
time_scale_factor=hparams['time_scale_factor'],
|
||||
backbone_type=self.pitch_backbone_type,
|
||||
backbone_args=self.pitch_backbone_args
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Invalid diffusion type: {self.diffusion_type}")
|
||||
if self.predict_variances:
|
||||
del self.variance_predictor
|
||||
if self.diffusion_type == 'ddpm':
|
||||
self.variance_predictor = self.build_adaptor(cls=MultiVarianceDiffusionONNX)
|
||||
elif self.diffusion_type == 'reflow':
|
||||
self.variance_predictor = self.build_adaptor(cls=MultiVarianceRectifiedFlowONNX)
|
||||
else:
|
||||
raise NotImplementedError(self.diffusion_type)
|
||||
|
||||
def build_smooth_op(self, device):
|
||||
smooth_kernel_size = round(hparams['midi_smooth_width'] * hparams['audio_sample_rate'] / hparams['hop_size'])
|
||||
smooth = nn.Conv1d(
|
||||
in_channels=1,
|
||||
out_channels=1,
|
||||
kernel_size=smooth_kernel_size,
|
||||
bias=False,
|
||||
padding='same',
|
||||
padding_mode='replicate'
|
||||
).eval()
|
||||
smooth_kernel = torch.sin(torch.from_numpy(
|
||||
np.linspace(0, 1, smooth_kernel_size).astype(np.float32) * np.pi
|
||||
))
|
||||
smooth_kernel /= smooth_kernel.sum()
|
||||
smooth.weight.data = smooth_kernel[None, None]
|
||||
self.smooth = smooth.to(device)
|
||||
|
||||
def embed_frozen_spk(self, encoder_out):
|
||||
if hparams['use_spk_id'] and hasattr(self, 'frozen_spk_embed'):
|
||||
encoder_out += self.frozen_spk_embed
|
||||
return encoder_out
|
||||
|
||||
def forward_linguistic_encoder_word(self, tokens, word_div, word_dur, languages=None):
|
||||
encoder_out, x_masks = self.fs2.forward_encoder_word(tokens, word_div, word_dur, languages=languages)
|
||||
encoder_out = self.embed_frozen_spk(encoder_out)
|
||||
return encoder_out, x_masks
|
||||
|
||||
def forward_linguistic_encoder_phoneme(self, tokens, ph_dur, languages=None):
|
||||
encoder_out, x_masks = self.fs2.forward_encoder_phoneme(tokens, ph_dur, languages=languages)
|
||||
encoder_out = self.embed_frozen_spk(encoder_out)
|
||||
return encoder_out, x_masks
|
||||
|
||||
def forward_dur_predictor(self, encoder_out, x_masks, ph_midi, spk_embed=None):
|
||||
return self.fs2.forward_dur_predictor(encoder_out, x_masks, ph_midi, spk_embed=spk_embed)
|
||||
|
||||
def forward_mel2x_gather(self, x_src, x_dur, x_dim=None, check_stretch_embed=False):
|
||||
mel2x = self.lr(x_dur)
|
||||
_mel2x = mel2x
|
||||
if x_dim is not None:
|
||||
x_src = F.pad(x_src, [0, 0, 1, 0])
|
||||
mel2x = mel2x[..., None].repeat([1, 1, x_dim])
|
||||
else:
|
||||
x_src = F.pad(x_src, [1, 0])
|
||||
x_cond = torch.gather(x_src, 1, mel2x)
|
||||
if self.use_stretch_embed and check_stretch_embed:
|
||||
stretch = torch.round(1000 * self.sr(_mel2x, x_dur))
|
||||
table = self.stretch_embed(torch.arange(0, 1001, device=stretch.device))
|
||||
stretch_embed = torch.index_select(table, 0, stretch.view(-1).long()).view_as(x_cond)
|
||||
x_cond += stretch_embed
|
||||
stretch_embed_rnn_out, _ = self.stretch_embed_rnn(x_cond)
|
||||
x_cond += stretch_embed_rnn_out
|
||||
return x_cond
|
||||
|
||||
def forward_pitch_preprocess(
|
||||
self, encoder_out, ph_dur,
|
||||
note_midi=None, note_rest=None, note_dur=None, note_glide=None,
|
||||
pitch=None, expr=None, retake=None, spk_embed=None
|
||||
):
|
||||
condition = self.forward_mel2x_gather(encoder_out, ph_dur, x_dim=self.hidden_size, check_stretch_embed=True)
|
||||
if self.use_melody_encoder:
|
||||
if self.melody_encoder.use_glide_embed and note_glide is None:
|
||||
note_glide = torch.LongTensor([[0]]).to(encoder_out.device)
|
||||
melody_encoder_out = self.melody_encoder(
|
||||
note_midi, note_rest, note_dur,
|
||||
glide=note_glide
|
||||
)
|
||||
melody_encoder_out = self.forward_mel2x_gather(melody_encoder_out, note_dur, x_dim=self.hidden_size)
|
||||
condition += melody_encoder_out
|
||||
if expr is None:
|
||||
retake_embed = self.pitch_retake_embed(retake.long())
|
||||
else:
|
||||
retake_true_embed = self.pitch_retake_embed(
|
||||
torch.ones(1, 1, dtype=torch.long, device=encoder_out.device)
|
||||
) # [B=1, T=1] => [B=1, T=1, H]
|
||||
retake_false_embed = self.pitch_retake_embed(
|
||||
torch.zeros(1, 1, dtype=torch.long, device=encoder_out.device)
|
||||
) # [B=1, T=1] => [B=1, T=1, H]
|
||||
expr = (expr * retake)[:, :, None] # [B, T, 1]
|
||||
retake_embed = expr * retake_true_embed + (1. - expr) * retake_false_embed
|
||||
pitch_cond = condition + retake_embed
|
||||
frame_midi_pitch = self.forward_mel2x_gather(note_midi, note_dur, x_dim=None)
|
||||
base_pitch = self.smooth(frame_midi_pitch)
|
||||
if self.use_melody_encoder:
|
||||
delta_pitch = (pitch - base_pitch) * ~retake
|
||||
if self.use_variance_scaling:
|
||||
pitch_cond += self.delta_pitch_embed(delta_pitch[:, :, None] / 12)
|
||||
else:
|
||||
pitch_cond += self.delta_pitch_embed(delta_pitch[:, :, None])
|
||||
else:
|
||||
base_pitch = base_pitch * retake + pitch * ~retake
|
||||
if self.use_variance_scaling:
|
||||
pitch_cond += self.base_pitch_embed(base_pitch[:, :, None] / 128)
|
||||
else:
|
||||
pitch_cond += self.base_pitch_embed(base_pitch[:, :, None])
|
||||
if hparams['use_spk_id'] and spk_embed is not None:
|
||||
pitch_cond += spk_embed
|
||||
return pitch_cond, base_pitch
|
||||
|
||||
def forward_pitch_reflow(
|
||||
self, pitch_cond, steps: int = 10
|
||||
):
|
||||
x_pred = self.pitch_predictor(pitch_cond, steps=steps)
|
||||
return x_pred
|
||||
|
||||
def forward_pitch_postprocess(self, x_pred, base_pitch):
|
||||
pitch_pred = self.pitch_predictor.clamp_spec(x_pred) + base_pitch
|
||||
return pitch_pred
|
||||
|
||||
def forward_variance_preprocess(
|
||||
self, encoder_out, ph_dur, pitch,
|
||||
variances: dict = None, retake=None, spk_embed=None
|
||||
):
|
||||
condition = self.forward_mel2x_gather(encoder_out, ph_dur, x_dim=self.hidden_size, check_stretch_embed=True)
|
||||
if self.use_variance_scaling:
|
||||
variance_cond = condition + self.pitch_embed(pitch[:, :, None] / 12)
|
||||
else:
|
||||
variance_cond = condition + self.pitch_embed(pitch[:, :, None])
|
||||
non_retake_masks = [
|
||||
v_retake.float() # [B, T, 1]
|
||||
for v_retake in (~retake).split(1, dim=2)
|
||||
]
|
||||
variance_embeds = [
|
||||
self.variance_embeds[v_name](variances[v_name][:, :, None] * self.variance_retake_scaling[v_name]) * v_masks
|
||||
for v_name, v_masks in zip(self.variance_prediction_list, non_retake_masks)
|
||||
]
|
||||
variance_cond += torch.stack(variance_embeds, dim=-1).sum(-1)
|
||||
if hparams['use_spk_id'] and spk_embed is not None:
|
||||
variance_cond += spk_embed
|
||||
return variance_cond
|
||||
|
||||
def forward_variance_reflow(self, variance_cond, steps: int = 10):
|
||||
xs_pred = self.variance_predictor(variance_cond, steps=steps)
|
||||
return xs_pred
|
||||
|
||||
def forward_variance_postprocess(self, xs_pred):
|
||||
if self.variance_predictor.num_feats == 1:
|
||||
xs_pred = [xs_pred]
|
||||
else:
|
||||
xs_pred = xs_pred.unbind(dim=1)
|
||||
variance_pred = self.variance_predictor.clamp_spec(xs_pred)
|
||||
return tuple(variance_pred)
|
||||
|
||||
def view_as_linguistic_encoder(self):
|
||||
model = copy.deepcopy(self)
|
||||
if self.predict_pitch:
|
||||
del model.pitch_predictor
|
||||
if self.use_melody_encoder:
|
||||
del model.melody_encoder
|
||||
if self.predict_variances:
|
||||
del model.variance_predictor
|
||||
model.fs2 = model.fs2.view_as_encoder()
|
||||
if self.predict_dur:
|
||||
model.forward = model.forward_linguistic_encoder_word
|
||||
else:
|
||||
model.forward = model.forward_linguistic_encoder_phoneme
|
||||
return model
|
||||
|
||||
def view_as_dur_predictor(self):
|
||||
assert self.predict_dur
|
||||
model = copy.deepcopy(self)
|
||||
if self.predict_pitch:
|
||||
del model.pitch_predictor
|
||||
if self.use_melody_encoder:
|
||||
del model.melody_encoder
|
||||
if self.predict_variances:
|
||||
del model.variance_predictor
|
||||
model.fs2 = model.fs2.view_as_dur_predictor()
|
||||
model.forward = model.forward_dur_predictor
|
||||
return model
|
||||
|
||||
def view_as_pitch_preprocess(self):
|
||||
model = copy.deepcopy(self)
|
||||
del model.fs2
|
||||
if self.predict_pitch:
|
||||
del model.pitch_predictor
|
||||
if self.predict_variances:
|
||||
del model.variance_predictor
|
||||
model.forward = model.forward_pitch_preprocess
|
||||
return model
|
||||
|
||||
def view_as_pitch_predictor(self):
|
||||
assert self.predict_pitch
|
||||
model = copy.deepcopy(self)
|
||||
del model.fs2
|
||||
del model.lr
|
||||
if self.use_melody_encoder:
|
||||
del model.melody_encoder
|
||||
if self.predict_variances:
|
||||
del model.variance_predictor
|
||||
model.forward = model.forward_pitch_reflow
|
||||
return model
|
||||
|
||||
def view_as_pitch_postprocess(self):
|
||||
model = copy.deepcopy(self)
|
||||
del model.fs2
|
||||
if self.use_melody_encoder:
|
||||
del model.melody_encoder
|
||||
if self.predict_variances:
|
||||
del model.variance_predictor
|
||||
model.forward = model.forward_pitch_postprocess
|
||||
return model
|
||||
|
||||
def view_as_variance_preprocess(self):
|
||||
model = copy.deepcopy(self)
|
||||
del model.fs2
|
||||
if self.predict_pitch:
|
||||
del model.pitch_predictor
|
||||
if self.use_melody_encoder:
|
||||
del model.melody_encoder
|
||||
if self.predict_variances:
|
||||
del model.variance_predictor
|
||||
model.forward = model.forward_variance_preprocess
|
||||
return model
|
||||
|
||||
def view_as_variance_predictor(self):
|
||||
assert self.predict_variances
|
||||
model = copy.deepcopy(self)
|
||||
del model.fs2
|
||||
del model.lr
|
||||
if self.predict_pitch:
|
||||
del model.pitch_predictor
|
||||
if self.use_melody_encoder:
|
||||
del model.melody_encoder
|
||||
model.forward = model.forward_variance_reflow
|
||||
return model
|
||||
|
||||
def view_as_variance_postprocess(self):
|
||||
model = copy.deepcopy(self)
|
||||
del model.fs2
|
||||
if self.predict_pitch:
|
||||
del model.pitch_predictor
|
||||
if self.use_melody_encoder:
|
||||
del model.melody_encoder
|
||||
model.forward = model.forward_variance_postprocess
|
||||
return model
|
||||
Reference in New Issue
Block a user