wehub-resource-sync
809 lines
35 KiB
Python
809 lines
35 KiB
Python
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}')
|