wehub/openvpi--diffsinger
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:35:17 +08:00
commit 344816a5d8
136 changed files with 25044 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
inference/dpm_solver_pytorch.py
+1305
View File
File diff suppressed because it is too large Load Diff
inference/ds_acoustic.py
+271
View File
@@ -0,0 +1,271 @@
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'])
inference/ds_variance.py
+468
View File
@@ -0,0 +1,468 @@
import copy
import json
import pathlib
from collections import OrderedDict
from typing import List, Tuple
import librosa
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import tqdm
from scipy import interpolate
from basics.base_svs_infer import BaseSVSInfer
from modules.fastspeech.param_adaptor import VARIANCE_CHECKLIST
from modules.fastspeech.tts_modules import (
LengthRegulator, RhythmRegulator,
mel2ph_to_dur
)
from modules.toplevel import DiffSingerVariance
from utils import load_ckpt
from utils.hparams import hparams
from utils.infer_utils import resample_align_curve
from utils.phoneme_utils import load_phoneme_dictionary
from utils.pitch_utils import interp_f0
class DiffSingerVarianceInfer(BaseSVSInfer):
def __init__(
self, device=None, ckpt_steps=None,
predictions: set = None
):
super().__init__(device=device)
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: DiffSingerVariance = self.build_model(ckpt_steps=ckpt_steps)
self.lr = LengthRegulator()
self.rr = RhythmRegulator()
smooth_kernel_size = round(hparams['midi_smooth_width'] / self.timestep)
self.smooth = nn.Conv1d(
in_channels=1,
out_channels=1,
kernel_size=smooth_kernel_size,
bias=False,
padding='same',
padding_mode='replicate'
).eval().to(self.device)
smooth_kernel = torch.sin(torch.from_numpy(
np.linspace(0, 1, smooth_kernel_size).astype(np.float32) * np.pi
).to(self.device))
smooth_kernel /= smooth_kernel.sum()
self.smooth.weight.data = smooth_kernel[None, None]
glide_types = hparams.get('glide_types', [])
assert 'none' not in glide_types, 'Type name \'none\' is reserved and should not appear in glide_types.'
self.glide_map = {
'none': 0,
**{
typename: idx + 1
for idx, typename in enumerate(glide_types)
}
}
self.auto_completion_mode = len(predictions) == 0
self.global_predict_dur = 'dur' in predictions and hparams['predict_dur']
self.global_predict_pitch = 'pitch' in predictions and hparams['predict_pitch']
self.variance_prediction_set = predictions.intersection(VARIANCE_CHECKLIST)
self.global_predict_variances = len(self.variance_prediction_set) > 0
def build_model(self, ckpt_steps=None):
model = DiffSingerVariance(
vocab_size=len(self.phoneme_dictionary)
).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
@torch.no_grad()
def preprocess_input(
self, param, idx=0,
load_dur: bool = False,
load_pitch: bool = False
):
"""
:param param: one segment in the .ds file
:param idx: index of the segment
:param load_dur: whether ph_dur is loaded
:param load_pitch: whether pitch is loaded
: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=1, T_ph]
batch['languages'] = languages
txt_tokens = torch.LongTensor([
self.phoneme_dictionary.encode(param['ph_seq'], lang=lang)
]).to(self.device) # [B=1, T_ph]
T_ph = txt_tokens.shape[1]
batch['tokens'] = txt_tokens
ph_num = torch.from_numpy(np.array([param['ph_num'].split()], np.int64)).to(self.device) # [B=1, T_w]
ph2word = self.lr(ph_num) # => [B=1, T_ph]
T_w = int(ph2word.max())
batch['ph2word'] = ph2word
note_midi = np.array(
[(librosa.note_to_midi(n, round_midi=False) if n != 'rest' else -1) for n in param['note_seq'].split()],
dtype=np.float32
)
note_rest = note_midi < 0
if np.all(note_rest):
# All rests, fill with constants
note_midi = np.full_like(note_midi, fill_value=60.)
else:
# Interpolate rest values
interp_func = interpolate.interp1d(
np.where(~note_rest)[0], note_midi[~note_rest],
kind='nearest', fill_value='extrapolate'
)
note_midi[note_rest] = interp_func(np.where(note_rest)[0])
note_midi = torch.from_numpy(note_midi).to(self.device)[None] # [B=1, T_n]
note_rest = torch.from_numpy(note_rest).to(self.device)[None] # [B=1, T_n]
T_n = note_midi.shape[1]
note_dur_sec = torch.from_numpy(np.array([param['note_dur'].split()], np.float32)).to(self.device) # [B=1, T_n]
note_acc = torch.round(torch.cumsum(note_dur_sec, dim=1) / self.timestep + 0.5).long()
note_dur = torch.diff(note_acc, dim=1, prepend=note_acc.new_zeros(1, 1))
mel2note = self.lr(note_dur) # [B=1, T_s]
T_s = mel2note.shape[1]
summary['words'] = T_w
summary['notes'] = T_n
summary['tokens'] = T_ph
summary['frames'] = T_s
summary['seconds'] = '%.2f' % (T_s * self.timestep)
if hparams['use_spk_id']:
ph_spk_mix_id, ph_spk_mix_value = self.load_speaker_mix(
param_src=param, summary_dst=summary, mix_mode='token', mix_length=T_ph
)
spk_mix_id, spk_mix_value = self.load_speaker_mix(
param_src=param, summary_dst=summary, mix_mode='frame', mix_length=T_s
)
batch['ph_spk_mix_id'] = ph_spk_mix_id
batch['ph_spk_mix_value'] = ph_spk_mix_value
batch['spk_mix_id'] = spk_mix_id
batch['spk_mix_value'] = spk_mix_value
if load_dur:
# Get mel2ph if ph_dur is needed
ph_dur_sec = torch.from_numpy(
np.array([param['ph_dur'].split()], np.float32)
).to(self.device) # [B=1, T_ph]
ph_acc = torch.round(torch.cumsum(ph_dur_sec, dim=1) / self.timestep + 0.5).long()
ph_dur = torch.diff(ph_acc, dim=1, prepend=ph_acc.new_zeros(1, 1))
mel2ph = self.lr(ph_dur, txt_tokens == 0)
if mel2ph.shape[1] != T_s: # Align phones with notes
mel2ph = F.pad(mel2ph, [0, T_s - mel2ph.shape[1]], value=mel2ph[0, -1])
ph_dur = mel2ph_to_dur(mel2ph, T_ph)
# Get word_dur from ph_dur and ph_num
word_dur = note_dur.new_zeros(1, T_w + 1).scatter_add(
1, ph2word, ph_dur
)[:, 1:] # => [B=1, T_w]
else:
ph_dur = None
mel2ph = None
# Get word_dur from note_dur and note_slur
is_slur = torch.BoolTensor([[int(s) for s in param['note_slur'].split()]]).to(self.device) # [B=1, T_n]
note2word = torch.cumsum(~is_slur, dim=1) # [B=1, T_n]
word_dur = note_dur.new_zeros(1, T_w + 1).scatter_add(
1, note2word, note_dur
)[:, 1:] # => [B=1, T_w]
batch['ph_dur'] = ph_dur
batch['mel2ph'] = mel2ph
mel2word = self.lr(word_dur) # [B=1, T_s]
if mel2word.shape[1] != T_s: # Align words with notes
mel2word = F.pad(mel2word, [0, T_s - mel2word.shape[1]], value=mel2word[0, -1])
word_dur = mel2ph_to_dur(mel2word, T_w)
batch['word_dur'] = word_dur
batch['note_midi'] = note_midi
batch['note_dur'] = note_dur
batch['note_rest'] = note_rest
if hparams.get('use_glide_embed', False) and param.get('note_glide') is not None:
batch['note_glide'] = torch.LongTensor(
[[self.glide_map.get(x, 0) for x in param['note_glide'].split()]]
).to(self.device)
else:
batch['note_glide'] = torch.zeros(1, T_n, dtype=torch.long, device=self.device)
batch['mel2note'] = mel2note
# Calculate and smoothen the frame-level MIDI pitch, which is a step function curve
frame_midi_pitch = torch.gather(
F.pad(note_midi, [1, 0]), 1, mel2note
) # => frame-level MIDI pitch, [B=1, T_s]
base_pitch = self.smooth(frame_midi_pitch)
batch['base_pitch'] = base_pitch
if ph_dur is not None:
# Phone durations are available, calculate phoneme-level MIDI.
mel2pdur = torch.gather(F.pad(ph_dur, [1, 0], value=1), 1, mel2ph) # frame-level phone duration
ph_midi = frame_midi_pitch.new_zeros(1, T_ph + 1).scatter_add(
1, mel2ph, frame_midi_pitch / mel2pdur
)[:, 1:]
else:
# Phone durations are not available, calculate word-level MIDI instead.
mel2wdur = torch.gather(F.pad(word_dur, [1, 0], value=1), 1, mel2word)
w_midi = frame_midi_pitch.new_zeros(1, T_w + 1).scatter_add(
1, mel2word, frame_midi_pitch / mel2wdur
)[:, 1:]
# Convert word-level MIDI to phoneme-level MIDI
ph_midi = torch.gather(F.pad(w_midi, [1, 0]), 1, ph2word)
ph_midi = ph_midi.round().long()
batch['midi'] = ph_midi
if load_pitch:
f0 = resample_align_curve(
np.array(param['f0_seq'].split(), np.float32),
original_timestep=float(param['f0_timestep']),
target_timestep=self.timestep,
align_length=T_s
)
batch['pitch'] = torch.from_numpy(
librosa.hz_to_midi(interp_f0(f0)[0]).astype(np.float32)
).to(self.device)[None]
if self.model.predict_dur:
if load_dur:
summary['ph_dur'] = 'manual'
elif self.auto_completion_mode or self.global_predict_dur:
summary['ph_dur'] = 'auto'
else:
summary['ph_dur'] = 'ignored'
if self.model.predict_pitch:
if load_pitch:
summary['pitch'] = 'manual'
elif self.auto_completion_mode or self.global_predict_pitch:
summary['pitch'] = 'auto'
# Load expressiveness
expr = param.get('expr', 1.)
if isinstance(expr, (int, float, bool)):
summary['expr'] = f'static({expr:.3f})'
batch['expr'] = torch.FloatTensor([expr]).to(self.device)[:, None] # [B=1, T=1]
else:
summary['expr'] = 'dynamic'
expr = resample_align_curve(
np.array(expr.split(), np.float32),
original_timestep=float(param['expr_timestep']),
target_timestep=self.timestep,
align_length=T_s
)
batch['expr'] = torch.from_numpy(expr.astype(np.float32)).to(self.device)[None]
else:
summary['pitch'] = 'ignored'
if self.model.predict_variances:
for v_name in self.model.variance_prediction_list:
if self.auto_completion_mode and param.get(v_name) is None or v_name in self.variance_prediction_set:
summary[v_name] = 'auto'
else:
summary[v_name] = 'ignored'
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']
midi = sample['midi']
ph2word = sample['ph2word']
word_dur = sample['word_dur']
ph_dur = sample['ph_dur']
mel2ph = sample['mel2ph']
note_midi = sample['note_midi']
note_rest = sample['note_rest']
note_dur = sample['note_dur']
note_glide = sample['note_glide']
mel2note = sample['mel2note']
base_pitch = sample['base_pitch']
expr = sample.get('expr')
pitch = sample.get('pitch')
if hparams['use_spk_id']:
ph_spk_mix_id = sample['ph_spk_mix_id']
ph_spk_mix_value = sample['ph_spk_mix_value']
spk_mix_id = sample['spk_mix_id']
spk_mix_value = sample['spk_mix_value']
ph_spk_mix_embed = torch.sum(
self.model.spk_embed(ph_spk_mix_id) * ph_spk_mix_value.unsqueeze(3), # => [B, T_ph, N, H]
dim=2, keepdim=False
) # => [B, T_ph, H]
spk_mix_embed = torch.sum(
self.model.spk_embed(spk_mix_id) * spk_mix_value.unsqueeze(3), # => [B, T_s, N, H]
dim=2, keepdim=False
) # [B, T_s, H]
else:
ph_spk_mix_embed = spk_mix_embed = None
dur_pred, pitch_pred, variance_pred = self.model(
txt_tokens, languages=sample.get('languages'),
midi=midi, ph2word=ph2word, word_dur=word_dur, ph_dur=ph_dur, mel2ph=mel2ph,
note_midi=note_midi, note_rest=note_rest, note_dur=note_dur, note_glide=note_glide, mel2note=mel2note,
base_pitch=base_pitch, pitch=pitch, pitch_expr=expr,
ph_spk_mix_embed=ph_spk_mix_embed, spk_mix_embed=spk_mix_embed,
infer=True
)
if dur_pred is not None:
dur_pred = self.rr(dur_pred, ph2word, word_dur)
if pitch_pred is not None:
pitch_pred = base_pitch + pitch_pred
return dur_pred, pitch_pred, variance_pred
def infer_once(self, param):
batch = self.preprocess_input(param)
dur_pred, pitch_pred, variance_pred = self.forward_model(batch)
if dur_pred is not None:
dur_pred = dur_pred[0].cpu().numpy()
if pitch_pred is not None:
pitch_pred = pitch_pred[0].cpu().numpy()
f0_pred = librosa.midi_to_hz(pitch_pred)
else:
f0_pred = None
variance_pred = {
k: v[0].cpu().numpy()
for k, v in variance_pred.items()
}
return dur_pred, f0_pred, variance_pred
def run_inference(
self, params,
out_dir: pathlib.Path = None,
title: str = None,
num_runs: int = 1,
seed: int = -1
):
batches = []
predictor_flags: List[Tuple[bool, bool, bool]] = []
for i, param in enumerate(params):
param: dict
if self.auto_completion_mode:
flag = (
self.model.fs2.predict_dur and param.get('ph_dur') is None,
self.model.predict_pitch and param.get('f0_seq') is None,
self.model.predict_variances and any(
param.get(v_name) is None for v_name in self.model.variance_prediction_list
)
)
else:
predict_variances = self.model.predict_variances and self.global_predict_variances
predict_pitch = self.model.predict_pitch and (
self.global_predict_pitch or (param.get('f0_seq') is None and predict_variances)
)
predict_dur = self.model.predict_dur and (
self.global_predict_dur or (param.get('ph_dur') is None and (predict_pitch or predict_variances))
)
flag = (predict_dur, predict_pitch, predict_variances)
predictor_flags.append(flag)
batches.append(self.preprocess_input(
param, idx=i,
load_dur=not flag[0] and (flag[1] or flag[2]),
load_pitch=not flag[1] and flag[2]
))
out_dir.mkdir(parents=True, exist_ok=True)
for i in range(num_runs):
results = []
for param, flag, batch in tqdm.tqdm(
zip(params, predictor_flags, 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)
param_copy = copy.deepcopy(param)
flag_saved = (
self.model.fs2.predict_dur,
self.model.predict_pitch,
self.model.predict_variances
)
(
self.model.fs2.predict_dur,
self.model.predict_pitch,
self.model.predict_variances
) = flag
dur_pred, pitch_pred, variance_pred = self.forward_model(batch)
(
self.model.fs2.predict_dur,
self.model.predict_pitch,
self.model.predict_variances
) = flag_saved
if dur_pred is not None and (self.auto_completion_mode or self.global_predict_dur):
dur_pred = dur_pred[0].cpu().numpy()
param_copy['ph_dur'] = ' '.join(str(round(dur, 6)) for dur in (dur_pred * self.timestep).tolist())
if pitch_pred is not None and (self.auto_completion_mode or self.global_predict_pitch):
pitch_pred = pitch_pred[0].cpu().numpy()
f0_pred = librosa.midi_to_hz(pitch_pred)
param_copy['f0_seq'] = ' '.join([str(round(freq, 1)) for freq in f0_pred.tolist()])
param_copy['f0_timestep'] = str(self.timestep)
variance_pred = {
k: v[0].cpu().numpy()
for k, v in variance_pred.items()
if (self.auto_completion_mode and param.get(k) is None) or k in self.variance_prediction_set
}
for v_name, v_pred in variance_pred.items():
param_copy[v_name] = ' '.join([str(round(v, 4)) for v in v_pred.tolist()])
param_copy[f'{v_name}_timestep'] = str(self.timestep)
# Restore ph_spk_mix and spk_mix
if 'ph_spk_mix' in param_copy and 'spk_mix' in param_copy:
if 'ph_spk_mix_backup' in param_copy:
if param_copy['ph_spk_mix_backup'] is None:
del param_copy['ph_spk_mix']
else:
param_copy['ph_spk_mix'] = param_copy['ph_spk_mix_backup']
del param['ph_spk_mix_backup']
if 'spk_mix_backup' in param_copy:
if param_copy['ph_spk_mix_backup'] is None:
del param_copy['spk_mix']
else:
param_copy['spk_mix'] = param_copy['spk_mix_backup']
del param['spk_mix_backup']
results.append(param_copy)
if num_runs > 1:
filename = f'{title}-{str(i).zfill(3)}.ds'
else:
filename = f'{title}.ds'
save_path = out_dir / filename
with open(save_path, 'w', encoding='utf8') as f:
print(f'| save params: {save_path}')
json.dump(results, f, ensure_ascii=False, indent=2)
inference/uni_pc.py
+731
View File
@@ -0,0 +1,731 @@
import torch
import torch.nn.functional as F
import math
class NoiseScheduleVP:
def __init__(
self,
schedule='discrete',
betas=None,
alphas_cumprod=None,
continuous_beta_0=0.1,
continuous_beta_1=20.,
dtype=torch.float32,
):
"""Create a wrapper class for the forward SDE (VP type).
***
Update: We support discrete-time diffusion models by implementing a picewise linear interpolation for log_alpha_t.
We recommend to use schedule='discrete' for the discrete-time diffusion models, especially for high-resolution images.
***
The forward SDE ensures that the condition distribution q_{t|0}(x_t | x_0) = N ( alpha_t * x_0, sigma_t^2 * I ).
We further define lambda_t = log(alpha_t) - log(sigma_t), which is the half-logSNR (described in the DPM-Solver paper).
Therefore, we implement the functions for computing alpha_t, sigma_t and lambda_t. For t in [0, T], we have:
log_alpha_t = self.marginal_log_mean_coeff(t)
sigma_t = self.marginal_std(t)
lambda_t = self.marginal_lambda(t)
Moreover, as lambda(t) is an invertible function, we also support its inverse function:
t = self.inverse_lambda(lambda_t)
===============================================================
We support both discrete-time DPMs (trained on n = 0, 1, ..., N-1) and continuous-time DPMs (trained on t in [t_0, T]).
1. For discrete-time DPMs:
For discrete-time DPMs trained on n = 0, 1, ..., N-1, we convert the discrete steps to continuous time steps by:
t_i = (i + 1) / N
e.g. for N = 1000, we have t_0 = 1e-3 and T = t_{N-1} = 1.
We solve the corresponding diffusion ODE from time T = 1 to time t_0 = 1e-3.
Args:
betas: A `torch.Tensor`. The beta array for the discrete-time DPM. (See the original DDPM paper for details)
alphas_cumprod: A `torch.Tensor`. The cumprod alphas for the discrete-time DPM. (See the original DDPM paper for details)
Note that we always have alphas_cumprod = cumprod(1 - betas). Therefore, we only need to set one of `betas` and `alphas_cumprod`.
**Important**: Please pay special attention for the args for `alphas_cumprod`:
The `alphas_cumprod` is the \hat{alpha_n} arrays in the notations of DDPM. Specifically, DDPMs assume that
q_{t_n | 0}(x_{t_n} | x_0) = N ( \sqrt{\hat{alpha_n}} * x_0, (1 - \hat{alpha_n}) * I ).
Therefore, the notation \hat{alpha_n} is different from the notation alpha_t in DPM-Solver. In fact, we have
alpha_{t_n} = \sqrt{\hat{alpha_n}},
and
log(alpha_{t_n}) = 0.5 * log(\hat{alpha_n}).
2. For continuous-time DPMs:
We support two types of VPSDEs: linear (DDPM) and cosine (improved-DDPM). The hyperparameters for the noise
schedule are the default settings in DDPM and improved-DDPM:
Args:
beta_min: A `float` number. The smallest beta for the linear schedule.
beta_max: A `float` number. The largest beta for the linear schedule.
cosine_s: A `float` number. The hyperparameter in the cosine schedule.
cosine_beta_max: A `float` number. The hyperparameter in the cosine schedule.
T: A `float` number. The ending time of the forward process.
===============================================================
Args:
schedule: A `str`. The noise schedule of the forward SDE. 'discrete' for discrete-time DPMs,
'linear' or 'cosine' for continuous-time DPMs.
Returns:
A wrapper object of the forward SDE (VP type).
===============================================================
Example:
# For discrete-time DPMs, given betas (the beta array for n = 0, 1, ..., N - 1):
>>> ns = NoiseScheduleVP('discrete', betas=betas)
# For discrete-time DPMs, given alphas_cumprod (the \hat{alpha_n} array for n = 0, 1, ..., N - 1):
>>> ns = NoiseScheduleVP('discrete', alphas_cumprod=alphas_cumprod)
# For continuous-time DPMs (VPSDE), linear schedule:
>>> ns = NoiseScheduleVP('linear', continuous_beta_0=0.1, continuous_beta_1=20.)
"""
if schedule not in ['discrete', 'linear', 'cosine']:
raise ValueError("Unsupported noise schedule {}. The schedule needs to be 'discrete' or 'linear' or 'cosine'".format(schedule))
self.schedule = schedule
if schedule == 'discrete':
if betas is not None:
log_alphas = 0.5 * torch.log(1 - betas).cumsum(dim=0)
else:
assert alphas_cumprod is not None
log_alphas = 0.5 * torch.log(alphas_cumprod)
self.total_N = len(log_alphas)
self.T = 1.
self.t_array = torch.linspace(0., 1., self.total_N + 1)[1:].reshape((1, -1)).to(dtype=dtype)
self.log_alpha_array = log_alphas.reshape((1, -1,)).to(dtype=dtype)
else:
self.total_N = 1000
self.beta_0 = continuous_beta_0
self.beta_1 = continuous_beta_1
self.cosine_s = 0.008
self.cosine_beta_max = 999.
self.cosine_t_max = math.atan(self.cosine_beta_max * (1. + self.cosine_s) / math.pi) * 2. * (1. + self.cosine_s) / math.pi - self.cosine_s
self.cosine_log_alpha_0 = math.log(math.cos(self.cosine_s / (1. + self.cosine_s) * math.pi / 2.))
self.schedule = schedule
if schedule == 'cosine':
# For the cosine schedule, T = 1 will have numerical issues. So we manually set the ending time T.
# Note that T = 0.9946 may be not the optimal setting. However, we find it works well.
self.T = 0.9946
else:
self.T = 1.
def marginal_log_mean_coeff(self, t):
"""
Compute log(alpha_t) of a given continuous-time label t in [0, T].
"""
if self.schedule == 'discrete':
return interpolate_fn(t.reshape((-1, 1)), self.t_array.to(t.device), self.log_alpha_array.to(t.device)).reshape((-1))
elif self.schedule == 'linear':
return -0.25 * t ** 2 * (self.beta_1 - self.beta_0) - 0.5 * t * self.beta_0
elif self.schedule == 'cosine':
log_alpha_fn = lambda s: torch.log(torch.cos((s + self.cosine_s) / (1. + self.cosine_s) * math.pi / 2.))
log_alpha_t = log_alpha_fn(t) - self.cosine_log_alpha_0
return log_alpha_t
def marginal_alpha(self, t):
"""
Compute alpha_t of a given continuous-time label t in [0, T].
"""
return torch.exp(self.marginal_log_mean_coeff(t))
def marginal_std(self, t):
"""
Compute sigma_t of a given continuous-time label t in [0, T].
"""
return torch.sqrt(1. - torch.exp(2. * self.marginal_log_mean_coeff(t)))
def marginal_lambda(self, t):
"""
Compute lambda_t = log(alpha_t) - log(sigma_t) of a given continuous-time label t in [0, T].
"""
log_mean_coeff = self.marginal_log_mean_coeff(t)
log_std = 0.5 * torch.log(1. - torch.exp(2. * log_mean_coeff))
return log_mean_coeff - log_std
def inverse_lambda(self, lamb):
"""
Compute the continuous-time label t in [0, T] of a given half-logSNR lambda_t.
"""
if self.schedule == 'linear':
tmp = 2. * (self.beta_1 - self.beta_0) * torch.logaddexp(-2. * lamb, torch.zeros((1,)).to(lamb))
Delta = self.beta_0**2 + tmp
return tmp / (torch.sqrt(Delta) + self.beta_0) / (self.beta_1 - self.beta_0)
elif self.schedule == 'discrete':
log_alpha = -0.5 * torch.logaddexp(torch.zeros((1,)).to(lamb.device), -2. * lamb)
t = interpolate_fn(log_alpha.reshape((-1, 1)), torch.flip(self.log_alpha_array.to(lamb.device), [1]), torch.flip(self.t_array.to(lamb.device), [1]))
return t.reshape((-1,))
else:
log_alpha = -0.5 * torch.logaddexp(-2. * lamb, torch.zeros((1,)).to(lamb))
t_fn = lambda log_alpha_t: torch.arccos(torch.exp(log_alpha_t + self.cosine_log_alpha_0)) * 2. * (1. + self.cosine_s) / math.pi - self.cosine_s
t = t_fn(log_alpha)
return t
def model_wrapper(
model,
noise_schedule,
model_type="noise",
model_kwargs={},
guidance_type="uncond",
condition=None,
unconditional_condition=None,
guidance_scale=1.,
classifier_fn=None,
classifier_kwargs={},
):
"""Create a wrapper function for the noise prediction model.
"""
def get_model_input_time(t_continuous):
"""
Convert the continuous-time `t_continuous` (in [epsilon, T]) to the model input time.
For discrete-time DPMs, we convert `t_continuous` in [1 / N, 1] to `t_input` in [0, 1000 * (N - 1) / N].
For continuous-time DPMs, we just use `t_continuous`.
"""
if noise_schedule.schedule == 'discrete':
return (t_continuous - 1. / noise_schedule.total_N) * noise_schedule.total_N
else:
return t_continuous
def noise_pred_fn(x, t_continuous, cond=None):
t_input = get_model_input_time(t_continuous)
if cond is None:
output = model(x, t_input, **model_kwargs)
else:
output = model(x, t_input, cond, **model_kwargs)
if model_type == "noise":
return output
elif model_type == "x_start":
alpha_t, sigma_t = noise_schedule.marginal_alpha(t_continuous), noise_schedule.marginal_std(t_continuous)
return (x - alpha_t * output) / sigma_t
elif model_type == "v":
alpha_t, sigma_t = noise_schedule.marginal_alpha(t_continuous), noise_schedule.marginal_std(t_continuous)
return alpha_t * output + sigma_t * x
elif model_type == "score":
sigma_t = noise_schedule.marginal_std(t_continuous)
return -sigma_t * output
def cond_grad_fn(x, t_input):
"""
Compute the gradient of the classifier, i.e. nabla_{x} log p_t(cond | x_t).
"""
with torch.enable_grad():
x_in = x.detach().requires_grad_(True)
log_prob = classifier_fn(x_in, t_input, condition, **classifier_kwargs)
return torch.autograd.grad(log_prob.sum(), x_in)[0]
def model_fn(x, t_continuous):
"""
The noise predicition model function that is used for DPM-Solver.
"""
if guidance_type == "uncond":
return noise_pred_fn(x, t_continuous)
elif guidance_type == "classifier":
assert classifier_fn is not None
t_input = get_model_input_time(t_continuous)
cond_grad = cond_grad_fn(x, t_input)
sigma_t = noise_schedule.marginal_std(t_continuous)
noise = noise_pred_fn(x, t_continuous)
return noise - guidance_scale * sigma_t * cond_grad
elif guidance_type == "classifier-free":
if guidance_scale == 1. or unconditional_condition is None:
return noise_pred_fn(x, t_continuous, cond=condition)
else:
x_in = torch.cat([x] * 2)
t_in = torch.cat([t_continuous] * 2)
c_in = torch.cat([unconditional_condition, condition])
noise_uncond, noise = noise_pred_fn(x_in, t_in, cond=c_in).chunk(2)
return noise_uncond + guidance_scale * (noise - noise_uncond)
assert model_type in ["noise", "x_start", "v"]
assert guidance_type in ["uncond", "classifier", "classifier-free"]
return model_fn
class UniPC:
def __init__(
self,
model_fn,
noise_schedule,
algorithm_type="data_prediction",
correcting_x0_fn=None,
correcting_xt_fn=None,
thresholding_max_val=1.,
dynamic_thresholding_ratio=0.995,
variant='bh1'
):
"""Construct a UniPC.
We support both data_prediction and noise_prediction.
"""
self.model = lambda x, t: model_fn(x, t.expand((x.shape[0])))
self.noise_schedule = noise_schedule
assert algorithm_type in ["data_prediction", "noise_prediction"]
if correcting_x0_fn == "dynamic_thresholding":
self.correcting_x0_fn = self.dynamic_thresholding_fn
else:
self.correcting_x0_fn = correcting_x0_fn
self.correcting_xt_fn = correcting_xt_fn
self.dynamic_thresholding_ratio = dynamic_thresholding_ratio
self.thresholding_max_val = thresholding_max_val
self.variant = variant
self.predict_x0 = algorithm_type == "data_prediction"
def dynamic_thresholding_fn(self, x0, t=None):
"""
The dynamic thresholding method.
"""
dims = x0.dim()
p = self.dynamic_thresholding_ratio
s = torch.quantile(torch.abs(x0).reshape((x0.shape[0], -1)), p, dim=1)
s = expand_dims(torch.maximum(s, self.thresholding_max_val * torch.ones_like(s).to(s.device)), dims)
x0 = torch.clamp(x0, -s, s) / s
return x0
def noise_prediction_fn(self, x, t):
"""
Return the noise prediction model.
"""
return self.model(x, t)
def data_prediction_fn(self, x, t):
"""
Return the data prediction model (with corrector).
"""
noise = self.noise_prediction_fn(x, t)
alpha_t, sigma_t = self.noise_schedule.marginal_alpha(t), self.noise_schedule.marginal_std(t)
x0 = (x - sigma_t * noise) / alpha_t
if self.correcting_x0_fn is not None:
x0 = self.correcting_x0_fn(x0)
return x0
def model_fn(self, x, t):
"""
Convert the model to the noise prediction model or the data prediction model.
"""
if self.predict_x0:
return self.data_prediction_fn(x, t)
else:
return self.noise_prediction_fn(x, t)
def get_time_steps(self, skip_type, t_T, t_0, N, device):
"""Compute the intermediate time steps for sampling.
"""
if skip_type == 'logSNR':
lambda_T = self.noise_schedule.marginal_lambda(torch.tensor(t_T).to(device))
lambda_0 = self.noise_schedule.marginal_lambda(torch.tensor(t_0).to(device))
logSNR_steps = torch.linspace(lambda_T.cpu().item(), lambda_0.cpu().item(), N + 1).to(device)
return self.noise_schedule.inverse_lambda(logSNR_steps)
elif skip_type == 'time_uniform':
return torch.linspace(t_T, t_0, N + 1).to(device)
elif skip_type == 'time_quadratic':
t_order = 2
t = torch.linspace(t_T**(1. / t_order), t_0**(1. / t_order), N + 1).pow(t_order).to(device)
return t
else:
raise ValueError("Unsupported skip_type {}, need to be 'logSNR' or 'time_uniform' or 'time_quadratic'".format(skip_type))
def get_orders_and_timesteps_for_singlestep_solver(self, steps, order, skip_type, t_T, t_0, device):
"""
Get the order of each step for sampling by the singlestep DPM-Solver.
"""
if order == 3:
K = steps // 3 + 1
if steps % 3 == 0:
orders = [3,] * (K - 2) + [2, 1]
elif steps % 3 == 1:
orders = [3,] * (K - 1) + [1]
else:
orders = [3,] * (K - 1) + [2]
elif order == 2:
if steps % 2 == 0:
K = steps // 2
orders = [2,] * K
else:
K = steps // 2 + 1
orders = [2,] * (K - 1) + [1]
elif order == 1:
K = steps
orders = [1,] * steps
else:
raise ValueError("'order' must be '1' or '2' or '3'.")
if skip_type == 'logSNR':
# To reproduce the results in DPM-Solver paper
timesteps_outer = self.get_time_steps(skip_type, t_T, t_0, K, device)
else:
timesteps_outer = self.get_time_steps(skip_type, t_T, t_0, steps, device)[torch.cumsum(torch.tensor([0,] + orders), 0).to(device)]
return timesteps_outer, orders
def denoise_to_zero_fn(self, x, s):
"""
Denoise at the final step, which is equivalent to solve the ODE from lambda_s to infty by first-order discretization.
"""
return self.data_prediction_fn(x, s)
def multistep_uni_pc_update(self, x, model_prev_list, t_prev_list, t, order, **kwargs):
if len(t.shape) == 0:
t = t.view(-1)
if 'bh' in self.variant:
return self.multistep_uni_pc_bh_update(x, model_prev_list, t_prev_list, t, order, **kwargs)
else:
assert self.variant == 'vary_coeff'
return self.multistep_uni_pc_vary_update(x, model_prev_list, t_prev_list, t, order, **kwargs)
def multistep_uni_pc_vary_update(self, x, model_prev_list, t_prev_list, t, order, use_corrector=True):
#print(f'using unified predictor-corrector with order {order} (solver type: vary coeff)')
ns = self.noise_schedule
assert order <= len(model_prev_list)
# first compute rks
t_prev_0 = t_prev_list[-1]
lambda_prev_0 = ns.marginal_lambda(t_prev_0)
lambda_t = ns.marginal_lambda(t)
model_prev_0 = model_prev_list[-1]
sigma_prev_0, sigma_t = ns.marginal_std(t_prev_0), ns.marginal_std(t)
log_alpha_t = ns.marginal_log_mean_coeff(t)
alpha_t = torch.exp(log_alpha_t)
h = lambda_t - lambda_prev_0
rks = []
D1s = []
for i in range(1, order):
t_prev_i = t_prev_list[-(i + 1)]
model_prev_i = model_prev_list[-(i + 1)]
lambda_prev_i = ns.marginal_lambda(t_prev_i)
rk = (lambda_prev_i - lambda_prev_0) / h
rks.append(rk)
D1s.append((model_prev_i - model_prev_0) / rk)
rks.append(1.)
rks = torch.tensor(rks, device=x.device)
K = len(rks)
# build C matrix
C = []
col = torch.ones_like(rks)
for k in range(1, K + 1):
C.append(col)
col = col * rks / (k + 1)
C = torch.stack(C, dim=1)
if len(D1s) > 0:
D1s = torch.stack(D1s, dim=1) # (B, K)
C_inv_p = torch.linalg.inv(C[:-1, :-1])
A_p = C_inv_p
if use_corrector:
#print('using corrector')
C_inv = torch.linalg.inv(C)
A_c = C_inv
hh = -h if self.predict_x0 else h
h_phi_1 = torch.expm1(hh)
h_phi_ks = []
factorial_k = 1
h_phi_k = h_phi_1
for k in range(1, K + 2):
h_phi_ks.append(h_phi_k)
h_phi_k = h_phi_k / hh - 1 / factorial_k
factorial_k *= (k + 1)
model_t = None
if self.predict_x0:
x_t_ = (
sigma_t / sigma_prev_0 * x
- alpha_t * h_phi_1 * model_prev_0
)
# now predictor
x_t = x_t_
if len(D1s) > 0:
# compute the residuals for predictor
for k in range(K - 1):
x_t = x_t - alpha_t * h_phi_ks[k + 1] * torch.einsum('bkchw,k->bchw', D1s, A_p[k])
# now corrector
if use_corrector:
model_t = self.model_fn(x_t, t)
D1_t = (model_t - model_prev_0)
x_t = x_t_
k = 0
for k in range(K - 1):
x_t = x_t - alpha_t * h_phi_ks[k + 1] * torch.einsum('bkchw,k->bchw', D1s, A_c[k][:-1])
x_t = x_t - alpha_t * h_phi_ks[K] * (D1_t * A_c[k][-1])
else:
log_alpha_prev_0, log_alpha_t = ns.marginal_log_mean_coeff(t_prev_0), ns.marginal_log_mean_coeff(t)
x_t_ = (
(torch.exp(log_alpha_t - log_alpha_prev_0)) * x
- (sigma_t * h_phi_1) * model_prev_0
)
# now predictor
x_t = x_t_
if len(D1s) > 0:
# compute the residuals for predictor
for k in range(K - 1):
x_t = x_t - sigma_t * h_phi_ks[k + 1] * torch.einsum('bkchw,k->bchw', D1s, A_p[k])
# now corrector
if use_corrector:
model_t = self.model_fn(x_t, t)
D1_t = (model_t - model_prev_0)
x_t = x_t_
k = 0
for k in range(K - 1):
x_t = x_t - sigma_t * h_phi_ks[k + 1] * torch.einsum('bkchw,k->bchw', D1s, A_c[k][:-1])
x_t = x_t - sigma_t * h_phi_ks[K] * (D1_t * A_c[k][-1])
return x_t, model_t
def multistep_uni_pc_bh_update(self, x, model_prev_list, t_prev_list, t, order, x_t=None, use_corrector=True):
#print(f'using unified predictor-corrector with order {order} (solver type: B(h))')
ns = self.noise_schedule
assert order <= len(model_prev_list)
# first compute rks
t_prev_0 = t_prev_list[-1]
lambda_prev_0 = ns.marginal_lambda(t_prev_0)
lambda_t = ns.marginal_lambda(t)
model_prev_0 = model_prev_list[-1]
sigma_prev_0, sigma_t = ns.marginal_std(t_prev_0), ns.marginal_std(t)
log_alpha_prev_0, log_alpha_t = ns.marginal_log_mean_coeff(t_prev_0), ns.marginal_log_mean_coeff(t)
alpha_t = torch.exp(log_alpha_t)
h = lambda_t - lambda_prev_0
rks = []
D1s = []
for i in range(1, order):
t_prev_i = t_prev_list[-(i + 1)]
model_prev_i = model_prev_list[-(i + 1)]
lambda_prev_i = ns.marginal_lambda(t_prev_i)
rk = (lambda_prev_i - lambda_prev_0) / h
rks.append(rk)
D1s.append((model_prev_i - model_prev_0) / rk)
rks.append(1.)
rks = torch.tensor(rks, device=x.device)
R = []
b = []
hh = -h if self.predict_x0 else h
h_phi_1 = torch.expm1(hh) # h\phi_1(h) = e^h - 1
h_phi_k = h_phi_1 / hh - 1
factorial_i = 1
if self.variant == 'bh1':
B_h = hh
elif self.variant == 'bh2':
B_h = torch.expm1(hh)
else:
raise NotImplementedError()
for i in range(1, order + 1):
R.append(torch.pow(rks, i - 1))
b.append(h_phi_k * factorial_i / B_h)
factorial_i *= (i + 1)
h_phi_k = h_phi_k / hh - 1 / factorial_i
R = torch.stack(R)
b = torch.cat(b)
# now predictor
use_predictor = len(D1s) > 0 and x_t is None
if len(D1s) > 0:
D1s = torch.stack(D1s, dim=1) # (B, K)
if x_t is None:
# for order 2, we use a simplified version
if order == 2:
rhos_p = torch.tensor([0.5], device=b.device)
else:
rhos_p = torch.linalg.solve(R[:-1, :-1], b[:-1])
else:
D1s = None
if use_corrector:
#print('using corrector')
# for order 1, we use a simplified version
if order == 1:
rhos_c = torch.tensor([0.5], device=b.device)
else:
rhos_c = torch.linalg.solve(R, b)
model_t = None
if self.predict_x0:
x_t_ = (
sigma_t / sigma_prev_0 * x
- alpha_t * h_phi_1 * model_prev_0
)
if x_t is None:
if use_predictor:
pred_res = torch.einsum('k,bkchw->bchw', rhos_p, D1s)
else:
pred_res = 0
x_t = x_t_ - alpha_t * B_h * pred_res
if use_corrector:
model_t = self.model_fn(x_t, t)
if D1s is not None:
corr_res = torch.einsum('k,bkchw->bchw', rhos_c[:-1], D1s)
else:
corr_res = 0
D1_t = (model_t - model_prev_0)
x_t = x_t_ - alpha_t * B_h * (corr_res + rhos_c[-1] * D1_t)
else:
x_t_ = (
torch.exp(log_alpha_t - log_alpha_prev_0) * x
- sigma_t * h_phi_1 * model_prev_0
)
if x_t is None:
if use_predictor:
pred_res = torch.einsum('k,bkchw->bchw', rhos_p, D1s)
else:
pred_res = 0
x_t = x_t_ - sigma_t * B_h * pred_res
if use_corrector:
model_t = self.model_fn(x_t, t)
if D1s is not None:
corr_res = torch.einsum('k,bkchw->bchw', rhos_c[:-1], D1s)
else:
corr_res = 0
D1_t = (model_t - model_prev_0)
x_t = x_t_ - sigma_t * B_h * (corr_res + rhos_c[-1] * D1_t)
return x_t, model_t
def sample(self, x, steps=20, t_start=None, t_end=None, order=2, skip_type='time_uniform',
method='multistep', lower_order_final=True, denoise_to_zero=False, atol=0.0078, rtol=0.05, return_intermediate=False,
):
"""
Compute the sample at time `t_end` by UniPC, given the initial `x` at time `t_start`.
"""
t_0 = 1. / self.noise_schedule.total_N if t_end is None else t_end
t_T = self.noise_schedule.T if t_start is None else t_start
assert t_0 > 0 and t_T > 0, "Time range needs to be greater than 0. For discrete-time DPMs, it needs to be in [1 / N, 1], where N is the length of betas array"
if return_intermediate:
assert method in ['multistep', 'singlestep', 'singlestep_fixed'], "Cannot use adaptive solver when saving intermediate values"
if self.correcting_xt_fn is not None:
assert method in ['multistep', 'singlestep', 'singlestep_fixed'], "Cannot use adaptive solver when correcting_xt_fn is not None"
device = x.device
intermediates = []
with torch.no_grad():
if method == 'multistep':
assert steps >= order
timesteps = self.get_time_steps(skip_type=skip_type, t_T=t_T, t_0=t_0, N=steps, device=device)
assert timesteps.shape[0] - 1 == steps
# Init the initial values.
step = 0
t = timesteps[step]
t_prev_list = [t]
model_prev_list = [self.model_fn(x, t)]
if self.correcting_xt_fn is not None:
x = self.correcting_xt_fn(x, t, step)
if return_intermediate:
intermediates.append(x)
# Init the first `order` values by lower order multistep UniPC.
for step in range(1, order):
t = timesteps[step]
x, model_x = self.multistep_uni_pc_update(x, model_prev_list, t_prev_list, t, step, use_corrector=True)
if model_x is None:
model_x = self.model_fn(x, t)
if self.correcting_xt_fn is not None:
x = self.correcting_xt_fn(x, t, step)
if return_intermediate:
intermediates.append(x)
t_prev_list.append(t)
model_prev_list.append(model_x)
# Compute the remaining values by `order`-th order multistep DPM-Solver.
for step in range(order, steps + 1):
t = timesteps[step]
if lower_order_final:
step_order = min(order, steps + 1 - step)
else:
step_order = order
if step == steps:
#print('do not run corrector at the last step')
use_corrector = False
else:
use_corrector = True
x, model_x = self.multistep_uni_pc_update(x, model_prev_list, t_prev_list, t, step_order, use_corrector=use_corrector)
if self.correcting_xt_fn is not None:
x = self.correcting_xt_fn(x, t, step)
if return_intermediate:
intermediates.append(x)
for i in range(order - 1):
t_prev_list[i] = t_prev_list[i + 1]
model_prev_list[i] = model_prev_list[i + 1]
t_prev_list[-1] = t
# We do not need to evaluate the final model value.
if step < steps:
if model_x is None:
model_x = self.model_fn(x, t)
model_prev_list[-1] = model_x
else:
raise ValueError("Got wrong method {}".format(method))
if denoise_to_zero:
t = torch.ones((1,)).to(device) * t_0
x = self.denoise_to_zero_fn(x, t)
if self.correcting_xt_fn is not None:
x = self.correcting_xt_fn(x, t, step + 1)
if return_intermediate:
intermediates.append(x)
if return_intermediate:
return x, intermediates
else:
return x
#############################################################
# other utility functions
#############################################################
def interpolate_fn(x, xp, yp):
"""
A piecewise linear function y = f(x), using xp and yp as keypoints.
We implement f(x) in a differentiable way (i.e. applicable for autograd).
The function f(x) is well-defined for all x-axis. (For x beyond the bounds of xp, we use the outmost points of xp to define the linear function.)
Args:
x: PyTorch tensor with shape [N, C], where N is the batch size, C is the number of channels (we use C = 1 for DPM-Solver).
xp: PyTorch tensor with shape [C, K], where K is the number of keypoints.
yp: PyTorch tensor with shape [C, K].
Returns:
The function values f(x), with shape [N, C].
"""
N, K = x.shape[0], xp.shape[1]
all_x = torch.cat([x.unsqueeze(2), xp.unsqueeze(0).repeat((N, 1, 1))], dim=2)
sorted_all_x, x_indices = torch.sort(all_x, dim=2)
x_idx = torch.argmin(x_indices, dim=2)
cand_start_idx = x_idx - 1
start_idx = torch.where(
torch.eq(x_idx, 0),
torch.tensor(1, device=x.device),
torch.where(
torch.eq(x_idx, K), torch.tensor(K - 2, device=x.device), cand_start_idx,
),
)
end_idx = torch.where(torch.eq(start_idx, cand_start_idx), start_idx + 2, start_idx + 1)
start_x = torch.gather(sorted_all_x, dim=2, index=start_idx.unsqueeze(2)).squeeze(2)
end_x = torch.gather(sorted_all_x, dim=2, index=end_idx.unsqueeze(2)).squeeze(2)
start_idx2 = torch.where(
torch.eq(x_idx, 0),
torch.tensor(0, device=x.device),
torch.where(
torch.eq(x_idx, K), torch.tensor(K - 2, device=x.device), cand_start_idx,
),
)
y_positions_expanded = yp.unsqueeze(0).expand(N, -1, -1)
start_y = torch.gather(y_positions_expanded, dim=2, index=start_idx2.unsqueeze(2)).squeeze(2)
end_y = torch.gather(y_positions_expanded, dim=2, index=(start_idx2 + 1).unsqueeze(2)).squeeze(2)
cand = start_y + (x - start_x) * (end_y - start_y) / (end_x - start_x)
return cand
def expand_dims(v, dims):
"""
Expand the tensor `v` to the dim `dims`.
Args:
`v`: a PyTorch tensor with shape [N].
`dim`: a `int`.
Returns:
a PyTorch tensor with shape [N, 1, 1, ..., 1] and the total dimension is `dims`.
"""
return v[(...,) + (None,)*(dims - 1)]
inference/val_nsf_hifigan.py
+78
View File
@@ -0,0 +1,78 @@
import os
import sys
import librosa
import numpy as np
import resampy
import torch
import torchcrepe
import tqdm
from utils.binarizer_utils import get_pitch_parselmouth, get_mel_torch
from modules.vocoders.nsf_hifigan import NsfHifiGAN
from utils.infer_utils import save_wav
from utils.hparams import set_hparams, hparams
sys.argv = [
'inference/svs/ds_acoustic.py',
'--config',
'configs/acoustic.yaml',
]
def get_pitch(wav_data, mel, hparams, threshold=0.3):
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# crepe只支持16khz采样率，需要重采样
wav16k = resampy.resample(wav_data, hparams['audio_sample_rate'], 16000)
wav16k_torch = torch.FloatTensor(wav16k).unsqueeze(0).to(device)
# 频率范围
f0_min = 40
f0_max = 800
# 重采样后按照hopsize=80,也就是5ms一帧分析f0
f0, pd = torchcrepe.predict(wav16k_torch, 16000, 80, f0_min, f0_max, pad=True, model='full', batch_size=1024,
device=device, return_periodicity=True)
# 滤波，去掉静音，设置uv阈值，参考原仓库readme
pd = torchcrepe.filter.median(pd, 3)
pd = torchcrepe.threshold.Silence(-60.)(pd, wav16k_torch, 16000, 80)
f0 = torchcrepe.threshold.At(threshold)(f0, pd)
f0 = torchcrepe.filter.mean(f0, 3)
# 将nan频率（uv部分）转换为0频率
f0 = torch.where(torch.isnan(f0), torch.full_like(f0, 0), f0)
# 去掉0频率，并线性插值
nzindex = torch.nonzero(f0[0]).squeeze()
f0 = torch.index_select(f0[0], dim=0, index=nzindex).cpu().numpy()
time_org = 0.005 * nzindex.cpu().numpy()
time_frame = np.arange(len(mel)) * hparams['hop_size'] / hparams['audio_sample_rate']
f0 = np.interp(time_frame, time_org, f0, left=f0[0], right=f0[-1])
return f0
set_hparams()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
vocoder = NsfHifiGAN()
in_path = 'path/to/input/wavs'
out_path = 'path/to/output/wavs'
os.makedirs(out_path, exist_ok=True)
for filename in tqdm.tqdm(os.listdir(in_path)):
if not filename.endswith('.wav'):
continue
wav, _ = librosa.load(os.path.join(in_path, filename), sr=hparams['audio_sample_rate'], mono=True)
mel = get_mel_torch(
wav, hparams['audio_sample_rate'], num_mel_bins=hparams['audio_num_mel_bins'],
hop_size=hparams['hop_size'], win_size=hparams['win_size'], fft_size=hparams['fft_size'],
fmin=hparams['fmin'], fmax=hparams['fmax'],
device=device
)
f0, _ = get_pitch_parselmouth(
wav, samplerate=hparams['audio_sample_rate'], length=len(mel),
hop_size=hparams['hop_size']
)
wav_out = vocoder.spec2wav(mel, f0=f0)
save_wav(wav_out, os.path.join(out_path, filename), hparams['audio_sample_rate'])