wehub-resource-sync
272 lines
12 KiB
Python
272 lines
12 KiB
Python
import json
|
|
import pathlib
|
|
from collections import OrderedDict
|
|
from typing import Dict
|
|
|
|
import numpy as np
|
|
import torch
|
|
import tqdm
|
|
|
|
from basics.base_svs_infer import BaseSVSInfer
|
|
from modules.fastspeech.param_adaptor import VARIANCE_CHECKLIST
|
|
from modules.fastspeech.tts_modules import LengthRegulator
|
|
from modules.toplevel import DiffSingerAcoustic, ShallowDiffusionOutput
|
|
from modules.vocoders.registry import VOCODERS
|
|
from utils import load_ckpt
|
|
from utils.hparams import hparams
|
|
from utils.infer_utils import cross_fade, resample_align_curve, save_wav
|
|
from utils.phoneme_utils import load_phoneme_dictionary
|
|
|
|
|
|
class DiffSingerAcousticInfer(BaseSVSInfer):
|
|
def __init__(self, device=None, load_model=True, load_vocoder=True, ckpt_steps=None):
|
|
super().__init__(device=device)
|
|
if load_model:
|
|
self.variance_checklist = []
|
|
|
|
self.variances_to_embed = set()
|
|
|
|
if hparams.get('use_energy_embed', False):
|
|
self.variances_to_embed.add('energy')
|
|
if hparams.get('use_breathiness_embed', False):
|
|
self.variances_to_embed.add('breathiness')
|
|
if hparams.get('use_voicing_embed', False):
|
|
self.variances_to_embed.add('voicing')
|
|
if hparams.get('use_tension_embed', False):
|
|
self.variances_to_embed.add('tension')
|
|
|
|
self.phoneme_dictionary = load_phoneme_dictionary()
|
|
if hparams['use_spk_id']:
|
|
with open(pathlib.Path(hparams['work_dir']) / 'spk_map.json', 'r', encoding='utf8') as f:
|
|
self.spk_map = json.load(f)
|
|
assert isinstance(self.spk_map, dict) and len(self.spk_map) > 0, 'Invalid or empty speaker map!'
|
|
assert len(self.spk_map) == len(set(self.spk_map.values())), 'Duplicate speaker id in speaker map!'
|
|
lang_map_fn = pathlib.Path(hparams['work_dir']) / 'lang_map.json'
|
|
if lang_map_fn.exists():
|
|
with open(lang_map_fn, 'r', encoding='utf8') as f:
|
|
self.lang_map = json.load(f)
|
|
self.model = self.build_model(ckpt_steps=ckpt_steps)
|
|
self.lr = LengthRegulator().to(self.device)
|
|
if load_vocoder:
|
|
self.vocoder = self.build_vocoder()
|
|
|
|
def build_model(self, ckpt_steps=None):
|
|
model = DiffSingerAcoustic(
|
|
vocab_size=len(self.phoneme_dictionary),
|
|
out_dims=hparams['audio_num_mel_bins']
|
|
).eval().to(self.device)
|
|
load_ckpt(model, hparams['work_dir'], ckpt_steps=ckpt_steps,
|
|
prefix_in_ckpt='model', strict=True, device=self.device)
|
|
return model
|
|
|
|
def build_vocoder(self):
|
|
if hparams['vocoder'] in VOCODERS:
|
|
vocoder = VOCODERS[hparams['vocoder']]()
|
|
else:
|
|
vocoder = VOCODERS[hparams['vocoder'].split('.')[-1]]()
|
|
vocoder.to_device(self.device)
|
|
return vocoder
|
|
|
|
def preprocess_input(self, param, idx=0):
|
|
"""
|
|
:param param: one segment in the .ds file
|
|
:param idx: index of the segment
|
|
:return: batch of the model inputs
|
|
"""
|
|
batch = {}
|
|
summary = OrderedDict()
|
|
|
|
lang = param.get('lang')
|
|
if lang is None:
|
|
assert len(self.lang_map) <= 1, (
|
|
"This is a multilingual model. "
|
|
"Please specify a language by --lang option."
|
|
)
|
|
else:
|
|
assert lang in self.lang_map, f'Unrecognized language name: \'{lang}\'.'
|
|
if hparams.get('use_lang_id', False):
|
|
languages = torch.LongTensor([
|
|
(
|
|
self.lang_map[lang if '/' not in p else p.split('/', maxsplit=1)[0]]
|
|
if self.phoneme_dictionary.is_cross_lingual(p if '/' in p else f'{lang}/{p}')
|
|
else 0
|
|
)
|
|
for p in param['ph_seq'].split()
|
|
]).to(self.device) # => [B, T_txt]
|
|
batch['languages'] = languages
|
|
txt_tokens = torch.LongTensor([
|
|
self.phoneme_dictionary.encode(param['ph_seq'], lang=lang)
|
|
]).to(self.device) # => [B, T_txt]
|
|
batch['tokens'] = txt_tokens
|
|
|
|
ph_dur = torch.from_numpy(np.array(param['ph_dur'].split(), np.float32)).to(self.device)
|
|
ph_acc = torch.round(torch.cumsum(ph_dur, dim=0) / self.timestep + 0.5).long()
|
|
durations = torch.diff(ph_acc, dim=0, prepend=torch.LongTensor([0]).to(self.device))[None] # => [B=1, T_txt]
|
|
mel2ph = self.lr(durations, txt_tokens == 0) # => [B=1, T]
|
|
batch['mel2ph'] = mel2ph
|
|
length = mel2ph.size(1) # => T
|
|
|
|
summary['tokens'] = txt_tokens.size(1)
|
|
summary['frames'] = length
|
|
summary['seconds'] = '%.2f' % (length * self.timestep)
|
|
|
|
if hparams['use_spk_id']:
|
|
spk_mix_id, spk_mix_value = self.load_speaker_mix(
|
|
param_src=param, summary_dst=summary, mix_mode='frame', mix_length=length
|
|
)
|
|
batch['spk_mix_id'] = spk_mix_id
|
|
batch['spk_mix_value'] = spk_mix_value
|
|
|
|
batch['f0'] = torch.from_numpy(resample_align_curve(
|
|
np.array(param['f0_seq'].split(), np.float32),
|
|
original_timestep=float(param['f0_timestep']),
|
|
target_timestep=self.timestep,
|
|
align_length=length
|
|
)).to(self.device)[None]
|
|
|
|
for v_name in VARIANCE_CHECKLIST:
|
|
if v_name in self.variances_to_embed:
|
|
batch[v_name] = torch.from_numpy(resample_align_curve(
|
|
np.array(param[v_name].split(), np.float32),
|
|
original_timestep=float(param[f'{v_name}_timestep']),
|
|
target_timestep=self.timestep,
|
|
align_length=length
|
|
)).to(self.device)[None]
|
|
summary[v_name] = 'manual'
|
|
|
|
if hparams['use_key_shift_embed']:
|
|
shift_min, shift_max = hparams['augmentation_args']['random_pitch_shifting']['range']
|
|
gender = param.get('gender')
|
|
if gender is None:
|
|
gender = 0.
|
|
if isinstance(gender, (int, float, bool)): # static gender value
|
|
summary['gender'] = f'static({gender:.3f})'
|
|
key_shift_value = gender * shift_max if gender >= 0 else gender * abs(shift_min)
|
|
batch['key_shift'] = torch.FloatTensor([key_shift_value]).to(self.device)[:, None] # => [B=1, T=1]
|
|
else:
|
|
summary['gender'] = 'dynamic'
|
|
gender_seq = resample_align_curve(
|
|
np.array(gender.split(), np.float32),
|
|
original_timestep=float(param['gender_timestep']),
|
|
target_timestep=self.timestep,
|
|
align_length=length
|
|
)
|
|
gender_mask = gender_seq >= 0
|
|
key_shift_seq = gender_seq * (gender_mask * shift_max + (1 - gender_mask) * abs(shift_min))
|
|
batch['key_shift'] = torch.clip(
|
|
torch.from_numpy(key_shift_seq.astype(np.float32)).to(self.device)[None], # => [B=1, T]
|
|
min=shift_min, max=shift_max
|
|
)
|
|
|
|
if hparams['use_speed_embed']:
|
|
if param.get('velocity') is None:
|
|
summary['velocity'] = 'default'
|
|
batch['speed'] = torch.FloatTensor([1.]).to(self.device)[:, None] # => [B=1, T=1]
|
|
else:
|
|
summary['velocity'] = 'manual'
|
|
speed_min, speed_max = hparams['augmentation_args']['random_time_stretching']['range']
|
|
speed_seq = resample_align_curve(
|
|
np.array(param['velocity'].split(), np.float32),
|
|
original_timestep=float(param['velocity_timestep']),
|
|
target_timestep=self.timestep,
|
|
align_length=length
|
|
)
|
|
batch['speed'] = torch.clip(
|
|
torch.from_numpy(speed_seq.astype(np.float32)).to(self.device)[None], # => [B=1, T]
|
|
min=speed_min, max=speed_max
|
|
)
|
|
|
|
print(f'[{idx}]\t' + ', '.join(f'{k}: {v}' for k, v in summary.items()))
|
|
|
|
return batch
|
|
|
|
@torch.no_grad()
|
|
def forward_model(self, sample):
|
|
txt_tokens = sample['tokens']
|
|
variances = {
|
|
v_name: sample.get(v_name)
|
|
for v_name in self.variances_to_embed
|
|
}
|
|
if hparams['use_spk_id']:
|
|
spk_mix_id = sample['spk_mix_id']
|
|
spk_mix_value = sample['spk_mix_value']
|
|
# perform mixing on spk embed
|
|
spk_mix_embed = torch.sum(
|
|
self.model.fs2.spk_embed(spk_mix_id) * spk_mix_value.unsqueeze(3), # => [B, T, N, H]
|
|
dim=2, keepdim=False
|
|
) # => [B, T, H]
|
|
else:
|
|
spk_mix_embed = None
|
|
mel_pred: ShallowDiffusionOutput = self.model(
|
|
txt_tokens, languages=sample.get('languages'),
|
|
mel2ph=sample['mel2ph'], f0=sample['f0'], **variances,
|
|
key_shift=sample.get('key_shift'), speed=sample.get('speed'),
|
|
spk_mix_embed=spk_mix_embed,
|
|
infer=True
|
|
)
|
|
return mel_pred.diff_out
|
|
|
|
@torch.no_grad()
|
|
def run_vocoder(self, spec, **kwargs):
|
|
y = self.vocoder.spec2wav_torch(spec, **kwargs)
|
|
return y[None]
|
|
|
|
def run_inference(
|
|
self, params,
|
|
out_dir: pathlib.Path = None,
|
|
title: str = None,
|
|
num_runs: int = 1,
|
|
spk_mix: Dict[str, float] = None,
|
|
seed: int = -1,
|
|
save_mel: bool = False
|
|
):
|
|
batches = [self.preprocess_input(param, idx=i) for i, param in enumerate(params)]
|
|
|
|
out_dir.mkdir(parents=True, exist_ok=True)
|
|
suffix = '.wav' if not save_mel else '.mel.pt'
|
|
for i in range(num_runs):
|
|
if save_mel:
|
|
result = []
|
|
else:
|
|
result = np.zeros(0)
|
|
current_length = 0
|
|
|
|
for param, batch in tqdm.tqdm(
|
|
zip(params, batches), desc='infer segments', total=len(params)
|
|
):
|
|
if 'seed' in param:
|
|
torch.manual_seed(param["seed"] & 0xffff_ffff)
|
|
torch.cuda.manual_seed_all(param["seed"] & 0xffff_ffff)
|
|
elif seed >= 0:
|
|
torch.manual_seed(seed & 0xffff_ffff)
|
|
torch.cuda.manual_seed_all(seed & 0xffff_ffff)
|
|
|
|
mel_pred = self.forward_model(batch)
|
|
if save_mel:
|
|
result.append({
|
|
'offset': param.get('offset', 0.),
|
|
'mel': mel_pred.cpu(),
|
|
'f0': batch['f0'].cpu()
|
|
})
|
|
else:
|
|
waveform_pred = self.run_vocoder(mel_pred, f0=batch['f0'])[0].cpu().numpy()
|
|
silent_length = round(param.get('offset', 0) * hparams['audio_sample_rate']) - current_length
|
|
if silent_length >= 0:
|
|
result = np.append(result, np.zeros(silent_length))
|
|
result = np.append(result, waveform_pred)
|
|
else:
|
|
result = cross_fade(result, waveform_pred, current_length + silent_length)
|
|
current_length = current_length + silent_length + waveform_pred.shape[0]
|
|
|
|
if num_runs > 1:
|
|
filename = f'{title}-{str(i).zfill(3)}{suffix}'
|
|
else:
|
|
filename = title + suffix
|
|
save_path = out_dir / filename
|
|
if save_mel:
|
|
print(f'| save mel: {save_path}')
|
|
torch.save(result, save_path)
|
|
else:
|
|
print(f'| save audio: {save_path}')
|
|
save_wav(result, save_path, hparams['audio_sample_rate'])
|