wehub-resource-sync
137 lines
6.4 KiB
Python
137 lines
6.4 KiB
Python
# coding=utf8
|
|
import numpy as np
|
|
import torch
|
|
from torch import Tensor
|
|
from typing import Tuple, Dict
|
|
|
|
from utils.hparams import hparams
|
|
from utils.infer_utils import resample_align_curve
|
|
|
|
|
|
class BaseSVSInfer:
|
|
"""
|
|
Base class for SVS inference models.
|
|
Subclasses should define:
|
|
1. *build_model*:
|
|
how to build the model;
|
|
2. *run_model*:
|
|
how to run the model (typically, generate a mel-spectrogram and
|
|
pass it to the pre-built vocoder);
|
|
3. *preprocess_input*:
|
|
how to preprocess user input.
|
|
4. *infer_once*
|
|
infer from raw inputs to the final outputs
|
|
"""
|
|
|
|
def __init__(self, device=None):
|
|
if device is None:
|
|
device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
|
self.device = device
|
|
self.timestep = hparams['hop_size'] / hparams['audio_sample_rate']
|
|
self.spk_map = {}
|
|
self.lang_map = {}
|
|
self.model: torch.nn.Module = None
|
|
|
|
def build_model(self, ckpt_steps=None) -> torch.nn.Module:
|
|
raise NotImplementedError()
|
|
|
|
def load_speaker_mix(self, param_src: dict, summary_dst: dict,
|
|
mix_mode: str = 'frame', mix_length: int = None) -> Tuple[Tensor, Tensor]:
|
|
"""
|
|
|
|
:param param_src: param dict
|
|
:param summary_dst: summary dict
|
|
:param mix_mode: 'token' or 'frame'
|
|
:param mix_length: total tokens or frames to mix
|
|
:return: spk_mix_id [B=1, 1, N], spk_mix_value [B=1, T, N]
|
|
"""
|
|
assert mix_mode == 'token' or mix_mode == 'frame'
|
|
param_key = 'spk_mix' if mix_mode == 'frame' else 'ph_spk_mix'
|
|
summary_solo_key = 'spk' if mix_mode == 'frame' else 'ph_spk'
|
|
spk_mix_map = param_src.get(param_key) # { spk_name: value } or { spk_name: "value value value ..." }
|
|
dynamic = False
|
|
if spk_mix_map is None:
|
|
assert len(self.spk_map) == 1, (
|
|
"This is a multi-speaker model. "
|
|
"Please specify a speaker or speaker mix by --spk option."
|
|
)
|
|
# Get the only speaker
|
|
for name in self.spk_map.keys():
|
|
spk_mix_map = {name: 1.0}
|
|
break
|
|
else:
|
|
for name in spk_mix_map:
|
|
assert name in self.spk_map, f'Speaker \'{name}\' not found.'
|
|
if len(spk_mix_map) == 1:
|
|
summary_dst[summary_solo_key] = list(spk_mix_map.keys())[0]
|
|
elif any([isinstance(val, str) for val in spk_mix_map.values()]):
|
|
print_mix = '|'.join(spk_mix_map.keys())
|
|
summary_dst[param_key] = f'dynamic({print_mix})'
|
|
dynamic = True
|
|
else:
|
|
print_mix = '|'.join([f'{n}:{"%.3f" % spk_mix_map[n]}' for n in spk_mix_map])
|
|
summary_dst[param_key] = f'static({print_mix})'
|
|
spk_mix_id_list = []
|
|
spk_mix_value_list = []
|
|
if dynamic:
|
|
for name, values in spk_mix_map.items():
|
|
spk_mix_id_list.append(self.spk_map[name])
|
|
if isinstance(values, str):
|
|
# this speaker has a variable proportion
|
|
if mix_mode == 'token':
|
|
cur_spk_mix_value = values.split()
|
|
assert len(cur_spk_mix_value) == mix_length, \
|
|
'Speaker mix checks failed. In dynamic token-level mix, ' \
|
|
'number of proportion values must equal number of tokens.'
|
|
cur_spk_mix_value = torch.from_numpy(
|
|
np.array(cur_spk_mix_value, 'float32')
|
|
).to(self.device)[None] # => [B=1, T]
|
|
else:
|
|
cur_spk_mix_value = torch.from_numpy(resample_align_curve(
|
|
np.array(values.split(), 'float32'),
|
|
original_timestep=float(param_src['spk_mix_timestep']),
|
|
target_timestep=self.timestep,
|
|
align_length=mix_length
|
|
)).to(self.device)[None] # => [B=1, T]
|
|
assert torch.all(cur_spk_mix_value >= 0.), \
|
|
f'Speaker mix checks failed.\n' \
|
|
f'Proportions of speaker \'{name}\' on some {mix_mode}s are negative.'
|
|
else:
|
|
# this speaker has a constant proportion
|
|
assert values >= 0., f'Speaker mix checks failed.\n' \
|
|
f'Proportion of speaker \'{name}\' is negative.'
|
|
cur_spk_mix_value = torch.full(
|
|
(1, mix_length), fill_value=values,
|
|
dtype=torch.float32, device=self.device
|
|
)
|
|
spk_mix_value_list.append(cur_spk_mix_value)
|
|
spk_mix_id = torch.LongTensor(spk_mix_id_list).to(self.device)[None, None] # => [B=1, 1, N]
|
|
spk_mix_value = torch.stack(spk_mix_value_list, dim=2) # [B=1, T] => [B=1, T, N]
|
|
spk_mix_value_sum = torch.sum(spk_mix_value, dim=2, keepdim=True) # => [B=1, T, 1]
|
|
assert torch.all(spk_mix_value_sum > 0.), \
|
|
f'Speaker mix checks failed.\n' \
|
|
f'Proportions of speaker mix on some frames sum to zero.'
|
|
spk_mix_value /= spk_mix_value_sum # normalize
|
|
else:
|
|
for name, value in spk_mix_map.items():
|
|
spk_mix_id_list.append(self.spk_map[name])
|
|
assert value >= 0., f'Speaker mix checks failed.\n' \
|
|
f'Proportion of speaker \'{name}\' is negative.'
|
|
spk_mix_value_list.append(value)
|
|
spk_mix_id = torch.LongTensor(spk_mix_id_list).to(self.device)[None, None] # => [B=1, 1, N]
|
|
spk_mix_value = torch.FloatTensor(spk_mix_value_list).to(self.device)[None, None] # => [B=1, 1, N]
|
|
spk_mix_value_sum = spk_mix_value.sum()
|
|
assert spk_mix_value_sum > 0., f'Speaker mix checks failed.\n' \
|
|
f'Proportions of speaker mix sum to zero.'
|
|
spk_mix_value /= spk_mix_value_sum # normalize
|
|
return spk_mix_id, spk_mix_value
|
|
|
|
def preprocess_input(self, param: dict, idx=0) -> Dict[str, torch.Tensor]:
|
|
raise NotImplementedError()
|
|
|
|
def forward_model(self, sample: Dict[str, torch.Tensor]):
|
|
raise NotImplementedError()
|
|
|
|
def run_inference(self, params, **kwargs):
|
|
raise NotImplementedError()
|