chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
+18
@@ -0,0 +1,18 @@
|
||||
.idea
|
||||
*.pyc
|
||||
__pycache__/
|
||||
*.sh
|
||||
local_tools/
|
||||
*.ckpt
|
||||
*.pth
|
||||
infer_out/
|
||||
*.onnx
|
||||
/data/*
|
||||
!/data/.gitkeep
|
||||
/checkpoints/*
|
||||
!/checkpoints/.gitkeep
|
||||
/venv/
|
||||
/artifacts/
|
||||
|
||||
.vscode
|
||||
.ipynb_checkpoints/
|
||||
@@ -0,0 +1,201 @@
|
||||
Apache License
|
||||
Version 2.0, January 2004
|
||||
http://www.apache.org/licenses/
|
||||
|
||||
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
|
||||
|
||||
1. Definitions.
|
||||
|
||||
"License" shall mean the terms and conditions for use, reproduction,
|
||||
and distribution as defined by Sections 1 through 9 of this document.
|
||||
|
||||
"Licensor" shall mean the copyright owner or entity authorized by
|
||||
the copyright owner that is granting the License.
|
||||
|
||||
"Legal Entity" shall mean the union of the acting entity and all
|
||||
other entities that control, are controlled by, or are under common
|
||||
control with that entity. For the purposes of this definition,
|
||||
"control" means (i) the power, direct or indirect, to cause the
|
||||
direction or management of such entity, whether by contract or
|
||||
otherwise, or (ii) ownership of fifty percent (50%) or more of the
|
||||
outstanding shares, or (iii) beneficial ownership of such entity.
|
||||
|
||||
"You" (or "Your") shall mean an individual or Legal Entity
|
||||
exercising permissions granted by this License.
|
||||
|
||||
"Source" form shall mean the preferred form for making modifications,
|
||||
including but not limited to software source code, documentation
|
||||
source, and configuration files.
|
||||
|
||||
"Object" form shall mean any form resulting from mechanical
|
||||
transformation or translation of a Source form, including but
|
||||
not limited to compiled object code, generated documentation,
|
||||
and conversions to other media types.
|
||||
|
||||
"Work" shall mean the work of authorship, whether in Source or
|
||||
Object form, made available under the License, as indicated by a
|
||||
copyright notice that is included in or attached to the work
|
||||
(an example is provided in the Appendix below).
|
||||
|
||||
"Derivative Works" shall mean any work, whether in Source or Object
|
||||
form, that is based on (or derived from) the Work and for which the
|
||||
editorial revisions, annotations, elaborations, or other modifications
|
||||
represent, as a whole, an original work of authorship. For the purposes
|
||||
of this License, Derivative Works shall not include works that remain
|
||||
separable from, or merely link (or bind by name) to the interfaces of,
|
||||
the Work and Derivative Works thereof.
|
||||
|
||||
"Contribution" shall mean any work of authorship, including
|
||||
the original version of the Work and any modifications or additions
|
||||
to that Work or Derivative Works thereof, that is intentionally
|
||||
submitted to Licensor for inclusion in the Work by the copyright owner
|
||||
or by an individual or Legal Entity authorized to submit on behalf of
|
||||
the copyright owner. For the purposes of this definition, "submitted"
|
||||
means any form of electronic, verbal, or written communication sent
|
||||
to the Licensor or its representatives, including but not limited to
|
||||
communication on electronic mailing lists, source code control systems,
|
||||
and issue tracking systems that are managed by, or on behalf of, the
|
||||
Licensor for the purpose of discussing and improving the Work, but
|
||||
excluding communication that is conspicuously marked or otherwise
|
||||
designated in writing by the copyright owner as "Not a Contribution."
|
||||
|
||||
"Contributor" shall mean Licensor and any individual or Legal Entity
|
||||
on behalf of whom a Contribution has been received by Licensor and
|
||||
subsequently incorporated within the Work.
|
||||
|
||||
2. Grant of Copyright License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
copyright license to reproduce, prepare Derivative Works of,
|
||||
publicly display, publicly perform, sublicense, and distribute the
|
||||
Work and such Derivative Works in Source or Object form.
|
||||
|
||||
3. Grant of Patent License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
(except as stated in this section) patent license to make, have made,
|
||||
use, offer to sell, sell, import, and otherwise transfer the Work,
|
||||
where such license applies only to those patent claims licensable
|
||||
by such Contributor that are necessarily infringed by their
|
||||
Contribution(s) alone or by combination of their Contribution(s)
|
||||
with the Work to which such Contribution(s) was submitted. If You
|
||||
institute patent litigation against any entity (including a
|
||||
cross-claim or counterclaim in a lawsuit) alleging that the Work
|
||||
or a Contribution incorporated within the Work constitutes direct
|
||||
or contributory patent infringement, then any patent licenses
|
||||
granted to You under this License for that Work shall terminate
|
||||
as of the date such litigation is filed.
|
||||
|
||||
4. Redistribution. You may reproduce and distribute copies of the
|
||||
Work or Derivative Works thereof in any medium, with or without
|
||||
modifications, and in Source or Object form, provided that You
|
||||
meet the following conditions:
|
||||
|
||||
(a) You must give any other recipients of the Work or
|
||||
Derivative Works a copy of this License; and
|
||||
|
||||
(b) You must cause any modified files to carry prominent notices
|
||||
stating that You changed the files; and
|
||||
|
||||
(c) You must retain, in the Source form of any Derivative Works
|
||||
that You distribute, all copyright, patent, trademark, and
|
||||
attribution notices from the Source form of the Work,
|
||||
excluding those notices that do not pertain to any part of
|
||||
the Derivative Works; and
|
||||
|
||||
(d) If the Work includes a "NOTICE" text file as part of its
|
||||
distribution, then any Derivative Works that You distribute must
|
||||
include a readable copy of the attribution notices contained
|
||||
within such NOTICE file, excluding those notices that do not
|
||||
pertain to any part of the Derivative Works, in at least one
|
||||
of the following places: within a NOTICE text file distributed
|
||||
as part of the Derivative Works; within the Source form or
|
||||
documentation, if provided along with the Derivative Works; or,
|
||||
within a display generated by the Derivative Works, if and
|
||||
wherever such third-party notices normally appear. The contents
|
||||
of the NOTICE file are for informational purposes only and
|
||||
do not modify the License. You may add Your own attribution
|
||||
notices within Derivative Works that You distribute, alongside
|
||||
or as an addendum to the NOTICE text from the Work, provided
|
||||
that such additional attribution notices cannot be construed
|
||||
as modifying the License.
|
||||
|
||||
You may add Your own copyright statement to Your modifications and
|
||||
may provide additional or different license terms and conditions
|
||||
for use, reproduction, or distribution of Your modifications, or
|
||||
for any such Derivative Works as a whole, provided Your use,
|
||||
reproduction, and distribution of the Work otherwise complies with
|
||||
the conditions stated in this License.
|
||||
|
||||
5. Submission of Contributions. Unless You explicitly state otherwise,
|
||||
any Contribution intentionally submitted for inclusion in the Work
|
||||
by You to the Licensor shall be under the terms and conditions of
|
||||
this License, without any additional terms or conditions.
|
||||
Notwithstanding the above, nothing herein shall supersede or modify
|
||||
the terms of any separate license agreement you may have executed
|
||||
with Licensor regarding such Contributions.
|
||||
|
||||
6. Trademarks. This License does not grant permission to use the trade
|
||||
names, trademarks, service marks, or product names of the Licensor,
|
||||
except as required for reasonable and customary use in describing the
|
||||
origin of the Work and reproducing the content of the NOTICE file.
|
||||
|
||||
7. Disclaimer of Warranty. Unless required by applicable law or
|
||||
agreed to in writing, Licensor provides the Work (and each
|
||||
Contributor provides its Contributions) on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
|
||||
implied, including, without limitation, any warranties or conditions
|
||||
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
|
||||
PARTICULAR PURPOSE. You are solely responsible for determining the
|
||||
appropriateness of using or redistributing the Work and assume any
|
||||
risks associated with Your exercise of permissions under this License.
|
||||
|
||||
8. Limitation of Liability. In no event and under no legal theory,
|
||||
whether in tort (including negligence), contract, or otherwise,
|
||||
unless required by applicable law (such as deliberate and grossly
|
||||
negligent acts) or agreed to in writing, shall any Contributor be
|
||||
liable to You for damages, including any direct, indirect, special,
|
||||
incidental, or consequential damages of any character arising as a
|
||||
result of this License or out of the use or inability to use the
|
||||
Work (including but not limited to damages for loss of goodwill,
|
||||
work stoppage, computer failure or malfunction, or any and all
|
||||
other commercial damages or losses), even if such Contributor
|
||||
has been advised of the possibility of such damages.
|
||||
|
||||
9. Accepting Warranty or Additional Liability. While redistributing
|
||||
the Work or Derivative Works thereof, You may choose to offer,
|
||||
and charge a fee for, acceptance of support, warranty, indemnity,
|
||||
or other liability obligations and/or rights consistent with this
|
||||
License. However, in accepting such obligations, You may act only
|
||||
on Your own behalf and on Your sole responsibility, not on behalf
|
||||
of any other Contributor, and only if You agree to indemnify,
|
||||
defend, and hold each Contributor harmless for any liability
|
||||
incurred by, or claims asserted against, such Contributor by reason
|
||||
of your accepting any such warranty or additional liability.
|
||||
|
||||
END OF TERMS AND CONDITIONS
|
||||
|
||||
APPENDIX: How to apply the Apache License to your work.
|
||||
|
||||
To apply the Apache License to your work, attach the following
|
||||
boilerplate notice, with the fields enclosed by brackets "[]"
|
||||
replaced with your own identifying information. (Don't include
|
||||
the brackets!) The text should be enclosed in the appropriate
|
||||
comment syntax for the file format. We also recommend that a
|
||||
file or class name and description of purpose be included on the
|
||||
same "printed page" as the copyright notice for easier
|
||||
identification within third-party archives.
|
||||
|
||||
Copyright 2023 Team OpenVPI
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
@@ -0,0 +1,77 @@
|
||||
# DiffSinger (OpenVPI maintained version)
|
||||
|
||||
[](https://arxiv.org/abs/2105.02446)
|
||||
[](https://github.com/openvpi/DiffSinger/releases)
|
||||
[](https://www.bilibili.com/video/BV1be411N7JA/)
|
||||
[](https://github.com/openvpi/DiffSinger/blob/main/LICENSE)
|
||||
|
||||
This is a refactored and enhanced version of _DiffSinger: Singing Voice Synthesis via Shallow Diffusion Mechanism_ based on the original [paper](https://arxiv.org/abs/2105.02446) and [implementation](https://github.com/MoonInTheRiver/DiffSinger), which provides:
|
||||
|
||||
- Cleaner code structure: useless and redundant files are removed and the others are re-organized.
|
||||
- Better sound quality: the sampling rate of synthesized audio are adapted to 44.1 kHz instead of the original 24 kHz.
|
||||
- Higher fidelity: improved acoustic models and diffusion sampling acceleration algorithms are integrated.
|
||||
- More controllability: introduced variance models and parameters for prediction and control of pitch, energy, breathiness, etc.
|
||||
- Production compatibility: functionalities are designed to match the requirements of production deployment and the SVS communities.
|
||||
|
||||
| Overview | Variance Model | Acoustic Model |
|
||||
|:-------------------------------------------------------------------------------------:|:-------------------------------------------------------------------------------------:|:-------------------------------------------------------------------------------------:|
|
||||
| <img src="docs/resources/arch-overview.jpg" alt="arch-overview" style="zoom: 60%;" /> | <img src="docs/resources/arch-variance.jpg" alt="arch-variance" style="zoom: 50%;" /> | <img src="docs/resources/arch-acoustic.jpg" alt="arch-acoustic" style="zoom: 60%;" /> |
|
||||
|
||||
## User Guidance
|
||||
|
||||
> Navigation page including all tutorials and resources (Chinese): https://diffsinger.com
|
||||
|
||||
- **Installation & basic usages**: See [Getting Started](docs/GettingStarted.md)
|
||||
- **Dataset creation pipelines & tools**: See [MakeDiffSinger](https://github.com/openvpi/MakeDiffSinger)
|
||||
- **Best practices & tutorials**: See [Best Practices](docs/BestPractices.md)
|
||||
- **Editing configurations**: See [Configuration Schemas](docs/ConfigurationSchemas.md)
|
||||
- **Deployment & production**: [OpenUTAU](https://github.com/stakira/OpenUtau), [DiffScope (under development)](https://github.com/diffscope/diffscope-project)
|
||||
- **Communication groups**: [QQ Group](http://qm.qq.com/cgi-bin/qm/qr?_wv=1027&k=fibG_dxuPW5maUJwe9_ya5-zFcIwaoOR&authKey=ZgLCG5EqQVUGCID1nfKei8tCnlQHAmD9koxebFXv5WfUchhLwWxb52o1pimNai5A&noverify=0&group_code=907879266) (907879266), [Discord server](https://discord.gg/wwbu2JUMjj)
|
||||
|
||||
## Progress & Roadmap
|
||||
|
||||
- **Progress since we forked into this repository**: See [Releases](https://github.com/openvpi/DiffSinger/releases)
|
||||
- **Roadmap for future releases**: See [Project Board](https://github.com/orgs/openvpi/projects/1)
|
||||
- **Thoughts, proposals & ideas**: See [Discussions](https://github.com/openvpi/DiffSinger/discussions)
|
||||
|
||||
## Architecture & Algorithms
|
||||
|
||||
TBD
|
||||
|
||||
## Development Resources
|
||||
|
||||
TBD
|
||||
|
||||
## References
|
||||
|
||||
### Original Paper & Implementation
|
||||
|
||||
- Paper: [DiffSinger: Singing Voice Synthesis via Shallow Diffusion Mechanism](https://arxiv.org/abs/2105.02446)
|
||||
- Implementation: [MoonInTheRiver/DiffSinger](https://github.com/MoonInTheRiver/DiffSinger)
|
||||
|
||||
### Generative Models & Algorithms
|
||||
|
||||
- Denoising Diffusion Probabilistic Models (DDPM): [paper](https://arxiv.org/abs/2006.11239), [implementation](https://github.com/hojonathanho/diffusion)
|
||||
- [DDIM](https://arxiv.org/abs/2010.02502) for diffusion sampling acceleration
|
||||
- [PNDM](https://arxiv.org/abs/2202.09778) for diffusion sampling acceleration
|
||||
- [DPM-Solver++](https://github.com/LuChengTHU/dpm-solver) for diffusion sampling acceleration
|
||||
- [UniPC](https://github.com/wl-zhao/UniPC) for diffusion sampling acceleration
|
||||
- Rectified Flow (RF): [paper](https://arxiv.org/abs/2209.03003), [implementation](https://github.com/gnobitab/RectifiedFlow)
|
||||
|
||||
### Dependencies & Submodules
|
||||
|
||||
- [RoPE](https://github.com/lucidrains/rotary-embedding-torch) for transformer encoder
|
||||
- [Mix-LN](https://www.isca-archive.org/interspeech_2024/hwang24_interspeech.pdf) for cross-speaker/lingual learning
|
||||
- [HiFi-GAN](https://github.com/jik876/hifi-gan) and [NSF](https://github.com/nii-yamagishilab/project-NN-Pytorch-scripts/tree/master/project/01-nsf) for waveform reconstruction
|
||||
- [pc-ddsp](https://github.com/yxlllc/pc-ddsp) for waveform reconstruction
|
||||
- [RMVPE](https://github.com/Dream-High/RMVPE) and yxlllc's [fork](https://github.com/yxlllc/RMVPE) for pitch extraction
|
||||
- [Vocal Remover](https://github.com/tsurumeso/vocal-remover) and yxlllc's [fork](https://github.com/yxlllc/vocal-remover) for harmonic-noise separation
|
||||
|
||||
## Disclaimer
|
||||
|
||||
Any organization or individual is prohibited from using any functionalities included in this repository to generate someone's speech without his/her consent, including but not limited to government leaders, political figures, and celebrities. If you do not comply with this item, you could be in violation of copyright laws.
|
||||
|
||||
## License
|
||||
|
||||
This forked DiffSinger repository is licensed under the [Apache 2.0 License](LICENSE).
|
||||
|
||||
@@ -0,0 +1,7 @@
|
||||
# WeHub 来源说明
|
||||
|
||||
- 原始项目：`openvpi/DiffSinger`
|
||||
- 原始仓库：https://github.com/openvpi/DiffSinger
|
||||
- 导入方式：上游默认分支的最新快照
|
||||
- 原作者、版权和许可证信息以原始仓库及本仓库 LICENSE 为准
|
||||
- 本文件仅用于记录来源，不代表 WeHub 是原项目作者
|
||||
@@ -0,0 +1,92 @@
|
||||
from copy import deepcopy
|
||||
|
||||
import librosa
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from basics.base_augmentation import BaseAugmentation, require_same_keys
|
||||
from basics.base_pe import BasePE
|
||||
from modules.fastspeech.param_adaptor import VARIANCE_CHECKLIST
|
||||
from modules.fastspeech.tts_modules import LengthRegulator
|
||||
from utils.binarizer_utils import get_mel_torch, get_mel2ph_torch
|
||||
from utils.hparams import hparams
|
||||
from utils.infer_utils import resample_align_curve
|
||||
|
||||
|
||||
class SpectrogramStretchAugmentation(BaseAugmentation):
|
||||
"""
|
||||
This class contains methods for frequency-domain and time-domain stretching augmentation.
|
||||
"""
|
||||
|
||||
def __init__(self, data_dirs: list, augmentation_args: dict, pe: BasePE = None):
|
||||
super().__init__(data_dirs, augmentation_args)
|
||||
self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
||||
self.lr = LengthRegulator().to(self.device)
|
||||
self.pe = pe
|
||||
|
||||
@require_same_keys
|
||||
def process_item(self, item: dict, key_shift=0., speed=1., replace_spk_id=None) -> dict:
|
||||
aug_item = deepcopy(item)
|
||||
waveform, _ = librosa.load(aug_item['wav_fn'], sr=hparams['audio_sample_rate'], mono=True)
|
||||
mel = get_mel_torch(
|
||||
waveform, 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'],
|
||||
keyshift=key_shift, speed=speed, device=self.device
|
||||
)
|
||||
|
||||
aug_item['mel'] = mel
|
||||
|
||||
if speed != 1. or hparams['use_speed_embed']:
|
||||
aug_item['length'] = mel.shape[0]
|
||||
aug_item['speed'] = int(np.round(hparams['hop_size'] * speed)) / hparams['hop_size'] # real speed
|
||||
aug_item['seconds'] /= aug_item['speed']
|
||||
aug_item['ph_dur'] /= aug_item['speed']
|
||||
aug_item['mel2ph'] = get_mel2ph_torch(
|
||||
self.lr, torch.from_numpy(aug_item['ph_dur']), aug_item['length'], self.timestep, device=self.device
|
||||
).cpu().numpy()
|
||||
|
||||
f0, _ = self.pe.get_pitch(
|
||||
waveform, samplerate=hparams['audio_sample_rate'], length=aug_item['length'],
|
||||
hop_size=hparams['hop_size'], f0_min=hparams['f0_min'], f0_max=hparams['f0_max'],
|
||||
speed=speed, interp_uv=True
|
||||
)
|
||||
aug_item['f0'] = f0.astype(np.float32)
|
||||
|
||||
# NOTE: variance curves are directly resampled according to speed,
|
||||
# despite how frequency-domain features change after the augmentation.
|
||||
# For acoustic models, this can bring more (but not much) difficulty
|
||||
# to learn how variance curves affect the mel spectrograms, since
|
||||
# they must realize how the augmentation causes the mismatch.
|
||||
#
|
||||
# This is a simple way to combine augmentation and variances. However,
|
||||
# dealing variance curves like this will decrease the accuracy of
|
||||
# variance controls. In most situations, not being ~100% accurate
|
||||
# will not ruin the user experience. For example, it does not matter
|
||||
# if the energy does not exactly equal the RMS; it is just fine
|
||||
# as long as higher energy can bring higher loudness and strength.
|
||||
# The neural networks itself cannot be 100% accurate, though.
|
||||
#
|
||||
# There are yet other choices to simulate variance curves:
|
||||
# 1. Re-extract the features from resampled waveforms;
|
||||
# 2. Re-extract the features from re-constructed waveforms using
|
||||
# the transformed mel spectrograms through the vocoder.
|
||||
# But there are actually no perfect ways to make them all accurate
|
||||
# and stable.
|
||||
for v_name in VARIANCE_CHECKLIST:
|
||||
if v_name in item:
|
||||
aug_item[v_name] = resample_align_curve(
|
||||
aug_item[v_name],
|
||||
original_timestep=self.timestep,
|
||||
target_timestep=self.timestep * aug_item['speed'],
|
||||
align_length=aug_item['length']
|
||||
)
|
||||
|
||||
if key_shift != 0. or hparams['use_key_shift_embed']:
|
||||
if replace_spk_id is None:
|
||||
aug_item['key_shift'] = key_shift
|
||||
else:
|
||||
aug_item['spk_id'] = replace_spk_id
|
||||
aug_item['f0'] *= 2 ** (key_shift / 12)
|
||||
|
||||
return aug_item
|
||||
@@ -0,0 +1,28 @@
|
||||
from utils.hparams import hparams
|
||||
|
||||
|
||||
class BaseAugmentation:
|
||||
"""
|
||||
Base class for data augmentation.
|
||||
All methods of this class should be thread-safe.
|
||||
1. *process_item*:
|
||||
Apply augmentation to one piece of data.
|
||||
"""
|
||||
def __init__(self, data_dirs: list, augmentation_args: dict):
|
||||
self.raw_data_dirs = data_dirs
|
||||
self.augmentation_args = augmentation_args
|
||||
self.timestep = hparams['hop_size'] / hparams['audio_sample_rate']
|
||||
|
||||
def process_item(self, item: dict, **kwargs) -> dict:
|
||||
raise NotImplementedError()
|
||||
|
||||
|
||||
def require_same_keys(func):
|
||||
def run(*args, **kwargs):
|
||||
item: dict = args[1]
|
||||
res: dict = func(*args, **kwargs)
|
||||
assert set(item.keys()) == set(res.keys()), 'Item keys mismatch after augmentation.\n' \
|
||||
f'Before: {sorted(item.keys())}\n' \
|
||||
f'After: {sorted(res.keys())}'
|
||||
return res
|
||||
return run
|
||||
@@ -0,0 +1,386 @@
|
||||
import json
|
||||
import pathlib
|
||||
import pickle
|
||||
import random
|
||||
import shutil
|
||||
import warnings
|
||||
from copy import deepcopy
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from tqdm import tqdm
|
||||
|
||||
from utils.hparams import hparams
|
||||
from utils.indexed_datasets import IndexedDatasetBuilder
|
||||
from utils.multiprocess_utils import chunked_multiprocess_run
|
||||
from utils.phoneme_utils import load_phoneme_dictionary
|
||||
from utils.plot import distribution_to_figure
|
||||
|
||||
|
||||
class BinarizationError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class BaseBinarizer:
|
||||
"""
|
||||
Base class for data processing.
|
||||
1. *process* and *process_data_split*:
|
||||
process entire data, generate the train-test split (support parallel processing);
|
||||
2. *process_item*:
|
||||
process singe piece of data;
|
||||
3. *get_pitch*:
|
||||
infer the pitch using some algorithm;
|
||||
4. *get_align*:
|
||||
get the alignment using 'mel2ph' format (see https://arxiv.org/abs/1905.09263).
|
||||
5. phoneme encoder, voice encoder, etc.
|
||||
|
||||
Subclasses should define:
|
||||
1. *load_metadata*:
|
||||
how to read multiple datasets from files;
|
||||
2. *train_item_names*, *valid_item_names*, *test_item_names*:
|
||||
how to split the dataset;
|
||||
3. load_ph_set:
|
||||
the phoneme set.
|
||||
"""
|
||||
|
||||
def __init__(self, datasets=None, data_attrs=None):
|
||||
if datasets is None:
|
||||
datasets = hparams['datasets']
|
||||
self.datasets = datasets
|
||||
self.raw_data_dirs = [pathlib.Path(ds['raw_data_dir']) for ds in self.datasets]
|
||||
self.binary_data_dir = pathlib.Path(hparams['binary_data_dir'])
|
||||
self.data_attrs = [] if data_attrs is None else data_attrs
|
||||
|
||||
self.binarization_args = hparams['binarization_args']
|
||||
self.augmentation_args = hparams.get('augmentation_args', {})
|
||||
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||||
|
||||
self.spk_map = {}
|
||||
self.spk_ids = None
|
||||
self.build_spk_map()
|
||||
|
||||
self.lang_map = {}
|
||||
self.dictionaries = hparams['dictionaries']
|
||||
self.build_lang_map()
|
||||
|
||||
self.items = {}
|
||||
self.item_names: list = None
|
||||
self._train_item_names: list = None
|
||||
self._valid_item_names: list = None
|
||||
|
||||
self.phoneme_dictionary = load_phoneme_dictionary()
|
||||
self.timestep = hparams['hop_size'] / hparams['audio_sample_rate']
|
||||
|
||||
def build_spk_map(self):
|
||||
spk_ids = [ds.get('spk_id') for ds in self.datasets]
|
||||
assigned_spk_ids = {spk_id for spk_id in spk_ids if spk_id is not None}
|
||||
idx = 0
|
||||
for i in range(len(spk_ids)):
|
||||
if spk_ids[i] is not None:
|
||||
continue
|
||||
while idx in assigned_spk_ids:
|
||||
idx += 1
|
||||
spk_ids[i] = idx
|
||||
assigned_spk_ids.add(idx)
|
||||
assert max(spk_ids) < hparams['num_spk'], \
|
||||
f'Index in spk_id sequence {spk_ids} is out of range. All values should be smaller than num_spk.'
|
||||
|
||||
for spk_id, dataset in zip(spk_ids, self.datasets):
|
||||
spk_name = dataset['speaker']
|
||||
if spk_name in self.spk_map and self.spk_map[spk_name] != spk_id:
|
||||
raise ValueError(f'Invalid speaker ID assignment. Name \'{spk_name}\' is assigned '
|
||||
f'with different speaker IDs: {self.spk_map[spk_name]} and {spk_id}.')
|
||||
self.spk_map[spk_name] = spk_id
|
||||
self.spk_ids = spk_ids
|
||||
|
||||
print("| spk_map: ", self.spk_map)
|
||||
|
||||
def build_lang_map(self):
|
||||
assert len(self.dictionaries.keys()) <= hparams['num_lang'], \
|
||||
'Number of languages must not be greater than num_lang!'
|
||||
for dataset in self.datasets:
|
||||
assert dataset['language'] in self.dictionaries, f'Unrecognized language name: {dataset["language"]}'
|
||||
|
||||
for lang_id, lang_name in enumerate(sorted(self.dictionaries.keys()), start=1):
|
||||
self.lang_map[lang_name] = lang_id
|
||||
|
||||
print("| lang_map: ", self.lang_map)
|
||||
|
||||
def load_meta_data(self, raw_data_dir: pathlib.Path, ds_id, spk, lang) -> dict:
|
||||
raise NotImplementedError()
|
||||
|
||||
def split_train_valid_set(self, prefixes: list):
|
||||
"""
|
||||
Split the dataset into training set and validation set.
|
||||
:return: train_item_names, valid_item_names
|
||||
"""
|
||||
prefixes = {str(pr): 1 for pr in prefixes}
|
||||
valid_item_names = {}
|
||||
# Add prefixes that specified speaker index and matches exactly item name to test set
|
||||
for prefix in deepcopy(prefixes):
|
||||
if prefix in self.item_names:
|
||||
valid_item_names[prefix] = 1
|
||||
prefixes.pop(prefix)
|
||||
# Add prefixes that exactly matches item name without speaker id to test set
|
||||
for prefix in deepcopy(prefixes):
|
||||
matched = False
|
||||
for name in self.item_names:
|
||||
if name.split(':')[-1] == prefix:
|
||||
valid_item_names[name] = 1
|
||||
matched = True
|
||||
if matched:
|
||||
prefixes.pop(prefix)
|
||||
# Add names with one of the remaining prefixes to test set
|
||||
for prefix in deepcopy(prefixes):
|
||||
matched = False
|
||||
for name in self.item_names:
|
||||
if name.startswith(prefix):
|
||||
valid_item_names[name] = 1
|
||||
matched = True
|
||||
if matched:
|
||||
prefixes.pop(prefix)
|
||||
for prefix in deepcopy(prefixes):
|
||||
matched = False
|
||||
for name in self.item_names:
|
||||
if name.split(':')[-1].startswith(prefix):
|
||||
valid_item_names[name] = 1
|
||||
matched = True
|
||||
if matched:
|
||||
prefixes.pop(prefix)
|
||||
|
||||
if len(prefixes) != 0:
|
||||
warnings.warn(
|
||||
f'The following rules in test_prefixes have no matching names in the dataset: {", ".join(prefixes.keys())}',
|
||||
category=UserWarning
|
||||
)
|
||||
warnings.filterwarnings('default')
|
||||
|
||||
valid_item_names = list(valid_item_names.keys())
|
||||
assert len(valid_item_names) > 0, 'Validation set is empty!'
|
||||
train_item_names = [x for x in self.item_names if x not in set(valid_item_names)]
|
||||
assert len(train_item_names) > 0, 'Training set is empty!'
|
||||
|
||||
return train_item_names, valid_item_names
|
||||
|
||||
@property
|
||||
def train_item_names(self):
|
||||
return self._train_item_names
|
||||
|
||||
@property
|
||||
def valid_item_names(self):
|
||||
return self._valid_item_names
|
||||
|
||||
def meta_data_iterator(self, prefix):
|
||||
if prefix == 'train':
|
||||
item_names = self.train_item_names
|
||||
else:
|
||||
item_names = self.valid_item_names
|
||||
for item_name in item_names:
|
||||
meta_data = self.items[item_name]
|
||||
yield item_name, meta_data
|
||||
|
||||
def process(self):
|
||||
# load each dataset
|
||||
test_prefixes = []
|
||||
for ds_id, dataset in enumerate(self.datasets):
|
||||
items = self.load_meta_data(
|
||||
pathlib.Path(dataset['raw_data_dir']),
|
||||
ds_id=ds_id, spk=dataset['speaker'], lang=dataset['language']
|
||||
)
|
||||
self.items.update(items)
|
||||
test_prefixes.extend(
|
||||
f'{ds_id}:{prefix}'
|
||||
for prefix in dataset.get('test_prefixes', [])
|
||||
)
|
||||
self.item_names = sorted(list(self.items.keys()))
|
||||
self._train_item_names, self._valid_item_names = self.split_train_valid_set(test_prefixes)
|
||||
|
||||
if self.binarization_args['shuffle']:
|
||||
random.shuffle(self.item_names)
|
||||
|
||||
self.binary_data_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Copy spk_map, lang_map and dictionary to binary data dir
|
||||
spk_map_fn = self.binary_data_dir / 'spk_map.json'
|
||||
with open(spk_map_fn, 'w', encoding='utf-8') as f:
|
||||
json.dump(self.spk_map, f, ensure_ascii=False)
|
||||
lang_map_fn = self.binary_data_dir / 'lang_map.json'
|
||||
with open(lang_map_fn, 'w', encoding='utf-8') as f:
|
||||
json.dump(self.lang_map, f, ensure_ascii=False)
|
||||
for lang, dict_path in hparams['dictionaries'].items():
|
||||
shutil.copy(dict_path, self.binary_data_dir / f'dictionary-{lang}.txt')
|
||||
self.check_coverage()
|
||||
|
||||
# Process valid set and train set
|
||||
try:
|
||||
self.process_dataset('valid')
|
||||
self.process_dataset(
|
||||
'train',
|
||||
num_workers=int(self.binarization_args['num_workers']),
|
||||
apply_augmentation=any(args['enabled'] for args in self.augmentation_args.values())
|
||||
)
|
||||
except KeyboardInterrupt:
|
||||
exit(-1)
|
||||
|
||||
def check_coverage(self):
|
||||
# Group by phonemes in the dictionary.
|
||||
ph_idx_required = set(range(1, len(self.phoneme_dictionary)))
|
||||
ph_idx_occurred = set()
|
||||
ph_idx_count_map = {
|
||||
idx: 0
|
||||
for idx in ph_idx_required
|
||||
}
|
||||
|
||||
# Load and count those phones that appear in the actual data
|
||||
for item_name in self.items:
|
||||
ph_idx_occurred.update(self.items[item_name]['ph_seq'])
|
||||
for idx in self.items[item_name]['ph_seq']:
|
||||
ph_idx_count_map[idx] += 1
|
||||
ph_count_map = {
|
||||
self.phoneme_dictionary.decode_one(idx, scalar=False): count
|
||||
for idx, count in ph_idx_count_map.items()
|
||||
}
|
||||
|
||||
def display_phoneme(phoneme):
|
||||
if isinstance(phoneme, tuple):
|
||||
return f'({", ".join(phoneme)})'
|
||||
return phoneme
|
||||
|
||||
print('===== Phoneme Distribution Summary =====')
|
||||
keys = sorted(ph_count_map.keys(), key=lambda v: v[0] if isinstance(v, tuple) else v)
|
||||
for i, key in enumerate(keys):
|
||||
if i == len(ph_count_map) - 1:
|
||||
end = '\n'
|
||||
elif i % 10 == 9:
|
||||
end = ',\n'
|
||||
else:
|
||||
end = ', '
|
||||
key_disp = display_phoneme(key)
|
||||
print(f'{key_disp}: {ph_count_map[key]}', end=end)
|
||||
|
||||
# Draw graph.
|
||||
xs = [display_phoneme(k) for k in keys]
|
||||
ys = [ph_count_map[k] for k in keys]
|
||||
plt = distribution_to_figure(
|
||||
title='Phoneme Distribution Summary',
|
||||
x_label='Phoneme', y_label='Number of occurrences',
|
||||
items=xs, values=ys, rotate=len(self.dictionaries) > 1
|
||||
)
|
||||
filename = self.binary_data_dir / 'phoneme_distribution.jpg'
|
||||
plt.savefig(fname=filename,
|
||||
bbox_inches='tight',
|
||||
pad_inches=0.25)
|
||||
print(f'| save summary to \'{filename}\'')
|
||||
|
||||
# Check unrecognizable or missing phonemes
|
||||
if ph_idx_occurred != ph_idx_required:
|
||||
missing_phones = sorted({
|
||||
self.phoneme_dictionary.decode_one(idx, scalar=False)
|
||||
for idx in ph_idx_required.difference(ph_idx_occurred)
|
||||
}, key=lambda v: v[0] if isinstance(v, tuple) else v)
|
||||
raise BinarizationError(
|
||||
f'The following phonemes are not covered in transcriptions: {missing_phones}'
|
||||
)
|
||||
|
||||
def process_dataset(self, prefix, num_workers=0, apply_augmentation=False):
|
||||
args = []
|
||||
builder = IndexedDatasetBuilder(self.binary_data_dir, prefix=prefix, allowed_attr=self.data_attrs)
|
||||
total_sec = {k: 0.0 for k in self.spk_map}
|
||||
total_raw_sec = {k: 0.0 for k in self.spk_map}
|
||||
extra_info = {'names': {}, 'ph_texts': {}, 'spk_ids': {}, 'spk_names': {}, 'lengths': {}}
|
||||
max_no = -1
|
||||
|
||||
for item_name, meta_data in self.meta_data_iterator(prefix):
|
||||
args.append([item_name, meta_data, self.binarization_args])
|
||||
|
||||
aug_map = self.arrange_data_augmentation(self.meta_data_iterator(prefix)) if apply_augmentation else {}
|
||||
|
||||
def postprocess(_item):
|
||||
nonlocal total_sec, total_raw_sec, extra_info, max_no
|
||||
if _item is None:
|
||||
return
|
||||
item_no = builder.add_item(_item)
|
||||
max_no = max(max_no, item_no)
|
||||
for k, v in _item.items():
|
||||
if isinstance(v, np.ndarray):
|
||||
if k not in extra_info:
|
||||
extra_info[k] = {}
|
||||
extra_info[k][item_no] = v.shape[0]
|
||||
extra_info['names'][item_no] = _item['name'].split(':', 1)[-1]
|
||||
extra_info['ph_texts'][item_no] = _item['ph_text']
|
||||
extra_info['spk_ids'][item_no] = _item['spk_id']
|
||||
extra_info['spk_names'][item_no] = _item['spk_name']
|
||||
extra_info['lengths'][item_no] = _item['length']
|
||||
total_raw_sec[_item['spk_name']] += _item['seconds']
|
||||
total_sec[_item['spk_name']] += _item['seconds']
|
||||
|
||||
for task in aug_map.get(_item['name'], []):
|
||||
aug_item = task['func'](_item, **task['kwargs'])
|
||||
aug_item_no = builder.add_item(aug_item)
|
||||
max_no = max(max_no, aug_item_no)
|
||||
for k, v in aug_item.items():
|
||||
if isinstance(v, np.ndarray):
|
||||
if k not in extra_info:
|
||||
extra_info[k] = {}
|
||||
extra_info[k][aug_item_no] = v.shape[0]
|
||||
extra_info['names'][aug_item_no] = aug_item['name'].split(':', 1)[-1]
|
||||
extra_info['ph_texts'][aug_item_no] = aug_item['ph_text']
|
||||
extra_info['spk_ids'][aug_item_no] = aug_item['spk_id']
|
||||
extra_info['spk_names'][aug_item_no] = aug_item['spk_name']
|
||||
extra_info['lengths'][aug_item_no] = aug_item['length']
|
||||
total_sec[aug_item['spk_name']] += aug_item['seconds']
|
||||
|
||||
try:
|
||||
if num_workers > 0:
|
||||
# code for parallel processing
|
||||
for item in tqdm(
|
||||
chunked_multiprocess_run(self.process_item, args, num_workers=num_workers),
|
||||
total=len(list(self.meta_data_iterator(prefix)))
|
||||
):
|
||||
postprocess(item)
|
||||
else:
|
||||
# code for single cpu processing
|
||||
for a in tqdm(args):
|
||||
item = self.process_item(*a)
|
||||
postprocess(item)
|
||||
for k in extra_info:
|
||||
assert set(extra_info[k]) == set(range(max_no + 1)), f'Item numbering is not consecutive.'
|
||||
extra_info[k] = list(map(lambda x: x[1], sorted(extra_info[k].items(), key=lambda x: x[0])))
|
||||
except KeyboardInterrupt:
|
||||
builder.finalize()
|
||||
raise
|
||||
|
||||
builder.finalize()
|
||||
if prefix == "train":
|
||||
extra_info.pop("names")
|
||||
extra_info.pop('ph_texts')
|
||||
extra_info.pop("spk_names")
|
||||
with open(self.binary_data_dir / f"{prefix}.meta", "wb") as f:
|
||||
# noinspection PyTypeChecker
|
||||
pickle.dump(extra_info, f)
|
||||
if apply_augmentation:
|
||||
print(f"| {prefix} total duration (before augmentation): {sum(total_raw_sec.values()):.2f}s")
|
||||
print(
|
||||
f"| {prefix} respective duration (before augmentation): "
|
||||
+ ', '.join(f'{k}={v:.2f}s' for k, v in total_raw_sec.items())
|
||||
)
|
||||
print(
|
||||
f"| {prefix} total duration (after augmentation): "
|
||||
f"{sum(total_sec.values()):.2f}s ({sum(total_sec.values()) / sum(total_raw_sec.values()):.2f}x)"
|
||||
)
|
||||
print(
|
||||
f"| {prefix} respective duration (after augmentation): "
|
||||
+ ', '.join(f'{k}={v:.2f}s' for k, v in total_sec.items())
|
||||
)
|
||||
else:
|
||||
print(f"| {prefix} total duration: {sum(total_raw_sec.values()):.2f}s")
|
||||
print(f"| {prefix} respective duration: " + ', '.join(f'{k}={v:.2f}s' for k, v in total_raw_sec.items()))
|
||||
|
||||
def arrange_data_augmentation(self, data_iterator):
|
||||
"""
|
||||
Code for all types of data augmentation should be added here.
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
def process_item(self, item_name, meta_data, binarization_args):
|
||||
raise NotImplementedError()
|
||||
@@ -0,0 +1,58 @@
|
||||
import os
|
||||
import pickle
|
||||
|
||||
import torch
|
||||
from torch.utils.data import Dataset
|
||||
|
||||
from utils.hparams import hparams
|
||||
from utils.indexed_datasets import IndexedDataset
|
||||
|
||||
|
||||
class BaseDataset(Dataset):
|
||||
"""
|
||||
Base class for datasets.
|
||||
1. *sizes*:
|
||||
clipped length if "max_frames" is set;
|
||||
2. *num_frames*:
|
||||
unclipped length.
|
||||
|
||||
Subclasses should define:
|
||||
1. *collate*:
|
||||
take the longest data, pad other data to the same length;
|
||||
2. *__getitem__*:
|
||||
the index function.
|
||||
"""
|
||||
|
||||
def __init__(self, prefix, size_key='lengths', preload=False):
|
||||
super().__init__()
|
||||
self.prefix = prefix
|
||||
self.data_dir = hparams['binary_data_dir']
|
||||
with open(os.path.join(self.data_dir, f'{self.prefix}.meta'), 'rb') as f:
|
||||
self.metadata = pickle.load(f)
|
||||
self.sizes = self.metadata[size_key]
|
||||
self._indexed_ds = IndexedDataset(self.data_dir, self.prefix)
|
||||
if preload:
|
||||
self.indexed_ds = [self._indexed_ds[i] for i in range(len(self._indexed_ds))]
|
||||
del self._indexed_ds
|
||||
else:
|
||||
self.indexed_ds = self._indexed_ds
|
||||
|
||||
def __getitem__(self, index):
|
||||
return {'_idx': index, **self.indexed_ds[index]}
|
||||
|
||||
def __len__(self):
|
||||
return len(self.sizes)
|
||||
|
||||
def num_frames(self, index):
|
||||
return self.sizes[index]
|
||||
|
||||
def size(self, index):
|
||||
"""Return an example's size as a float or tuple. This value is used when
|
||||
filtering a dataset with ``--max-positions``."""
|
||||
return self.sizes[index]
|
||||
|
||||
def collater(self, samples):
|
||||
return {
|
||||
'size': len(samples),
|
||||
'indices': torch.LongTensor([s['_idx'] for s in samples])
|
||||
}
|
||||
@@ -0,0 +1,86 @@
|
||||
import json
|
||||
import pathlib
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
from typing import Union
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from utils.hparams import hparams
|
||||
|
||||
|
||||
class BaseExporter:
|
||||
def __init__(
|
||||
self,
|
||||
device: Union[str, torch.device] = None,
|
||||
cache_dir: Path = None,
|
||||
**kwargs
|
||||
):
|
||||
self.device = device if device is not None else torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||||
self.cache_dir: Path = cache_dir.resolve() if cache_dir is not None \
|
||||
else Path(__file__).parent.parent / 'deployment' / 'cache'
|
||||
self.cache_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# noinspection PyMethodMayBeStatic
|
||||
def build_spk_map(self) -> dict:
|
||||
if hparams['use_spk_id']:
|
||||
with open(Path(hparams['work_dir']) / 'spk_map.json', 'r', encoding='utf8') as f:
|
||||
spk_map = json.load(f)
|
||||
assert isinstance(spk_map, dict) and len(spk_map) > 0, 'Invalid or empty speaker map!'
|
||||
assert len(spk_map) == len(set(spk_map.values())), 'Duplicate speaker id in speaker map!'
|
||||
return spk_map
|
||||
else:
|
||||
return {}
|
||||
|
||||
# noinspection PyMethodMayBeStatic
|
||||
def build_lang_map(self) -> dict:
|
||||
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:
|
||||
lang_map = json.load(f)
|
||||
assert isinstance(lang_map, dict) and len(lang_map) > 0, 'Invalid or empty language map!'
|
||||
assert len(lang_map) == len(set(lang_map.values())), 'Duplicate language id in language map!'
|
||||
return lang_map
|
||||
else:
|
||||
return {}
|
||||
|
||||
def build_model(self) -> nn.Module:
|
||||
"""
|
||||
Creates an instance of nn.Module and load its state dict on the target device.
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
def export_model(self, path: Path):
|
||||
"""
|
||||
Exports the model to ONNX format.
|
||||
:param path: the target model path
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
# noinspection PyMethodMayBeStatic
|
||||
def export_dictionaries(self, path: Path):
|
||||
dicts = hparams.get('dictionaries')
|
||||
if dicts is not None:
|
||||
for lang in dicts.keys():
|
||||
fn = f'dictionary-{lang}.txt'
|
||||
shutil.copy(pathlib.Path(hparams['work_dir']) / fn, path)
|
||||
print(f'| export dictionary => {path / fn}')
|
||||
else:
|
||||
fn = 'dictionary.txt'
|
||||
shutil.copy(pathlib.Path(hparams['work_dir']) / fn, path)
|
||||
print(f'| export dictionary => {path / fn}')
|
||||
|
||||
def export_attachments(self, path: Path):
|
||||
"""
|
||||
Exports related files and configs (e.g. the dictionary) to the target directory.
|
||||
:param path: the target directory
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
def export(self, path: Path):
|
||||
"""
|
||||
Exports all the artifacts to the target directory.
|
||||
:param path: the target directory
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
@@ -0,0 +1,18 @@
|
||||
from torch import nn
|
||||
|
||||
|
||||
class CategorizedModule(nn.Module):
|
||||
@property
|
||||
def category(self):
|
||||
raise NotImplementedError()
|
||||
|
||||
def check_category(self, category):
|
||||
if category is None:
|
||||
raise RuntimeError('Category is not specified in this checkpoint.\n'
|
||||
'If this is a checkpoint in the old format, please consider '
|
||||
'migrating it to the new format via the following command:\n'
|
||||
'python scripts/migrate.py ckpt <INPUT_CKPT> <OUTPUT_CKPT>')
|
||||
elif category != self.category:
|
||||
raise RuntimeError('Category mismatches!\n'
|
||||
f'This checkpoint is of the category \'{category}\', '
|
||||
f'but a checkpoint of category \'{self.category}\' is required.')
|
||||
@@ -0,0 +1,7 @@
|
||||
class BasePE:
|
||||
def get_pitch(
|
||||
self, waveform, samplerate, length,
|
||||
*, hop_size, f0_min=65, f0_max=1100,
|
||||
speed=1, interp_uv=False
|
||||
):
|
||||
raise NotImplementedError()
|
||||
@@ -0,0 +1,136 @@
|
||||
# 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()
|
||||
@@ -0,0 +1,514 @@
|
||||
import logging
|
||||
import os
|
||||
import pathlib
|
||||
import shutil
|
||||
import sys
|
||||
from typing import Dict
|
||||
|
||||
import matplotlib
|
||||
|
||||
import utils
|
||||
|
||||
matplotlib.use('Agg')
|
||||
|
||||
import torch.utils.data
|
||||
from torchmetrics import Metric, MeanMetric
|
||||
import lightning.pytorch as pl
|
||||
from lightning.pytorch.utilities.rank_zero import rank_zero_debug, rank_zero_info, rank_zero_only
|
||||
|
||||
from basics.base_module import CategorizedModule
|
||||
from utils.hparams import hparams
|
||||
from utils.training_utils import (
|
||||
DsModelCheckpoint, DsTQDMProgressBar,
|
||||
DsBatchSampler, DsTensorBoardLogger,
|
||||
get_latest_checkpoint_path, get_strategy
|
||||
)
|
||||
from utils.phoneme_utils import load_phoneme_dictionary
|
||||
|
||||
torch.multiprocessing.set_sharing_strategy(os.getenv('TORCH_SHARE_STRATEGY', 'file_system'))
|
||||
|
||||
log_format = '%(asctime)s %(message)s'
|
||||
logging.basicConfig(stream=sys.stdout, level=logging.INFO,
|
||||
format=log_format, datefmt='%m/%d %I:%M:%S %p')
|
||||
|
||||
|
||||
class BaseTask(pl.LightningModule):
|
||||
"""
|
||||
Base class for training tasks.
|
||||
1. *load_ckpt*:
|
||||
load checkpoint;
|
||||
2. *training_step*:
|
||||
record and log the loss;
|
||||
3. *optimizer_step*:
|
||||
run backwards step;
|
||||
4. *start*:
|
||||
load training configs, backup code, log to tensorboard, start training;
|
||||
5. *configure_ddp* and *init_ddp_connection*:
|
||||
start parallel training.
|
||||
|
||||
Subclasses should define:
|
||||
1. *build_model*, *build_optimizer*, *build_scheduler*:
|
||||
how to build the model, the optimizer and the training scheduler;
|
||||
2. *_training_step*:
|
||||
one training step of the model;
|
||||
3. *on_validation_end* and *_on_validation_end*:
|
||||
postprocess the validation output.
|
||||
"""
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self.max_batch_frames = hparams['max_batch_frames']
|
||||
self.max_batch_size = hparams['max_batch_size']
|
||||
self.max_val_batch_frames = hparams['max_val_batch_frames']
|
||||
if self.max_val_batch_frames == -1:
|
||||
hparams['max_val_batch_frames'] = self.max_val_batch_frames = self.max_batch_frames
|
||||
self.max_val_batch_size = hparams['max_val_batch_size']
|
||||
if self.max_val_batch_size == -1:
|
||||
hparams['max_val_batch_size'] = self.max_val_batch_size = self.max_batch_size
|
||||
|
||||
self.training_sampler = None
|
||||
self.skip_immediate_validation = False
|
||||
self.skip_immediate_ckpt_save = False
|
||||
|
||||
self.phoneme_dictionary = load_phoneme_dictionary()
|
||||
self.build_model()
|
||||
|
||||
self.valid_losses: Dict[str, Metric] = {}
|
||||
self.valid_metrics: Dict[str, Metric] = {}
|
||||
|
||||
def _finish_init(self):
|
||||
self.register_validation_loss('total_loss')
|
||||
self.build_losses_and_metrics()
|
||||
assert len(self.valid_losses) > 0, "No validation loss registered. Please check your configuration file."
|
||||
|
||||
###########
|
||||
# Training, validation and testing
|
||||
###########
|
||||
def setup(self, stage):
|
||||
self.train_dataset = self.dataset_cls('train')
|
||||
self.valid_dataset = self.dataset_cls('valid')
|
||||
self.num_replicas = (self.trainer.distributed_sampler_kwargs or {}).get('num_replicas', 1)
|
||||
|
||||
def get_need_freeze_state_dict_key(self, model_state_dict) -> list:
|
||||
key_list = []
|
||||
for i in hparams['frozen_params']:
|
||||
for j in model_state_dict:
|
||||
if j.startswith(i):
|
||||
key_list.append(j)
|
||||
return list(set(key_list))
|
||||
|
||||
def freeze_params(self) -> None:
|
||||
model_state_dict = self.state_dict().keys()
|
||||
freeze_key = self.get_need_freeze_state_dict_key(model_state_dict=model_state_dict)
|
||||
|
||||
for i in freeze_key:
|
||||
params=self.get_parameter(i)
|
||||
|
||||
params.requires_grad = False
|
||||
|
||||
def unfreeze_all_params(self) -> None:
|
||||
for i in self.model.parameters():
|
||||
i.requires_grad = True
|
||||
|
||||
def load_finetune_ckpt(
|
||||
self, state_dict
|
||||
) -> None:
|
||||
adapt_shapes = hparams['finetune_strict_shapes']
|
||||
if not adapt_shapes:
|
||||
cur_model_state_dict = self.state_dict()
|
||||
unmatched_keys = []
|
||||
for key, param in state_dict.items():
|
||||
if key in cur_model_state_dict:
|
||||
new_param = cur_model_state_dict[key]
|
||||
if new_param.shape != param.shape:
|
||||
unmatched_keys.append(key)
|
||||
print('| Unmatched keys: ', key, new_param.shape, param.shape)
|
||||
for key in unmatched_keys:
|
||||
del state_dict[key]
|
||||
self.load_state_dict(state_dict, strict=False)
|
||||
|
||||
def load_pre_train_model(self):
|
||||
pre_train_ckpt_path = hparams['finetune_ckpt_path']
|
||||
blacklist = hparams['finetune_ignored_params']
|
||||
# whitelist=hparams['pre_train_whitelist']
|
||||
if blacklist is None:
|
||||
blacklist = []
|
||||
# if whitelist is None:
|
||||
# raise RuntimeError("")
|
||||
|
||||
if pre_train_ckpt_path is not None:
|
||||
ckpt = torch.load(pre_train_ckpt_path)
|
||||
# if ckpt.get('category') is None:
|
||||
# raise RuntimeError("")
|
||||
|
||||
if isinstance(self.model, CategorizedModule):
|
||||
self.model.check_category(ckpt.get('category'))
|
||||
|
||||
state_dict = {}
|
||||
for i in ckpt['state_dict']:
|
||||
# if 'diffusion' in i:
|
||||
# if i in rrrr:
|
||||
# continue
|
||||
skip = False
|
||||
for b in blacklist:
|
||||
if i.startswith(b):
|
||||
skip = True
|
||||
break
|
||||
|
||||
if skip:
|
||||
continue
|
||||
|
||||
state_dict[i] = ckpt['state_dict'][i]
|
||||
print(i)
|
||||
return state_dict
|
||||
else:
|
||||
raise RuntimeError("")
|
||||
|
||||
def _build_model(self):
|
||||
raise NotImplementedError()
|
||||
|
||||
def build_model(self):
|
||||
self.model = self._build_model()
|
||||
# utils.load_warp(self)
|
||||
self.unfreeze_all_params()
|
||||
if hparams['freezing_enabled']:
|
||||
self.freeze_params()
|
||||
if hparams['finetune_enabled'] and get_latest_checkpoint_path(pathlib.Path(hparams['work_dir'])) is None:
|
||||
self.load_finetune_ckpt(self.load_pre_train_model())
|
||||
self.print_arch()
|
||||
|
||||
@rank_zero_only
|
||||
def print_arch(self):
|
||||
utils.print_arch(self.model)
|
||||
|
||||
def build_losses_and_metrics(self):
|
||||
raise NotImplementedError()
|
||||
|
||||
def register_validation_metric(self, name: str, metric: Metric):
|
||||
assert isinstance(metric, Metric)
|
||||
self.valid_metrics[name] = metric
|
||||
|
||||
def register_validation_loss(self, name: str, Aggregator: Metric = MeanMetric):
|
||||
assert issubclass(Aggregator, Metric)
|
||||
self.valid_losses[name] = Aggregator()
|
||||
|
||||
def run_model(self, sample, infer=False):
|
||||
"""
|
||||
steps:
|
||||
1. run the full model
|
||||
2. calculate losses if not infer
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
def on_train_epoch_start(self):
|
||||
if self.training_sampler is not None:
|
||||
self.training_sampler.set_epoch(self.current_epoch)
|
||||
|
||||
def _training_step(self, sample):
|
||||
"""
|
||||
:return: total loss: torch.Tensor, loss_log: dict, other_log: dict
|
||||
"""
|
||||
losses = self.run_model(sample)
|
||||
total_loss = sum(losses.values())
|
||||
return total_loss, {**losses, 'batch_size': float(sample['size'])}
|
||||
|
||||
def training_step(self, sample, batch_idx):
|
||||
total_loss, log_outputs = self._training_step(sample)
|
||||
|
||||
# logs to progress bar
|
||||
self.log_dict(log_outputs, prog_bar=True, logger=False, on_step=True, on_epoch=False)
|
||||
self.log('lr', self.lr_schedulers().get_last_lr()[0], prog_bar=True, logger=False, on_step=True, on_epoch=False)
|
||||
# logs to tensorboard
|
||||
if self.global_step % hparams['log_interval'] == 0:
|
||||
tb_log = {f'training/{k}': v for k, v in log_outputs.items()}
|
||||
tb_log['training/lr'] = self.lr_schedulers().get_last_lr()[0]
|
||||
self.logger.log_metrics(tb_log, step=self.global_step)
|
||||
|
||||
return total_loss
|
||||
|
||||
# def on_before_optimizer_step(self, *args, **kwargs):
|
||||
# self.log_dict(grad_norm(self, norm_type=2))
|
||||
|
||||
def _on_validation_start(self):
|
||||
pass
|
||||
|
||||
def on_validation_start(self):
|
||||
if self.skip_immediate_validation:
|
||||
rank_zero_debug("Skip validation")
|
||||
return
|
||||
self._on_validation_start()
|
||||
for metric in self.valid_losses.values():
|
||||
metric.to(self.device)
|
||||
metric.reset()
|
||||
for metric in self.valid_metrics.values():
|
||||
metric.to(self.device)
|
||||
metric.reset()
|
||||
|
||||
def _validation_step(self, sample, batch_idx):
|
||||
"""
|
||||
|
||||
:param sample:
|
||||
:param batch_idx:
|
||||
:return: loss_log: dict, weight: int
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
def validation_step(self, sample, batch_idx):
|
||||
"""
|
||||
|
||||
:param sample:
|
||||
:param batch_idx:
|
||||
"""
|
||||
if self.skip_immediate_validation:
|
||||
rank_zero_debug("Skip validation")
|
||||
return
|
||||
if sample['size'] > 0:
|
||||
with torch.autocast(self.device.type, enabled=False):
|
||||
losses, weight = self._validation_step(sample, batch_idx)
|
||||
losses = {
|
||||
'total_loss': sum(losses.values()),
|
||||
**losses
|
||||
}
|
||||
for k, v in losses.items():
|
||||
self.valid_losses[k].update(v, weight=weight)
|
||||
|
||||
def _on_validation_epoch_end(self):
|
||||
pass
|
||||
|
||||
def on_validation_epoch_end(self):
|
||||
if self.skip_immediate_validation:
|
||||
self.skip_immediate_validation = False
|
||||
self.skip_immediate_ckpt_save = True
|
||||
return
|
||||
self._on_validation_epoch_end()
|
||||
loss_vals = {k: v.compute() for k, v in self.valid_losses.items()}
|
||||
metric_vals = {k: v.compute() for k, v in self.valid_metrics.items()}
|
||||
self.log('val_loss', loss_vals['total_loss'], on_epoch=True, prog_bar=True, logger=False, sync_dist=True)
|
||||
self.logger.log_metrics({f'validation/{k}': v for k, v in loss_vals.items()}, step=self.global_step)
|
||||
self.logger.log_metrics({f'metrics/{k}': v for k, v in metric_vals.items()}, step=self.global_step)
|
||||
|
||||
# noinspection PyMethodMayBeStatic
|
||||
def build_scheduler(self, optimizer):
|
||||
from utils import build_lr_scheduler_from_config
|
||||
|
||||
scheduler_args = hparams['lr_scheduler_args']
|
||||
assert scheduler_args['scheduler_cls'] != ''
|
||||
scheduler = build_lr_scheduler_from_config(optimizer, scheduler_args)
|
||||
return scheduler
|
||||
|
||||
# noinspection PyMethodMayBeStatic
|
||||
def build_optimizer(self, model):
|
||||
from utils import build_object_from_class_name
|
||||
|
||||
optimizer_args = hparams['optimizer_args']
|
||||
assert optimizer_args['optimizer_cls'] != ''
|
||||
if 'beta1' in optimizer_args and 'beta2' in optimizer_args and 'betas' not in optimizer_args:
|
||||
optimizer_args['betas'] = (optimizer_args['beta1'], optimizer_args['beta2'])
|
||||
optimizer = build_object_from_class_name(
|
||||
optimizer_args['optimizer_cls'],
|
||||
torch.optim.Optimizer,
|
||||
model if optimizer_args['optimizer_cls'] == 'modules.optimizer.muon.Muon_AdamW' else model.parameters(),
|
||||
**optimizer_args
|
||||
)
|
||||
return optimizer
|
||||
|
||||
def configure_optimizers(self):
|
||||
optm = self.build_optimizer(self.model)
|
||||
scheduler = self.build_scheduler(optm)
|
||||
if scheduler is None:
|
||||
return optm
|
||||
return {
|
||||
"optimizer": optm,
|
||||
"lr_scheduler": {
|
||||
"scheduler": scheduler,
|
||||
"interval": "step",
|
||||
"frequency": 1
|
||||
}
|
||||
}
|
||||
|
||||
def train_dataloader(self):
|
||||
self.training_sampler = DsBatchSampler(
|
||||
self.train_dataset,
|
||||
max_batch_frames=self.max_batch_frames,
|
||||
max_batch_size=self.max_batch_size,
|
||||
num_replicas=self.num_replicas,
|
||||
rank=self.global_rank,
|
||||
sort_by_similar_size=hparams['sort_by_len'],
|
||||
size_reversed=True,
|
||||
required_batch_count_multiple=hparams['accumulate_grad_batches'],
|
||||
shuffle_sample=True,
|
||||
shuffle_batch=True
|
||||
)
|
||||
return torch.utils.data.DataLoader(
|
||||
self.train_dataset,
|
||||
collate_fn=self.train_dataset.collater,
|
||||
batch_sampler=self.training_sampler,
|
||||
num_workers=hparams['ds_workers'],
|
||||
prefetch_factor=hparams['dataloader_prefetch_factor'],
|
||||
pin_memory=True,
|
||||
persistent_workers=True
|
||||
)
|
||||
|
||||
def val_dataloader(self):
|
||||
sampler = DsBatchSampler(
|
||||
self.valid_dataset,
|
||||
max_batch_frames=self.max_val_batch_frames,
|
||||
max_batch_size=self.max_val_batch_size,
|
||||
num_replicas=self.num_replicas,
|
||||
rank=self.global_rank,
|
||||
shuffle_sample=False,
|
||||
shuffle_batch=False,
|
||||
disallow_empty_batch=False,
|
||||
pad_batch_assignment=False
|
||||
)
|
||||
return torch.utils.data.DataLoader(
|
||||
self.valid_dataset,
|
||||
collate_fn=self.valid_dataset.collater,
|
||||
batch_sampler=sampler,
|
||||
num_workers=hparams['ds_workers'],
|
||||
prefetch_factor=hparams['dataloader_prefetch_factor'],
|
||||
persistent_workers=True
|
||||
)
|
||||
|
||||
def test_dataloader(self):
|
||||
return self.val_dataloader()
|
||||
|
||||
def on_test_start(self):
|
||||
self.on_validation_start()
|
||||
|
||||
def test_step(self, sample, batch_idx):
|
||||
return self.validation_step(sample, batch_idx)
|
||||
|
||||
def on_test_end(self):
|
||||
return self.on_validation_end()
|
||||
|
||||
###########
|
||||
# Running configuration
|
||||
###########
|
||||
|
||||
@classmethod
|
||||
def start(cls):
|
||||
task = cls()
|
||||
|
||||
# if pre_train is not None:
|
||||
# task.load_state_dict(pre_train,strict=False)
|
||||
# print("load success-------------------------------------------------------------------")
|
||||
|
||||
work_dir = pathlib.Path(hparams['work_dir'])
|
||||
trainer = pl.Trainer(
|
||||
accelerator=hparams['pl_trainer_accelerator'],
|
||||
devices=hparams['pl_trainer_devices'],
|
||||
num_nodes=hparams['pl_trainer_num_nodes'],
|
||||
strategy=get_strategy(
|
||||
hparams['pl_trainer_devices'],
|
||||
hparams['pl_trainer_num_nodes'],
|
||||
hparams['pl_trainer_accelerator'],
|
||||
hparams['pl_trainer_strategy'],
|
||||
hparams['pl_trainer_precision'],
|
||||
),
|
||||
precision=hparams['pl_trainer_precision'],
|
||||
callbacks=[
|
||||
DsModelCheckpoint(
|
||||
dirpath=work_dir,
|
||||
filename='model_ckpt_steps_{step}',
|
||||
auto_insert_metric_name=False,
|
||||
monitor='step',
|
||||
mode='max',
|
||||
save_last=False,
|
||||
# every_n_train_steps=hparams['val_check_interval'],
|
||||
save_top_k=hparams['num_ckpt_keep'],
|
||||
permanent_ckpt_start=hparams['permanent_ckpt_start'],
|
||||
permanent_ckpt_interval=hparams['permanent_ckpt_interval'],
|
||||
verbose=True
|
||||
),
|
||||
# LearningRateMonitor(logging_interval='step'),
|
||||
DsTQDMProgressBar(),
|
||||
],
|
||||
logger=DsTensorBoardLogger(
|
||||
save_dir=str(work_dir),
|
||||
name='lightning_logs',
|
||||
version='latest'
|
||||
),
|
||||
gradient_clip_val=hparams['clip_grad_norm'],
|
||||
val_check_interval=hparams['val_check_interval'] * hparams['accumulate_grad_batches'],
|
||||
# so this is global_steps
|
||||
check_val_every_n_epoch=None,
|
||||
log_every_n_steps=1,
|
||||
max_steps=hparams['max_updates'],
|
||||
use_distributed_sampler=False,
|
||||
num_sanity_val_steps=hparams['num_sanity_val_steps'],
|
||||
accumulate_grad_batches=hparams['accumulate_grad_batches']
|
||||
)
|
||||
if not hparams['infer']: # train
|
||||
@rank_zero_only
|
||||
def train_payload_copy():
|
||||
# Copy files to work_dir
|
||||
binary_dir = pathlib.Path(hparams['binary_data_dir'])
|
||||
spk_map_dst = work_dir / 'spk_map.json'
|
||||
spk_map_src = binary_dir / 'spk_map.json'
|
||||
shutil.copy(spk_map_src, spk_map_dst)
|
||||
print(f'| Copied spk map to {spk_map_dst}.')
|
||||
lang_map_dst = work_dir / 'lang_map.json'
|
||||
lang_map_src = binary_dir / 'lang_map.json'
|
||||
shutil.copy(lang_map_src, lang_map_dst)
|
||||
print(f'| Copied lang map to {lang_map_dst}.')
|
||||
for lang in hparams['dictionaries'].keys():
|
||||
dict_dst = work_dir / f'dictionary-{lang}.txt'
|
||||
dict_src = binary_dir / f'dictionary-{lang}.txt'
|
||||
shutil.copy(dict_src, dict_dst)
|
||||
print(f'| Copied dictionary for language \'{lang}\' to {dict_dst}.')
|
||||
|
||||
train_payload_copy()
|
||||
trainer.fit(task, ckpt_path=get_latest_checkpoint_path(work_dir))
|
||||
else:
|
||||
trainer.test(task)
|
||||
|
||||
def on_save_checkpoint(self, checkpoint):
|
||||
if isinstance(self.model, CategorizedModule):
|
||||
checkpoint['category'] = self.model.category
|
||||
checkpoint['trainer_stage'] = self.trainer.state.stage.value
|
||||
|
||||
def on_load_checkpoint(self, checkpoint):
|
||||
from lightning.pytorch.trainer.states import RunningStage
|
||||
from utils import simulate_lr_scheduler
|
||||
if checkpoint.get('trainer_stage', '') == RunningStage.VALIDATING.value:
|
||||
self.skip_immediate_validation = True
|
||||
|
||||
optimizer_args = hparams['optimizer_args']
|
||||
scheduler_args = hparams['lr_scheduler_args']
|
||||
|
||||
if 'beta1' in optimizer_args and 'beta2' in optimizer_args and 'betas' not in optimizer_args:
|
||||
optimizer_args['betas'] = (optimizer_args['beta1'], optimizer_args['beta2'])
|
||||
|
||||
if checkpoint.get('optimizer_states', None):
|
||||
opt_states = checkpoint['optimizer_states']
|
||||
assert len(opt_states) == 1 # only support one optimizer
|
||||
opt_state = opt_states[0]
|
||||
for param_group in opt_state['param_groups']:
|
||||
for k, v in optimizer_args.items():
|
||||
if k in param_group and param_group[k] != v:
|
||||
if 'lr_schedulers' in checkpoint and checkpoint['lr_schedulers'] and k == 'lr':
|
||||
continue
|
||||
rank_zero_info(f'| Overriding optimizer parameter {k} from checkpoint: {param_group[k]} -> {v}')
|
||||
param_group[k] = v
|
||||
if 'initial_lr' in param_group and param_group['initial_lr'] != optimizer_args['lr']:
|
||||
rank_zero_info(
|
||||
f'| Overriding optimizer parameter initial_lr from checkpoint: {param_group["initial_lr"]} -> {optimizer_args["lr"]}'
|
||||
)
|
||||
param_group['initial_lr'] = optimizer_args['lr']
|
||||
|
||||
if checkpoint.get('lr_schedulers', None):
|
||||
assert checkpoint.get('optimizer_states', False)
|
||||
assert len(checkpoint['lr_schedulers']) == 1 # only support one scheduler
|
||||
checkpoint['lr_schedulers'][0] = simulate_lr_scheduler(
|
||||
optimizer_args, scheduler_args,
|
||||
step_count=checkpoint['global_step'],
|
||||
num_param_groups=len(checkpoint['optimizer_states'][0]['param_groups'])
|
||||
)
|
||||
for param_group, new_lr in zip(
|
||||
checkpoint['optimizer_states'][0]['param_groups'],
|
||||
checkpoint['lr_schedulers'][0]['_last_lr'],
|
||||
):
|
||||
if param_group['lr'] != new_lr:
|
||||
rank_zero_info(f'| Overriding optimizer parameter lr from checkpoint: {param_group["lr"]} -> {new_lr}')
|
||||
param_group['lr'] = new_lr
|
||||
@@ -0,0 +1,23 @@
|
||||
class BaseVocoder:
|
||||
def to_device(self, device):
|
||||
"""
|
||||
|
||||
:param device: torch.device or str
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
def get_device(self):
|
||||
"""
|
||||
|
||||
:return: device: torch.device or str
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
def spec2wav(self, mel, **kwargs):
|
||||
"""
|
||||
|
||||
:param mel: [T, 80]
|
||||
:return: wav: [T']
|
||||
"""
|
||||
|
||||
raise NotImplementedError()
|
||||
@@ -0,0 +1,144 @@
|
||||
base_config:
|
||||
- configs/base.yaml
|
||||
|
||||
task_cls: training.acoustic_task.AcousticTask
|
||||
|
||||
dictionaries: {}
|
||||
extra_phonemes: []
|
||||
merged_phoneme_groups: []
|
||||
datasets: []
|
||||
|
||||
vocoder: NsfHifiGAN
|
||||
vocoder_ckpt: checkpoints/pc_nsf_hifigan_44.1k_hop512_128bin_2025.02/model.ckpt
|
||||
audio_sample_rate: 44100
|
||||
audio_num_mel_bins: 128
|
||||
hop_size: 512 # Hop size.
|
||||
fft_size: 2048 # FFT size.
|
||||
win_size: 2048 # FFT size.
|
||||
fmin: 40
|
||||
fmax: 16000
|
||||
|
||||
binarization_args:
|
||||
shuffle: true
|
||||
num_workers: 0
|
||||
augmentation_args:
|
||||
random_pitch_shifting:
|
||||
enabled: false
|
||||
range: [-5., 5.]
|
||||
scale: 0.75
|
||||
fixed_pitch_shifting:
|
||||
enabled: false
|
||||
targets: [-5., 5.]
|
||||
scale: 0.5
|
||||
random_time_stretching:
|
||||
enabled: false
|
||||
range: [0.5, 2.]
|
||||
scale: 0.75
|
||||
|
||||
binary_data_dir: 'data/opencpop/binary'
|
||||
binarizer_cls: preprocessing.acoustic_binarizer.AcousticBinarizer
|
||||
spec_min: [-12]
|
||||
spec_max: [0]
|
||||
mel_vmin: -14.
|
||||
mel_vmax: 4.
|
||||
mel_base: 'e'
|
||||
energy_smooth_width: 0.06
|
||||
breathiness_smooth_width: 0.06
|
||||
voicing_smooth_width: 0.06
|
||||
tension_smooth_width: 0.06
|
||||
|
||||
use_lang_id: false
|
||||
num_lang: 1
|
||||
use_spk_id: false
|
||||
num_spk: 1
|
||||
use_mix_ln: false
|
||||
mix_ln_layer: [0, 2]
|
||||
use_energy_embed: false
|
||||
use_breathiness_embed: false
|
||||
use_voicing_embed: false
|
||||
use_tension_embed: false
|
||||
use_key_shift_embed: false
|
||||
use_speed_embed: false
|
||||
|
||||
diffusion_type: reflow
|
||||
time_scale_factor: 1000
|
||||
timesteps: 1000
|
||||
max_beta: 0.02
|
||||
enc_ffn_kernel_size: 3
|
||||
use_rope: true
|
||||
rope_interleaved: false
|
||||
use_stretch_embed: true
|
||||
use_variance_scaling: true
|
||||
rel_pos: true
|
||||
sampling_algorithm: euler
|
||||
sampling_steps: 20
|
||||
diff_accelerator: ddim
|
||||
diff_speedup: 10
|
||||
hidden_size: 384
|
||||
backbone_type: 'lynxnet2'
|
||||
backbone_args:
|
||||
num_channels: 1024
|
||||
num_layers: 6
|
||||
kernel_size: 31
|
||||
dropout_rate: 0.0
|
||||
use_conditioner_cache: true
|
||||
glu_type: 'atanglu'
|
||||
main_loss_type: l2
|
||||
main_loss_log_norm: false
|
||||
schedule_type: 'linear'
|
||||
|
||||
# shallow diffusion
|
||||
use_shallow_diffusion: true
|
||||
T_start: 0.4
|
||||
T_start_infer: 0.4
|
||||
K_step: 400
|
||||
K_step_infer: 400
|
||||
|
||||
shallow_diffusion_args:
|
||||
train_aux_decoder: true
|
||||
train_diffusion: true
|
||||
val_gt_start: false
|
||||
aux_decoder_arch: convnext
|
||||
aux_decoder_args:
|
||||
num_channels: 512
|
||||
num_layers: 6
|
||||
kernel_size: 7
|
||||
dropout_rate: 0.1
|
||||
aux_decoder_grad: 0.1
|
||||
|
||||
lambda_aux_mel_loss: 0.2
|
||||
|
||||
# train and eval
|
||||
num_sanity_val_steps: 1
|
||||
optimizer_args:
|
||||
optimizer_cls: modules.optimizer.muon.Muon_AdamW
|
||||
lr: 0.0006
|
||||
muon_args:
|
||||
weight_decay: 0.1
|
||||
adamw_args:
|
||||
weight_decay: 0.0
|
||||
lr_scheduler_args:
|
||||
step_size: 5000
|
||||
gamma: 0.8
|
||||
max_batch_frames: 50000
|
||||
max_batch_size: 64
|
||||
dataset_size_key: 'lengths'
|
||||
val_with_vocoder: true
|
||||
val_check_interval: 4000
|
||||
num_valid_plots: 10
|
||||
max_updates: 100000
|
||||
num_ckpt_keep: 8
|
||||
permanent_ckpt_start: 60000
|
||||
permanent_ckpt_interval: 10000
|
||||
|
||||
finetune_enabled: false
|
||||
finetune_ckpt_path: null
|
||||
|
||||
finetune_ignored_params:
|
||||
- model.fs2.encoder.embed_tokens
|
||||
- model.fs2.txt_embed
|
||||
- model.fs2.spk_embed
|
||||
finetune_strict_shapes: true
|
||||
|
||||
freezing_enabled: false
|
||||
frozen_params: []
|
||||
@@ -0,0 +1,94 @@
|
||||
# task
|
||||
task_cls: null
|
||||
|
||||
#############
|
||||
# dataset
|
||||
#############
|
||||
sort_by_len: true
|
||||
datasets: []
|
||||
binary_data_dir: null
|
||||
binarizer_cls: null
|
||||
binarization_args:
|
||||
shuffle: false
|
||||
num_workers: 0
|
||||
|
||||
audio_sample_rate: 44100
|
||||
hop_size: 512
|
||||
win_size: 2048
|
||||
fft_size: 2048 # Extra window size is filled with 0 paddings to match this parameter
|
||||
sampler_frame_count_grid: 6
|
||||
ds_workers: 4
|
||||
dataloader_prefetch_factor: 2
|
||||
|
||||
#########
|
||||
# model
|
||||
#########
|
||||
hidden_size: 384
|
||||
dropout: 0.1
|
||||
use_pos_embed: true
|
||||
enc_layers: 4
|
||||
num_heads: 2
|
||||
enc_ffn_kernel_size: 3
|
||||
ffn_act: gelu
|
||||
use_spk_id: false
|
||||
|
||||
###########
|
||||
# optimization
|
||||
###########
|
||||
optimizer_args:
|
||||
optimizer_cls: modules.optimizer.muon.Muon_AdamW
|
||||
lr: 0.0006
|
||||
beta1: 0.9
|
||||
beta2: 0.98
|
||||
weight_decay: 0
|
||||
lr_scheduler_args:
|
||||
scheduler_cls: torch.optim.lr_scheduler.StepLR
|
||||
step_size: 5000
|
||||
gamma: 0.8
|
||||
clip_grad_norm: 1
|
||||
|
||||
###########
|
||||
# train and eval
|
||||
###########
|
||||
num_ckpt_keep: 5
|
||||
accumulate_grad_batches: 1
|
||||
log_interval: 100
|
||||
num_sanity_val_steps: 1 # steps of validation at the beginning
|
||||
val_check_interval: 4000
|
||||
max_updates: 100000
|
||||
max_batch_frames: 50000
|
||||
max_batch_size: 100000
|
||||
max_val_batch_frames: 60000
|
||||
max_val_batch_size: 1
|
||||
pe: parselmouth
|
||||
pe_ckpt: 'checkpoints/rmvpe/model.pt'
|
||||
hnsep: vr
|
||||
hnsep_ckpt: 'checkpoints/vr/model.pt'
|
||||
f0_min: 65
|
||||
f0_max: 1100
|
||||
num_valid_plots: 10
|
||||
|
||||
###########
|
||||
# pytorch lightning
|
||||
# Read https://lightning.ai/docs/pytorch/stable/common/trainer.html#trainer-class-api for possible values
|
||||
###########
|
||||
pl_trainer_accelerator: 'auto'
|
||||
pl_trainer_devices: 'auto'
|
||||
pl_trainer_precision: '16-mixed'
|
||||
pl_trainer_num_nodes: 1
|
||||
pl_trainer_strategy:
|
||||
name: auto
|
||||
process_group_backend: nccl
|
||||
find_unused_parameters: false
|
||||
nccl_p2p: true
|
||||
|
||||
###########
|
||||
# finetune
|
||||
###########
|
||||
finetune_enabled: false
|
||||
finetune_ckpt_path: null
|
||||
finetune_ignored_params: []
|
||||
finetune_strict_shapes: true
|
||||
|
||||
freezing_enabled: false
|
||||
frozen_params: []
|
||||
@@ -0,0 +1,128 @@
|
||||
base_config:
|
||||
- configs/acoustic.yaml
|
||||
|
||||
dictionaries:
|
||||
zh: dictionaries/opencpop-extension.txt
|
||||
extra_phonemes: []
|
||||
merged_phoneme_groups: []
|
||||
|
||||
datasets:
|
||||
- raw_data_dir: data/xxx1/raw
|
||||
speaker: speaker1
|
||||
spk_id: 0
|
||||
language: zh
|
||||
test_prefixes:
|
||||
- wav1
|
||||
- wav2
|
||||
- wav3
|
||||
- wav4
|
||||
- wav5
|
||||
- raw_data_dir: data/xxx2/raw
|
||||
speaker: speaker2
|
||||
spk_id: 1
|
||||
language: zh
|
||||
test_prefixes:
|
||||
- wav1
|
||||
- wav2
|
||||
- wav3
|
||||
- wav4
|
||||
- wav5
|
||||
|
||||
binary_data_dir: data/xxx/binary
|
||||
binarization_args:
|
||||
num_workers: 0
|
||||
pe: parselmouth
|
||||
pe_ckpt: 'checkpoints/rmvpe/model.pt'
|
||||
hnsep: vr
|
||||
hnsep_ckpt: 'checkpoints/vr/model.pt'
|
||||
vocoder: NsfHifiGAN
|
||||
vocoder_ckpt: checkpoints/nsf_hifigan_44.1k_hop512_128bin_2024.02/model.ckpt
|
||||
|
||||
use_lang_id: false
|
||||
num_lang: 1
|
||||
use_spk_id: false
|
||||
num_spk: 1
|
||||
|
||||
use_mix_ln: false
|
||||
mix_ln_layer: [0, 2]
|
||||
|
||||
# NOTICE: before enabling variance embeddings, please read the docs at
|
||||
# https://github.com/openvpi/DiffSinger/tree/main/docs/BestPractices.md#choosing-variance-parameters
|
||||
use_energy_embed: false
|
||||
use_breathiness_embed: false
|
||||
use_voicing_embed: false
|
||||
use_tension_embed: false
|
||||
|
||||
use_key_shift_embed: true
|
||||
use_speed_embed: true
|
||||
|
||||
augmentation_args:
|
||||
random_pitch_shifting:
|
||||
enabled: true
|
||||
range: [-5., 5.]
|
||||
scale: 0.75
|
||||
fixed_pitch_shifting:
|
||||
enabled: false
|
||||
targets: [-5., 5.]
|
||||
scale: 0.5
|
||||
random_time_stretching:
|
||||
enabled: true
|
||||
range: [0.5, 2.]
|
||||
scale: 0.75
|
||||
|
||||
# diffusion and shallow diffusion
|
||||
diffusion_type: reflow
|
||||
enc_ffn_kernel_size: 3
|
||||
use_rope: true
|
||||
rope_interleaved: false
|
||||
use_stretch_embed: true
|
||||
use_variance_scaling: true
|
||||
use_shallow_diffusion: true
|
||||
T_start: 0.4
|
||||
T_start_infer: 0.4
|
||||
K_step: 300
|
||||
K_step_infer: 300
|
||||
hidden_size: 384
|
||||
backbone_type: 'lynxnet2'
|
||||
backbone_args:
|
||||
num_channels: 1024
|
||||
num_layers: 6
|
||||
kernel_size: 31
|
||||
dropout_rate: 0.0
|
||||
use_conditioner_cache: true
|
||||
glu_type: 'atanglu'
|
||||
shallow_diffusion_args:
|
||||
train_aux_decoder: true
|
||||
train_diffusion: true
|
||||
val_gt_start: false
|
||||
aux_decoder_arch: convnext
|
||||
aux_decoder_args:
|
||||
num_channels: 512
|
||||
num_layers: 6
|
||||
kernel_size: 7
|
||||
dropout_rate: 0.1
|
||||
aux_decoder_grad: 0.1
|
||||
lambda_aux_mel_loss: 0.2
|
||||
|
||||
optimizer_args:
|
||||
optimizer_cls: modules.optimizer.muon.Muon_AdamW
|
||||
lr: 0.0006
|
||||
muon_args:
|
||||
weight_decay: 0.1
|
||||
adamw_args:
|
||||
weight_decay: 0.0
|
||||
lr_scheduler_args:
|
||||
step_size: 5000
|
||||
gamma: 0.8
|
||||
max_batch_frames: 50000
|
||||
max_batch_size: 64
|
||||
max_updates: 100000
|
||||
|
||||
num_valid_plots: 10
|
||||
val_with_vocoder: true
|
||||
val_check_interval: 4000
|
||||
num_ckpt_keep: 8
|
||||
permanent_ckpt_start: 60000
|
||||
permanent_ckpt_interval: 10000
|
||||
pl_trainer_devices: 'auto'
|
||||
pl_trainer_precision: '16-mixed'
|
||||
@@ -0,0 +1,139 @@
|
||||
base_config:
|
||||
- configs/variance.yaml
|
||||
|
||||
dictionaries:
|
||||
zh: dictionaries/opencpop-extension.txt
|
||||
extra_phonemes: []
|
||||
merged_phoneme_groups: []
|
||||
|
||||
datasets:
|
||||
- raw_data_dir: data/xxx1/raw
|
||||
speaker: speaker1
|
||||
spk_id: 0
|
||||
language: zh
|
||||
test_prefixes:
|
||||
- wav1
|
||||
- wav2
|
||||
- wav3
|
||||
- wav4
|
||||
- wav5
|
||||
- raw_data_dir: data/xxx2/raw
|
||||
speaker: speaker2
|
||||
spk_id: 1
|
||||
language: zh
|
||||
test_prefixes:
|
||||
- wav1
|
||||
- wav2
|
||||
- wav3
|
||||
- wav4
|
||||
- wav5
|
||||
|
||||
binary_data_dir: data/xxx/binary
|
||||
binarization_args:
|
||||
num_workers: 0
|
||||
pe: parselmouth
|
||||
pe_ckpt: 'checkpoints/rmvpe/model.pt'
|
||||
hnsep: vr
|
||||
hnsep_ckpt: 'checkpoints/vr/model.pt'
|
||||
|
||||
use_lang_id: false
|
||||
num_lang: 1
|
||||
use_spk_id: false
|
||||
num_spk: 1
|
||||
# NOTICE: before enabling variance modules, please read the docs at
|
||||
# https://github.com/openvpi/DiffSinger/tree/main/docs/BestPractices.md#mutual-influence-between-variance-modules
|
||||
predict_dur: true
|
||||
predict_pitch: false
|
||||
# NOTICE: before enabling variance predictions, please read the docs at
|
||||
# https://github.com/openvpi/DiffSinger/tree/main/docs/BestPractices.md#choosing-variance-parameters
|
||||
predict_energy: false
|
||||
predict_breathiness: true
|
||||
predict_voicing: true
|
||||
predict_tension: false
|
||||
|
||||
energy_db_min: -96.0
|
||||
energy_db_max: -12.0
|
||||
|
||||
breathiness_db_min: -96.0
|
||||
breathiness_db_max: -20.0
|
||||
|
||||
voicing_db_min: -96.0
|
||||
voicing_db_max: -12.0
|
||||
|
||||
tension_logit_min: -10.0
|
||||
tension_logit_max: 10.0
|
||||
|
||||
enc_ffn_kernel_size: 3
|
||||
use_rope: true
|
||||
rope_interleaved: false
|
||||
use_stretch_embed: false
|
||||
use_variance_scaling: true
|
||||
hidden_size: 384
|
||||
dur_prediction_args:
|
||||
arch: resnet
|
||||
hidden_size: 256
|
||||
dropout: 0.1
|
||||
num_layers: 5
|
||||
kernel_size: 3
|
||||
log_offset: 1.0
|
||||
loss_type: mse
|
||||
lambda_pdur_loss: 0.3
|
||||
lambda_wdur_loss: 1.0
|
||||
lambda_sdur_loss: 3.0
|
||||
|
||||
use_melody_encoder: true
|
||||
melody_encoder_args:
|
||||
hidden_size: 128
|
||||
enc_layers: 4
|
||||
use_glide_embed: false
|
||||
glide_types: [up, down]
|
||||
glide_embed_scale: 11.313708498984760 # sqrt(128)
|
||||
|
||||
diffusion_type: reflow
|
||||
|
||||
pitch_prediction_args:
|
||||
pitd_norm_min: -8.0
|
||||
pitd_norm_max: 8.0
|
||||
pitd_clip_min: -12.0
|
||||
pitd_clip_max: 12.0
|
||||
repeat_bins: 64
|
||||
backbone_type: 'lynxnet2'
|
||||
backbone_args:
|
||||
num_layers: 6
|
||||
num_channels: 512
|
||||
dropout_rate: 0.0
|
||||
use_conditioner_cache: true
|
||||
glu_type: 'atanglu'
|
||||
|
||||
variances_prediction_args:
|
||||
total_repeat_bins: 72
|
||||
backbone_type: 'lynxnet2'
|
||||
backbone_args:
|
||||
num_layers: 6
|
||||
num_channels: 384
|
||||
dropout_rate: 0.0
|
||||
use_conditioner_cache: true
|
||||
glu_type: 'atanglu'
|
||||
|
||||
lambda_dur_loss: 1.0
|
||||
lambda_pitch_loss: 1.0
|
||||
lambda_var_loss: 1.0
|
||||
|
||||
optimizer_args:
|
||||
optimizer_cls: torch.optim.AdamW
|
||||
lr: 0.0006
|
||||
lr_scheduler_args:
|
||||
scheduler_cls: torch.optim.lr_scheduler.StepLR
|
||||
step_size: 5000
|
||||
gamma: 0.75
|
||||
max_batch_frames: 80000
|
||||
max_batch_size: 48
|
||||
max_updates: 60000
|
||||
|
||||
num_valid_plots: 10
|
||||
val_check_interval: 4000
|
||||
num_ckpt_keep: 8
|
||||
permanent_ckpt_start: 30000
|
||||
permanent_ckpt_interval: 10000
|
||||
pl_trainer_devices: 'auto'
|
||||
pl_trainer_precision: '16-mixed'
|
||||
@@ -0,0 +1,142 @@
|
||||
base_config:
|
||||
- configs/variance.yaml
|
||||
|
||||
dictionaries:
|
||||
zh: dictionaries/opencpop-extension.txt
|
||||
extra_phonemes: []
|
||||
merged_phoneme_groups: []
|
||||
|
||||
datasets:
|
||||
- raw_data_dir: data/xxx1/raw
|
||||
speaker: speaker1
|
||||
spk_id: 0
|
||||
language: zh
|
||||
test_prefixes:
|
||||
- wav1
|
||||
- wav2
|
||||
- wav3
|
||||
- wav4
|
||||
- wav5
|
||||
- raw_data_dir: data/xxx2/raw
|
||||
speaker: speaker2
|
||||
spk_id: 1
|
||||
language: zh
|
||||
test_prefixes:
|
||||
- wav1
|
||||
- wav2
|
||||
- wav3
|
||||
- wav4
|
||||
- wav5
|
||||
|
||||
binary_data_dir: data/xxx/binary
|
||||
binarization_args:
|
||||
num_workers: 0
|
||||
pe: parselmouth
|
||||
pe_ckpt: 'checkpoints/rmvpe/model.pt'
|
||||
hnsep: vr
|
||||
hnsep_ckpt: 'checkpoints/vr/model.pt'
|
||||
|
||||
use_lang_id: false
|
||||
num_lang: 1
|
||||
use_spk_id: false
|
||||
num_spk: 1
|
||||
# NOTICE: before enabling variance modules, please read the docs at
|
||||
# https://github.com/openvpi/DiffSinger/tree/main/docs/BestPractices.md#mutual-influence-between-variance-modules
|
||||
predict_dur: false
|
||||
predict_pitch: false
|
||||
# NOTICE: before enabling variance predictions, please read the docs at
|
||||
# https://github.com/openvpi/DiffSinger/tree/main/docs/BestPractices.md#choosing-variance-parameters
|
||||
predict_energy: false
|
||||
predict_breathiness: false
|
||||
predict_voicing: false
|
||||
predict_tension: false
|
||||
|
||||
energy_db_min: -96.0
|
||||
energy_db_max: -12.0
|
||||
|
||||
breathiness_db_min: -96.0
|
||||
breathiness_db_max: -20.0
|
||||
|
||||
voicing_db_min: -96.0
|
||||
voicing_db_max: -12.0
|
||||
|
||||
tension_logit_min: -10.0
|
||||
tension_logit_max: 10.0
|
||||
|
||||
enc_ffn_kernel_size: 3
|
||||
use_rope: true
|
||||
rope_interleaved: false
|
||||
use_stretch_embed: false
|
||||
use_variance_scaling: true
|
||||
hidden_size: 384
|
||||
dur_prediction_args:
|
||||
arch: resnet
|
||||
hidden_size: 256
|
||||
dropout: 0.1
|
||||
num_layers: 5
|
||||
kernel_size: 3
|
||||
log_offset: 1.0
|
||||
loss_type: mse
|
||||
lambda_pdur_loss: 0.3
|
||||
lambda_wdur_loss: 1.0
|
||||
lambda_sdur_loss: 3.0
|
||||
|
||||
use_melody_encoder: true
|
||||
melody_encoder_args:
|
||||
hidden_size: 128
|
||||
enc_layers: 4
|
||||
use_glide_embed: false
|
||||
glide_types: [up, down]
|
||||
glide_embed_scale: 11.313708498984760 # sqrt(128)
|
||||
|
||||
diffusion_type: reflow
|
||||
|
||||
pitch_prediction_args:
|
||||
pitd_norm_min: -8.0
|
||||
pitd_norm_max: 8.0
|
||||
pitd_clip_min: -12.0
|
||||
pitd_clip_max: 12.0
|
||||
repeat_bins: 64
|
||||
backbone_type: 'lynxnet2'
|
||||
backbone_args:
|
||||
num_layers: 6
|
||||
num_channels: 512
|
||||
dropout_rate: 0.0
|
||||
use_conditioner_cache: true
|
||||
glu_type: 'atanglu'
|
||||
|
||||
variances_prediction_args:
|
||||
total_repeat_bins: 72
|
||||
backbone_type: 'lynxnet2'
|
||||
backbone_args:
|
||||
num_layers: 6
|
||||
num_channels: 384
|
||||
dropout_rate: 0.0
|
||||
use_conditioner_cache: true
|
||||
glu_type: 'atanglu'
|
||||
|
||||
lambda_dur_loss: 1.0
|
||||
lambda_pitch_loss: 1.0
|
||||
lambda_var_loss: 1.0
|
||||
|
||||
optimizer_args:
|
||||
optimizer_cls: modules.optimizer.muon.Muon_AdamW
|
||||
lr: 0.0006
|
||||
muon_args:
|
||||
weight_decay: 0.1
|
||||
adamw_args:
|
||||
weight_decay: 0.0
|
||||
lr_scheduler_args:
|
||||
step_size: 5000
|
||||
gamma: 0.8
|
||||
max_batch_frames: 80000
|
||||
max_batch_size: 48
|
||||
max_updates: 80000
|
||||
|
||||
num_valid_plots: 10
|
||||
val_check_interval: 2000
|
||||
num_ckpt_keep: 8
|
||||
permanent_ckpt_start: 30000
|
||||
permanent_ckpt_interval: 10000
|
||||
pl_trainer_devices: 'auto'
|
||||
pl_trainer_precision: '16-mixed'
|
||||
@@ -0,0 +1,152 @@
|
||||
base_config:
|
||||
- configs/base.yaml
|
||||
|
||||
task_cls: training.variance_task.VarianceTask
|
||||
|
||||
dictionaries: {}
|
||||
extra_phonemes: []
|
||||
merged_phoneme_groups: []
|
||||
datasets: []
|
||||
|
||||
audio_sample_rate: 44100
|
||||
hop_size: 512 # Hop size.
|
||||
fft_size: 2048 # FFT size.
|
||||
win_size: 2048 # FFT size.
|
||||
midi_smooth_width: 0.06 # in seconds
|
||||
|
||||
binarization_args:
|
||||
shuffle: true
|
||||
num_workers: 0
|
||||
prefer_ds: false
|
||||
|
||||
binary_data_dir: 'data/opencpop_variance/binary'
|
||||
binarizer_cls: preprocessing.variance_binarizer.VarianceBinarizer
|
||||
|
||||
use_lang_id: false
|
||||
num_lang: 1
|
||||
use_spk_id: false
|
||||
num_spk: 1
|
||||
|
||||
predict_dur: true
|
||||
predict_pitch: true
|
||||
predict_energy: false
|
||||
predict_breathiness: false
|
||||
predict_voicing: false
|
||||
predict_tension: false
|
||||
|
||||
enc_ffn_kernel_size: 3
|
||||
use_rope: true
|
||||
rope_interleaved: false
|
||||
use_stretch_embed: false
|
||||
use_variance_scaling: true
|
||||
rel_pos: true
|
||||
hidden_size: 384
|
||||
|
||||
dur_prediction_args:
|
||||
arch: resnet
|
||||
hidden_size: 256
|
||||
dropout: 0.1
|
||||
num_layers: 5
|
||||
kernel_size: 3
|
||||
log_offset: 1.0
|
||||
loss_type: mse
|
||||
lambda_pdur_loss: 0.3
|
||||
lambda_wdur_loss: 1.0
|
||||
lambda_sdur_loss: 3.0
|
||||
|
||||
use_melody_encoder: true
|
||||
melody_encoder_args:
|
||||
hidden_size: 128
|
||||
enc_layers: 4
|
||||
use_glide_embed: false
|
||||
glide_types: [up, down]
|
||||
glide_embed_scale: 11.313708498984760 # sqrt(128)
|
||||
|
||||
pitch_prediction_args:
|
||||
pitd_norm_min: -8.0
|
||||
pitd_norm_max: 8.0
|
||||
pitd_clip_min: -12.0
|
||||
pitd_clip_max: 12.0
|
||||
repeat_bins: 64
|
||||
backbone_type: 'lynxnet2'
|
||||
backbone_args:
|
||||
num_layers: 6
|
||||
num_channels: 512
|
||||
dropout_rate: 0.0
|
||||
use_conditioner_cache: true
|
||||
glu_type: 'atanglu'
|
||||
|
||||
energy_db_min: -96.0
|
||||
energy_db_max: -12.0
|
||||
energy_smooth_width: 0.06
|
||||
|
||||
breathiness_db_min: -96.0
|
||||
breathiness_db_max: -20.0
|
||||
breathiness_smooth_width: 0.06
|
||||
voicing_db_min: -96.0
|
||||
voicing_db_max: -12.0
|
||||
voicing_smooth_width: 0.06
|
||||
|
||||
tension_logit_min: -10.0
|
||||
tension_logit_max: 10.0
|
||||
tension_smooth_width: 0.06
|
||||
|
||||
variances_prediction_args:
|
||||
total_repeat_bins: 72
|
||||
backbone_type: 'lynxnet2'
|
||||
backbone_args:
|
||||
num_layers: 6
|
||||
num_channels: 384
|
||||
dropout_rate: 0.0
|
||||
use_conditioner_cache: true
|
||||
glu_type: 'atanglu'
|
||||
|
||||
lambda_dur_loss: 1.0
|
||||
lambda_pitch_loss: 1.0
|
||||
lambda_var_loss: 1.0
|
||||
|
||||
diffusion_type: reflow # ddpm
|
||||
time_scale_factor: 1000
|
||||
schedule_type: 'linear'
|
||||
K_step: 1000
|
||||
timesteps: 1000
|
||||
max_beta: 0.02
|
||||
main_loss_type: l2
|
||||
main_loss_log_norm: true
|
||||
sampling_algorithm: euler
|
||||
sampling_steps: 20
|
||||
diff_accelerator: ddim
|
||||
diff_speedup: 10
|
||||
|
||||
# train and eval
|
||||
num_sanity_val_steps: 1
|
||||
optimizer_args:
|
||||
optimizer_cls: modules.optimizer.muon.Muon_AdamW
|
||||
lr: 0.0006
|
||||
muon_args:
|
||||
weight_decay: 0.1
|
||||
adamw_args:
|
||||
weight_decay: 0.0
|
||||
lr_scheduler_args:
|
||||
step_size: 5000
|
||||
gamma: 0.8
|
||||
max_batch_frames: 80000
|
||||
max_batch_size: 48
|
||||
dataset_size_key: 'lengths'
|
||||
val_check_interval: 4000
|
||||
num_valid_plots: 10
|
||||
max_updates: 80000
|
||||
num_ckpt_keep: 5
|
||||
permanent_ckpt_start: 30000
|
||||
permanent_ckpt_interval: 10000
|
||||
|
||||
finetune_enabled: false
|
||||
finetune_ckpt_path: null
|
||||
finetune_ignored_params:
|
||||
- model.spk_embed
|
||||
- model.fs2.txt_embed
|
||||
- model.fs2.encoder.embed_tokens
|
||||
finetune_strict_shapes: true
|
||||
|
||||
freezing_enabled: false
|
||||
frozen_params: []
|
||||
@@ -0,0 +1,7 @@
|
||||
*.ds
|
||||
*.onnx
|
||||
*.npy
|
||||
*.wav
|
||||
temp/
|
||||
cache/
|
||||
assets/
|
||||
@@ -0,0 +1,32 @@
|
||||
import numpy as np
|
||||
import onnxruntime as ort
|
||||
import tqdm
|
||||
|
||||
n_tokens = 10
|
||||
n_frames = 100
|
||||
n_runs = 20
|
||||
speedup = 20
|
||||
provider = 'DmlExecutionProvider'
|
||||
|
||||
tokens = np.array([[1] * n_tokens], dtype=np.int64)
|
||||
durations = np.array([[n_frames // n_tokens] * n_tokens], dtype=np.int64)
|
||||
f0 = np.array([[440.] * n_frames], dtype=np.float32)
|
||||
speedup = np.array(speedup, dtype=np.int64)
|
||||
|
||||
session = ort.InferenceSession('model1.onnx', providers=[provider])
|
||||
for _ in tqdm.tqdm(range(n_runs)):
|
||||
session.run(['mel'], {
|
||||
'tokens': tokens,
|
||||
'durations': durations,
|
||||
'f0': f0,
|
||||
'speedup': speedup
|
||||
})
|
||||
|
||||
session = ort.InferenceSession('model2.onnx', providers=[provider])
|
||||
for _ in tqdm.tqdm(range(n_runs)):
|
||||
session.run(['mel'], {
|
||||
'tokens': tokens,
|
||||
'durations': durations,
|
||||
'f0': f0,
|
||||
'speedup': speedup
|
||||
})
|
||||
@@ -0,0 +1,16 @@
|
||||
import numpy as np
|
||||
import onnxruntime as ort
|
||||
import tqdm
|
||||
|
||||
n_frames = 1000
|
||||
n_runs = 20
|
||||
mel = np.random.randn(1, n_frames, 128).astype(np.float32)
|
||||
f0 = np.random.randn(1, n_frames).astype(np.float32) + 440.
|
||||
provider = 'DmlExecutionProvider'
|
||||
|
||||
session = ort.InferenceSession('nsf_hifigan.onnx', providers=[provider])
|
||||
for _ in tqdm.tqdm(range(n_runs)):
|
||||
session.run(['waveform'], {
|
||||
'mel': mel,
|
||||
'f0': f0
|
||||
})
|
||||
@@ -0,0 +1,3 @@
|
||||
from .acoustic_exporter import DiffSingerAcousticExporter
|
||||
from .variance_exporter import DiffSingerVarianceExporter
|
||||
from .nsf_hifigan_exporter import NSFHiFiGANExporter
|
||||
@@ -0,0 +1,424 @@
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Union, List, Tuple, Dict
|
||||
|
||||
import onnx
|
||||
import onnxsim
|
||||
import torch
|
||||
import yaml
|
||||
|
||||
from basics.base_exporter import BaseExporter
|
||||
from deployment.modules.toplevel import DiffSingerAcousticONNX
|
||||
from modules.fastspeech.param_adaptor import VARIANCE_CHECKLIST
|
||||
from utils import load_ckpt, onnx_helper, remove_suffix
|
||||
from utils.hparams import hparams
|
||||
from utils.phoneme_utils import load_phoneme_dictionary
|
||||
|
||||
|
||||
class DiffSingerAcousticExporter(BaseExporter):
|
||||
def __init__(
|
||||
self,
|
||||
device: Union[str, torch.device] = 'cpu',
|
||||
cache_dir: Path = None,
|
||||
ckpt_steps: int = None,
|
||||
freeze_gender: float = None,
|
||||
freeze_velocity: bool = False,
|
||||
export_spk: List[Tuple[str, Dict[str, float]]] = None,
|
||||
freeze_spk: Tuple[str, Dict[str, float]] = None
|
||||
):
|
||||
super().__init__(device=device, cache_dir=cache_dir)
|
||||
# Basic attributes
|
||||
self.model_name: str = hparams['exp_name']
|
||||
self.ckpt_steps: int = ckpt_steps
|
||||
self.spk_map: dict = self.build_spk_map()
|
||||
self.lang_map: dict = self.build_lang_map()
|
||||
self.phoneme_dictionary = load_phoneme_dictionary()
|
||||
self.use_lang_id = hparams.get('use_lang_id', False) and len(self.phoneme_dictionary.cross_lingual_phonemes) > 0
|
||||
self.model = self.build_model()
|
||||
self.fs2_aux_cache_path = self.cache_dir / (
|
||||
'fs2_aux.onnx' if self.model.use_shallow_diffusion else 'fs2.onnx'
|
||||
)
|
||||
self.diffusion_cache_path = self.cache_dir / 'diffusion.onnx'
|
||||
|
||||
# Attributes for logging
|
||||
self.model_class_name = remove_suffix(self.model.__class__.__name__, 'ONNX')
|
||||
fs2_aux_cls_logging = [remove_suffix(self.model.fs2.__class__.__name__, 'ONNX')]
|
||||
if self.model.use_shallow_diffusion:
|
||||
fs2_aux_cls_logging.append(remove_suffix(
|
||||
self.model.aux_decoder.decoder.__class__.__name__, 'ONNX'
|
||||
))
|
||||
self.fs2_aux_class_name = ', '.join(fs2_aux_cls_logging)
|
||||
self.aux_decoder_class_name = remove_suffix(
|
||||
self.model.aux_decoder.decoder.__class__.__name__, 'ONNX'
|
||||
) if self.model.use_shallow_diffusion else None
|
||||
self.backbone_class_name = remove_suffix(self.model.diffusion.backbone.__class__.__name__, 'ONNX')
|
||||
self.diffusion_class_name = remove_suffix(self.model.diffusion.__class__.__name__, 'ONNX')
|
||||
|
||||
# Attributes for exporting
|
||||
self.expose_gender = freeze_gender is None
|
||||
self.expose_velocity = not freeze_velocity
|
||||
self.freeze_spk: Tuple[str, Dict[str, float]] = freeze_spk \
|
||||
if hparams['use_spk_id'] else None
|
||||
self.export_spk: List[Tuple[str, Dict[str, float]]] = export_spk \
|
||||
if hparams['use_spk_id'] and export_spk is not None else []
|
||||
if hparams['use_key_shift_embed'] and not self.expose_gender:
|
||||
shift_min, shift_max = hparams['augmentation_args']['random_pitch_shifting']['range']
|
||||
key_shift = freeze_gender * shift_max if freeze_gender >= 0. else freeze_gender * abs(shift_min)
|
||||
key_shift = max(min(key_shift, shift_max), shift_min) # clip key shift
|
||||
self.model.fs2.register_buffer('frozen_key_shift', torch.FloatTensor([key_shift]).to(self.device))
|
||||
if hparams['use_spk_id']:
|
||||
if not self.export_spk and self.freeze_spk is None:
|
||||
# In case the user did not specify any speaker settings:
|
||||
if len(self.spk_map) == 1:
|
||||
# If there is only one speaker, freeze him/her.
|
||||
first_spk = next(iter(self.spk_map.keys()))
|
||||
self.freeze_spk = (first_spk, {first_spk: 1.0})
|
||||
else:
|
||||
# If there are multiple speakers, export them all.
|
||||
self.export_spk = [(name, {name: 1.0}) for name in self.spk_map.keys()]
|
||||
if self.freeze_spk is not None:
|
||||
self.model.fs2.register_buffer('frozen_spk_embed', self._perform_spk_mix(self.freeze_spk[1]))
|
||||
|
||||
def build_model(self) -> DiffSingerAcousticONNX:
|
||||
model = DiffSingerAcousticONNX(
|
||||
vocab_size=len(self.phoneme_dictionary),
|
||||
out_dims=hparams['audio_num_mel_bins'],
|
||||
cross_lingual_token_idx=sorted({
|
||||
self.phoneme_dictionary.encode_one(p)
|
||||
for p in self.phoneme_dictionary.cross_lingual_phonemes
|
||||
})
|
||||
).eval().to(self.device)
|
||||
load_ckpt(model, hparams['work_dir'], ckpt_steps=self.ckpt_steps,
|
||||
prefix_in_ckpt='model', strict=True, device=self.device)
|
||||
return model
|
||||
|
||||
def export(self, path: Path):
|
||||
path.mkdir(parents=True, exist_ok=True)
|
||||
model_name = self.model_name
|
||||
if self.freeze_spk is not None:
|
||||
model_name += '.' + self.freeze_spk[0]
|
||||
self.export_model(path / f'{model_name}.onnx')
|
||||
self.export_attachments(path)
|
||||
|
||||
def export_model(self, path: Path):
|
||||
self._torch_export_model()
|
||||
fs2_aux_onnx = self._optimize_fs2_aux_graph(onnx.load(self.fs2_aux_cache_path))
|
||||
diffusion_onnx = self._optimize_diffusion_graph(onnx.load(self.diffusion_cache_path))
|
||||
model_onnx = self._merge_fs2_aux_diffusion_graphs(fs2_aux_onnx, diffusion_onnx)
|
||||
onnx.save(model_onnx, path)
|
||||
self.fs2_aux_cache_path.unlink()
|
||||
self.diffusion_cache_path.unlink()
|
||||
print(f'| export model => {path}')
|
||||
|
||||
def export_attachments(self, path: Path):
|
||||
for spk in self.export_spk:
|
||||
self._export_spk_embed(
|
||||
path / f'{self.model_name}.{spk[0]}.emb',
|
||||
self._perform_spk_mix(spk[1])
|
||||
)
|
||||
self.export_dictionaries(path)
|
||||
self._export_phonemes(path)
|
||||
|
||||
model_name = self.model_name
|
||||
if self.freeze_spk is not None:
|
||||
model_name += '.' + self.freeze_spk[0]
|
||||
dsconfig = {
|
||||
# basic configs
|
||||
'phonemes': f'{self.model_name}.phonemes.json',
|
||||
'languages': f'{self.model_name}.languages.json',
|
||||
'use_lang_id': self.use_lang_id,
|
||||
'acoustic': f'{model_name}.onnx',
|
||||
'hidden_size': hparams['hidden_size'],
|
||||
'vocoder': 'pc_nsf_hifigan_44.1k_hop512_128bin_2025.02',
|
||||
}
|
||||
# multi-speaker
|
||||
if len(self.export_spk) > 0:
|
||||
dsconfig['speakers'] = [f'{self.model_name}.{spk[0]}' for spk in self.export_spk]
|
||||
# parameters
|
||||
if self.expose_gender:
|
||||
dsconfig['augmentation_args'] = {
|
||||
'random_pitch_shifting': {
|
||||
'range': hparams['augmentation_args']['random_pitch_shifting']['range']
|
||||
}
|
||||
}
|
||||
dsconfig['use_key_shift_embed'] = self.expose_gender
|
||||
dsconfig['use_speed_embed'] = self.expose_velocity
|
||||
for variance in VARIANCE_CHECKLIST:
|
||||
dsconfig[f'use_{variance}_embed'] = (variance in self.model.fs2.variance_embed_list)
|
||||
# sampling acceleration and shallow diffusion
|
||||
dsconfig['use_continuous_acceleration'] = True
|
||||
dsconfig['use_variable_depth'] = self.model.use_shallow_diffusion
|
||||
dsconfig['max_depth'] = 1 - self.model.diffusion.t_start
|
||||
# mel specification
|
||||
dsconfig['sample_rate'] = hparams['audio_sample_rate']
|
||||
dsconfig['hop_size'] = hparams['hop_size']
|
||||
dsconfig['win_size'] = hparams['win_size']
|
||||
dsconfig['fft_size'] = hparams['fft_size']
|
||||
dsconfig['num_mel_bins'] = hparams['audio_num_mel_bins']
|
||||
dsconfig['mel_fmin'] = hparams['fmin']
|
||||
dsconfig['mel_fmax'] = hparams['fmax'] if hparams['fmax'] is not None else hparams['audio_sample_rate'] / 2
|
||||
dsconfig['mel_base'] = 'e'
|
||||
dsconfig['mel_scale'] = 'slaney'
|
||||
config_path = path / 'dsconfig.yaml'
|
||||
with open(config_path, 'w', encoding='utf8') as fw:
|
||||
yaml.safe_dump(dsconfig, fw, sort_keys=False)
|
||||
print(f'| export configs => {config_path} **PLEASE EDIT BEFORE USE**')
|
||||
|
||||
@torch.no_grad()
|
||||
def _torch_export_model(self):
|
||||
# Prepare inputs for FastSpeech2 and aux decoder tracing
|
||||
n_frames = 10
|
||||
tokens = torch.LongTensor([[1]]).to(self.device)
|
||||
durations = torch.LongTensor([[n_frames]]).to(self.device)
|
||||
f0 = torch.FloatTensor([[440.] * n_frames]).to(self.device)
|
||||
variances = {
|
||||
v_name: torch.zeros(1, n_frames, dtype=torch.float32, device=self.device)
|
||||
for v_name in self.model.fs2.variance_embed_list
|
||||
}
|
||||
kwargs: Dict[str, torch.Tensor] = {}
|
||||
arguments = (tokens, durations, f0, variances, kwargs)
|
||||
input_names = ['tokens', 'durations', 'f0'] + self.model.fs2.variance_embed_list
|
||||
dynamic_axes = {
|
||||
'tokens': {
|
||||
1: 'n_tokens'
|
||||
},
|
||||
'durations': {
|
||||
1: 'n_tokens'
|
||||
},
|
||||
'f0': {
|
||||
1: 'n_frames'
|
||||
},
|
||||
**{
|
||||
v_name: {
|
||||
1: 'n_frames'
|
||||
}
|
||||
for v_name in self.model.fs2.variance_embed_list
|
||||
}
|
||||
}
|
||||
if hparams['use_key_shift_embed']:
|
||||
if self.expose_gender:
|
||||
kwargs['gender'] = torch.rand((1, n_frames), dtype=torch.float32, device=self.device)
|
||||
input_names.append('gender')
|
||||
dynamic_axes['gender'] = {
|
||||
1: 'n_frames'
|
||||
}
|
||||
if hparams['use_speed_embed']:
|
||||
if self.expose_velocity:
|
||||
kwargs['velocity'] = torch.rand((1, n_frames), dtype=torch.float32, device=self.device)
|
||||
input_names.append('velocity')
|
||||
dynamic_axes['velocity'] = {
|
||||
1: 'n_frames'
|
||||
}
|
||||
if hparams['use_spk_id'] and not self.freeze_spk:
|
||||
kwargs['spk_embed'] = torch.rand(
|
||||
(1, n_frames, hparams['hidden_size']),
|
||||
dtype=torch.float32, device=self.device
|
||||
)
|
||||
input_names.append('spk_embed')
|
||||
dynamic_axes['spk_embed'] = {
|
||||
1: 'n_frames'
|
||||
}
|
||||
if self.use_lang_id:
|
||||
kwargs['languages'] = torch.zeros_like(tokens)
|
||||
input_names.append('languages')
|
||||
dynamic_axes['languages'] = {
|
||||
1: 'n_tokens'
|
||||
}
|
||||
dynamic_axes['condition'] = {
|
||||
1: 'n_frames'
|
||||
}
|
||||
|
||||
# PyTorch ONNX export for FastSpeech2 and aux decoder
|
||||
output_names = ['condition']
|
||||
if self.model.use_shallow_diffusion:
|
||||
output_names.append('aux_mel')
|
||||
dynamic_axes['aux_mel'] = {
|
||||
1: 'n_frames'
|
||||
}
|
||||
print(f'Exporting {self.fs2_aux_class_name}...')
|
||||
torch.onnx.export(
|
||||
self.model.view_as_fs2_aux(),
|
||||
arguments,
|
||||
self.fs2_aux_cache_path,
|
||||
input_names=input_names,
|
||||
output_names=output_names,
|
||||
dynamic_axes=dynamic_axes,
|
||||
opset_version=17,
|
||||
**onnx_helper.TORCHSCRIPT_EXPORT_KWARGS
|
||||
)
|
||||
|
||||
condition = torch.rand((1, n_frames, hparams['hidden_size']), device=self.device)
|
||||
|
||||
# Prepare inputs for backbone tracing and GaussianDiffusion scripting
|
||||
shape = (1, 1, hparams['audio_num_mel_bins'], n_frames)
|
||||
noise = torch.randn(shape, device=self.device)
|
||||
x_aux = torch.randn((1, n_frames, hparams['audio_num_mel_bins']), device=self.device)
|
||||
dummy_time = (torch.rand((1,), device=self.device) * self.model.diffusion.time_scale_factor).float()
|
||||
dummy_depth = torch.tensor(0.1, device=self.device)
|
||||
dummy_steps = 5
|
||||
|
||||
print(f'Tracing {self.backbone_class_name} backbone...')
|
||||
if self.model.diffusion_type == 'ddpm':
|
||||
major_mel_decoder = self.model.view_as_diffusion()
|
||||
elif self.model.diffusion_type == 'reflow':
|
||||
major_mel_decoder = self.model.view_as_reflow()
|
||||
else:
|
||||
raise ValueError(f'Invalid diffusion type: {self.model.diffusion_type}')
|
||||
major_mel_decoder.diffusion.set_backbone(
|
||||
torch.jit.trace(
|
||||
major_mel_decoder.diffusion.backbone,
|
||||
(
|
||||
noise,
|
||||
dummy_time,
|
||||
condition.transpose(1, 2)
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
print(f'Scripting {self.diffusion_class_name}...')
|
||||
diffusion_inputs = [
|
||||
condition,
|
||||
*([x_aux, dummy_depth] if self.model.use_shallow_diffusion else [])
|
||||
]
|
||||
major_mel_decoder = torch.jit.script(
|
||||
major_mel_decoder,
|
||||
example_inputs=[
|
||||
(
|
||||
*diffusion_inputs,
|
||||
1 # p_sample branch
|
||||
),
|
||||
(
|
||||
*diffusion_inputs,
|
||||
dummy_steps # p_sample_plms branch
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
# PyTorch ONNX export for GaussianDiffusion
|
||||
print(f'Exporting {self.diffusion_class_name}...')
|
||||
torch.onnx.export(
|
||||
major_mel_decoder,
|
||||
(
|
||||
*diffusion_inputs,
|
||||
dummy_steps
|
||||
),
|
||||
self.diffusion_cache_path,
|
||||
input_names=[
|
||||
'condition',
|
||||
*(['x_aux', 'depth'] if self.model.use_shallow_diffusion else []),
|
||||
'steps'
|
||||
],
|
||||
output_names=[
|
||||
'mel'
|
||||
],
|
||||
dynamic_axes={
|
||||
'condition': {
|
||||
1: 'n_frames'
|
||||
},
|
||||
**({'x_aux': {1: 'n_frames'}} if self.model.use_shallow_diffusion else {}),
|
||||
'mel': {
|
||||
1: 'n_frames'
|
||||
}
|
||||
},
|
||||
opset_version=17,
|
||||
**onnx_helper.TORCHSCRIPT_EXPORT_KWARGS
|
||||
)
|
||||
|
||||
@torch.no_grad()
|
||||
def _perform_spk_mix(self, spk_mix: Dict[str, float]):
|
||||
spk_mix_ids = []
|
||||
spk_mix_values = []
|
||||
for name, value in spk_mix.items():
|
||||
spk_mix_ids.append(self.spk_map[name])
|
||||
assert value >= 0., f'Speaker mix checks failed.\n' \
|
||||
f'Proportion of speaker \'{name}\' is negative.'
|
||||
spk_mix_values.append(value)
|
||||
spk_mix_id_N = torch.LongTensor(spk_mix_ids).to(self.device)[None] # => [1, N]
|
||||
spk_mix_value_N = torch.FloatTensor(spk_mix_values).to(self.device)[None] # => [1, N]
|
||||
spk_mix_value_sum = spk_mix_value_N.sum()
|
||||
assert spk_mix_value_sum > 0., 'Speaker mix checks failed.\n' \
|
||||
'Proportions of speaker mix sum to zero.'
|
||||
spk_mix_value_N /= spk_mix_value_sum # normalize
|
||||
spk_mix_embed = torch.sum(
|
||||
self.model.fs2.spk_embed(spk_mix_id_N) * spk_mix_value_N.unsqueeze(2), # => [1, N, H]
|
||||
dim=1, keepdim=True
|
||||
) # => [1, 1, H]
|
||||
return spk_mix_embed
|
||||
|
||||
def _optimize_fs2_aux_graph(self, fs2: onnx.ModelProto) -> onnx.ModelProto:
|
||||
print(f'Running ONNX Simplifier on {self.fs2_aux_class_name}...')
|
||||
fs2, check = onnxsim.simplify(fs2, include_subgraph=True)
|
||||
assert check, 'Simplified ONNX model could not be validated'
|
||||
onnx_helper.model_reorder_io_list(
|
||||
fs2, 'input',
|
||||
target_name='languages', insert_after_name='tokens'
|
||||
)
|
||||
print(f'| optimize graph: {self.fs2_aux_class_name}')
|
||||
return fs2
|
||||
|
||||
def _optimize_diffusion_graph(self, diffusion: onnx.ModelProto) -> onnx.ModelProto:
|
||||
onnx_helper.model_override_io_shapes(diffusion, output_shapes={
|
||||
'mel': (1, 'n_frames', hparams['audio_num_mel_bins'])
|
||||
})
|
||||
print(f'Running ONNX Simplifier #1 on {self.diffusion_class_name}...')
|
||||
diffusion, check = onnxsim.simplify(diffusion, include_subgraph=True)
|
||||
assert check, 'Simplified ONNX model could not be validated'
|
||||
onnx_helper.graph_fold_back_to_squeeze(diffusion.graph)
|
||||
onnx_helper.graph_extract_conditioner_projections(
|
||||
graph=diffusion.graph, op_type='Conv',
|
||||
weight_pattern=r'diffusion\..*\.conditioner_projection\.weight',
|
||||
alias_prefix='/diffusion/backbone/cache'
|
||||
)
|
||||
onnx_helper.graph_remove_unused_values(diffusion.graph)
|
||||
print(f'Running ONNX Simplifier #2 on {self.diffusion_class_name}...')
|
||||
diffusion, check = onnxsim.simplify(
|
||||
diffusion,
|
||||
include_subgraph=True
|
||||
)
|
||||
assert check, 'Simplified ONNX model could not be validated'
|
||||
print(f'| optimize graph: {self.diffusion_class_name}')
|
||||
return diffusion
|
||||
|
||||
def _merge_fs2_aux_diffusion_graphs(self, fs2: onnx.ModelProto, diffusion: onnx.ModelProto) -> onnx.ModelProto:
|
||||
onnx_helper.model_add_prefixes(
|
||||
fs2, dim_prefix=('fs2aux.' if self.model.use_shallow_diffusion else 'fs2.'),
|
||||
ignored_pattern=r'(n_tokens)|(n_frames)'
|
||||
)
|
||||
onnx_helper.model_add_prefixes(diffusion, dim_prefix='diffusion.', ignored_pattern='n_frames')
|
||||
print(f'Merging {self.fs2_aux_class_name} and {self.diffusion_class_name} '
|
||||
f'back into {self.model_class_name}...')
|
||||
merged = onnx.compose.merge_models(
|
||||
fs2, diffusion, io_map=[
|
||||
('condition', 'condition'),
|
||||
*([('aux_mel', 'x_aux')] if self.model.use_shallow_diffusion else []),
|
||||
],
|
||||
prefix1='', prefix2='', doc_string='',
|
||||
producer_name=fs2.producer_name, producer_version=fs2.producer_version,
|
||||
domain=fs2.domain, model_version=fs2.model_version
|
||||
)
|
||||
merged.graph.name = fs2.graph.name
|
||||
|
||||
print(f'Running ONNX Simplifier on {self.model_class_name}...')
|
||||
merged, check = onnxsim.simplify(
|
||||
merged,
|
||||
include_subgraph=True
|
||||
)
|
||||
assert check, 'Simplified ONNX model could not be validated'
|
||||
print(f'| optimize graph: {self.model_class_name}')
|
||||
|
||||
return merged
|
||||
|
||||
# noinspection PyMethodMayBeStatic
|
||||
def _export_spk_embed(self, path: Path, spk_embed: torch.Tensor):
|
||||
with open(path, 'wb') as f:
|
||||
f.write(spk_embed.cpu().numpy().tobytes())
|
||||
print(f'| export spk embed => {path}')
|
||||
|
||||
def _export_phonemes(self, path: Path):
|
||||
ph_path = path / f'{self.model_name}.phonemes.json'
|
||||
self.phoneme_dictionary.dump(ph_path)
|
||||
print(f'| export phonemes => {ph_path}')
|
||||
lang_path = path / f'{self.model_name}.languages.json'
|
||||
with open(lang_path, 'w', encoding='utf8') as f:
|
||||
json.dump(self.lang_map, f, ensure_ascii=False, indent=2)
|
||||
print(f'| export languages => {lang_path}')
|
||||
@@ -0,0 +1,127 @@
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Union
|
||||
|
||||
import onnx
|
||||
import onnxsim
|
||||
import torch
|
||||
import yaml
|
||||
from torch import nn
|
||||
|
||||
from basics.base_exporter import BaseExporter
|
||||
from deployment.modules.nsf_hifigan import NSFHiFiGANONNX
|
||||
from utils import load_ckpt, onnx_helper, remove_suffix
|
||||
from utils.hparams import hparams
|
||||
|
||||
|
||||
class NSFHiFiGANExporter(BaseExporter):
|
||||
def __init__(
|
||||
self,
|
||||
device: Union[str, torch.device] = 'cpu',
|
||||
cache_dir: Path = None,
|
||||
model_path: Path = None,
|
||||
model_name: str = 'nsf_hifigan'
|
||||
):
|
||||
super().__init__(device=device, cache_dir=cache_dir)
|
||||
self.model_path = model_path
|
||||
self.model_name = model_name
|
||||
self.vocoder_pitch_controllable = False
|
||||
self.model = self.build_model()
|
||||
self.model_class_name = remove_suffix(self.model.__class__.__name__, 'ONNX')
|
||||
self.model_cache_path = (self.cache_dir / self.model_name).with_suffix('.onnx')
|
||||
|
||||
def build_model(self) -> nn.Module:
|
||||
config_path = self.model_path.with_name('config.json')
|
||||
with open(config_path, 'r', encoding='utf8') as f:
|
||||
config = json.load(f)
|
||||
assert hparams.get('mel_base') == 'e', (
|
||||
"Mel base must be set to \'e\' according to 2nd stage of the migration plan. "
|
||||
"See https://github.com/openvpi/DiffSinger/releases/tag/v2.3.0 for more details."
|
||||
)
|
||||
model = NSFHiFiGANONNX(config).eval().to(self.device)
|
||||
self.vocoder_pitch_controllable = config.get("pc_aug", False)
|
||||
load_ckpt(model.generator, str(self.model_path),
|
||||
prefix_in_ckpt=None, key_in_ckpt='generator',
|
||||
strict=True, device=self.device)
|
||||
model.generator.remove_weight_norm()
|
||||
return model
|
||||
|
||||
def export(self, path: Path):
|
||||
path.mkdir(parents=True, exist_ok=True)
|
||||
self.export_model(path / self.model_cache_path.name)
|
||||
self.export_attachments(path)
|
||||
|
||||
def export_model(self, path: Path):
|
||||
self._torch_export_model()
|
||||
model_onnx = self._optimize_model_graph(onnx.load(self.model_cache_path))
|
||||
onnx.save(model_onnx, path)
|
||||
self.model_cache_path.unlink()
|
||||
print(f'| export model => {path}')
|
||||
|
||||
def export_attachments(self, path: Path):
|
||||
config_path = path / 'vocoder.yaml'
|
||||
with open(config_path, 'w', encoding='utf8') as fw:
|
||||
yaml.safe_dump({
|
||||
# basic configs
|
||||
'name': self.model_name,
|
||||
'model': self.model_cache_path.name,
|
||||
# mel specifications
|
||||
'sample_rate': hparams['audio_sample_rate'],
|
||||
'hop_size': hparams['hop_size'],
|
||||
'win_size': hparams['win_size'],
|
||||
'fft_size': hparams['fft_size'],
|
||||
'num_mel_bins': hparams['audio_num_mel_bins'],
|
||||
'mel_fmin': hparams['fmin'],
|
||||
'mel_fmax': hparams['fmax'] if hparams['fmax'] is not None else hparams['audio_sample_rate'] / 2,
|
||||
'mel_base': 'e',
|
||||
'mel_scale': 'slaney',
|
||||
'pitch_controllable': self.vocoder_pitch_controllable,
|
||||
# Some old vocoder versions may have severe performance issues on CUDA;
|
||||
# the issues were fixed in newer versions, and this flag is to distinguish them
|
||||
'force_on_cpu': False,
|
||||
}, fw, sort_keys=False)
|
||||
print(f'| export configs => {config_path} **PLEASE EDIT BEFORE USE**')
|
||||
|
||||
@torch.no_grad()
|
||||
def _torch_export_model(self):
|
||||
# Prepare inputs for NSFHiFiGAN
|
||||
n_frames = 10
|
||||
mel = torch.randn((1, n_frames, hparams['audio_num_mel_bins']), dtype=torch.float32, device=self.device)
|
||||
f0 = torch.randn((1, n_frames), dtype=torch.float32, device=self.device) + 440.
|
||||
|
||||
# PyTorch ONNX export for NSFHiFiGAN
|
||||
print(f'Exporting {self.model_class_name}...')
|
||||
torch.onnx.export(
|
||||
self.model,
|
||||
(
|
||||
mel,
|
||||
f0
|
||||
),
|
||||
self.model_cache_path,
|
||||
input_names=[
|
||||
'mel',
|
||||
'f0'
|
||||
],
|
||||
output_names=[
|
||||
'waveform'
|
||||
],
|
||||
dynamic_axes={
|
||||
'mel': {
|
||||
1: 'n_frames'
|
||||
},
|
||||
'f0': {
|
||||
1: 'n_frames'
|
||||
},
|
||||
'waveform': {
|
||||
1: 'n_samples'
|
||||
}
|
||||
},
|
||||
opset_version=17,
|
||||
**onnx_helper.TORCHSCRIPT_EXPORT_KWARGS
|
||||
)
|
||||
|
||||
def _optimize_model_graph(self, model: onnx.ModelProto) -> onnx.ModelProto:
|
||||
print(f'Running ONNX simplifier for {self.model_class_name}...')
|
||||
model, check = onnxsim.simplify(model, include_subgraph=True)
|
||||
assert check, 'Simplified ONNX model could not be validated'
|
||||
return model
|
||||
@@ -0,0 +1,808 @@
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Union, List, Tuple, Dict
|
||||
|
||||
import onnx
|
||||
import onnxsim
|
||||
import torch
|
||||
import yaml
|
||||
|
||||
from basics.base_exporter import BaseExporter
|
||||
from deployment.modules.toplevel import DiffSingerVarianceONNX
|
||||
from modules.fastspeech.param_adaptor import VARIANCE_CHECKLIST
|
||||
from utils import load_ckpt, onnx_helper, remove_suffix
|
||||
from utils.hparams import hparams
|
||||
from utils.phoneme_utils import load_phoneme_dictionary
|
||||
|
||||
|
||||
class DiffSingerVarianceExporter(BaseExporter):
|
||||
def __init__(
|
||||
self,
|
||||
device: Union[str, torch.device] = 'cpu',
|
||||
cache_dir: Path = None,
|
||||
ckpt_steps: int = None,
|
||||
freeze_glide: bool = False,
|
||||
freeze_expr: bool = False,
|
||||
export_spk: List[Tuple[str, Dict[str, float]]] = None,
|
||||
freeze_spk: Tuple[str, Dict[str, float]] = None
|
||||
):
|
||||
super().__init__(device=device, cache_dir=cache_dir)
|
||||
# Basic attributes
|
||||
self.model_name: str = hparams['exp_name']
|
||||
self.ckpt_steps: int = ckpt_steps
|
||||
self.spk_map: dict = self.build_spk_map()
|
||||
self.lang_map: dict = self.build_lang_map()
|
||||
self.phoneme_dictionary = load_phoneme_dictionary()
|
||||
self.use_lang_id = hparams.get('use_lang_id', False) and len(self.phoneme_dictionary.cross_lingual_phonemes) > 0
|
||||
self.model = self.build_model()
|
||||
self.linguistic_encoder_cache_path = self.cache_dir / 'linguistic.onnx'
|
||||
self.dur_predictor_cache_path = self.cache_dir / 'dur.onnx'
|
||||
self.pitch_preprocess_cache_path = self.cache_dir / 'pitch_pre.onnx'
|
||||
self.pitch_predictor_cache_path = self.cache_dir / 'pitch.onnx'
|
||||
self.pitch_postprocess_cache_path = self.cache_dir / 'pitch_post.onnx'
|
||||
self.variance_preprocess_cache_path = self.cache_dir / 'variance_pre.onnx'
|
||||
self.multi_var_predictor_cache_path = self.cache_dir / 'variance.onnx'
|
||||
self.variance_postprocess_cache_path = self.cache_dir / 'variance_post.onnx'
|
||||
|
||||
# Attributes for logging
|
||||
self.fs2_class_name = remove_suffix(self.model.fs2.__class__.__name__, 'ONNX')
|
||||
self.dur_predictor_class_name = \
|
||||
remove_suffix(self.model.fs2.dur_predictor.__class__.__name__, 'ONNX') \
|
||||
if self.model.predict_dur else None
|
||||
self.pitch_backbone_class_name = \
|
||||
remove_suffix(self.model.pitch_predictor.backbone.__class__.__name__, 'ONNX') \
|
||||
if self.model.predict_pitch else None
|
||||
self.pitch_predictor_class_name = \
|
||||
remove_suffix(self.model.pitch_predictor.__class__.__name__, 'ONNX') \
|
||||
if self.model.predict_pitch else None
|
||||
self.variance_backbone_class_name = \
|
||||
remove_suffix(self.model.variance_predictor.backbone.__class__.__name__, 'ONNX') \
|
||||
if self.model.predict_variances else None
|
||||
self.multi_var_predictor_class_name = \
|
||||
remove_suffix(self.model.variance_predictor.__class__.__name__, 'ONNX') \
|
||||
if self.model.predict_variances else None
|
||||
|
||||
# Attributes for exporting
|
||||
self.expose_expr = not freeze_expr
|
||||
self.freeze_glide = freeze_glide
|
||||
self.freeze_spk: Tuple[str, Dict[str, float]] = freeze_spk \
|
||||
if hparams['use_spk_id'] else None
|
||||
self.export_spk: List[Tuple[str, Dict[str, float]]] = export_spk \
|
||||
if hparams['use_spk_id'] and export_spk is not None else []
|
||||
if hparams['use_spk_id']:
|
||||
if not self.export_spk and self.freeze_spk is None:
|
||||
# In case the user did not specify any speaker settings:
|
||||
if len(self.spk_map) == 1:
|
||||
# If there is only one speaker, freeze him/her.
|
||||
first_spk = next(iter(self.spk_map.keys()))
|
||||
self.freeze_spk = (first_spk, {first_spk: 1.0})
|
||||
else:
|
||||
# If there are multiple speakers, export them all.
|
||||
self.export_spk = [(name, {name: 1.0}) for name in self.spk_map.keys()]
|
||||
if self.freeze_spk is not None:
|
||||
self.model.register_buffer('frozen_spk_embed', self._perform_spk_mix(self.freeze_spk[1]))
|
||||
|
||||
def build_model(self) -> DiffSingerVarianceONNX:
|
||||
model = DiffSingerVarianceONNX(
|
||||
vocab_size=len(self.phoneme_dictionary),
|
||||
cross_lingual_token_idx=sorted({
|
||||
self.phoneme_dictionary.encode_one(p)
|
||||
for p in self.phoneme_dictionary.cross_lingual_phonemes
|
||||
})
|
||||
).eval().to(self.device)
|
||||
load_ckpt(model, hparams['work_dir'], ckpt_steps=self.ckpt_steps,
|
||||
prefix_in_ckpt='model', strict=True, device=self.device)
|
||||
model.build_smooth_op(self.device)
|
||||
return model
|
||||
|
||||
def export(self, path: Path):
|
||||
path.mkdir(parents=True, exist_ok=True)
|
||||
model_name = self.model_name
|
||||
if self.freeze_spk is not None:
|
||||
model_name += '.' + self.freeze_spk[0]
|
||||
self.export_model(path, model_name)
|
||||
self.export_attachments(path)
|
||||
|
||||
def export_model(self, path: Path, model_name: str = None):
|
||||
self._torch_export_model()
|
||||
linguistic_onnx = self._optimize_linguistic_graph(onnx.load(self.linguistic_encoder_cache_path))
|
||||
linguistic_path = path / f'{model_name}.linguistic.onnx'
|
||||
onnx.save(linguistic_onnx, linguistic_path)
|
||||
print(f'| export linguistic encoder => {linguistic_path}')
|
||||
self.linguistic_encoder_cache_path.unlink()
|
||||
if self.model.predict_dur:
|
||||
dur_predictor_onnx = self._optimize_dur_predictor_graph(onnx.load(self.dur_predictor_cache_path))
|
||||
dur_predictor_path = path / f'{model_name}.dur.onnx'
|
||||
onnx.save(dur_predictor_onnx, dur_predictor_path)
|
||||
self.dur_predictor_cache_path.unlink()
|
||||
print(f'| export dur predictor => {dur_predictor_path}')
|
||||
if self.model.predict_pitch:
|
||||
pitch_predictor_onnx = self._optimize_merge_pitch_predictor_graph(
|
||||
onnx.load(self.pitch_preprocess_cache_path),
|
||||
onnx.load(self.pitch_predictor_cache_path),
|
||||
onnx.load(self.pitch_postprocess_cache_path)
|
||||
)
|
||||
pitch_predictor_path = path / f'{model_name}.pitch.onnx'
|
||||
onnx.save(pitch_predictor_onnx, pitch_predictor_path)
|
||||
self.pitch_preprocess_cache_path.unlink()
|
||||
self.pitch_predictor_cache_path.unlink()
|
||||
self.pitch_postprocess_cache_path.unlink()
|
||||
print(f'| export pitch predictor => {pitch_predictor_path}')
|
||||
if self.model.predict_variances:
|
||||
variance_predictor_onnx = self._optimize_merge_variance_predictor_graph(
|
||||
onnx.load(self.variance_preprocess_cache_path),
|
||||
onnx.load(self.multi_var_predictor_cache_path),
|
||||
onnx.load(self.variance_postprocess_cache_path)
|
||||
)
|
||||
variance_predictor_path = path / f'{model_name}.variance.onnx'
|
||||
onnx.save(variance_predictor_onnx, variance_predictor_path)
|
||||
self.variance_preprocess_cache_path.unlink()
|
||||
self.multi_var_predictor_cache_path.unlink()
|
||||
self.variance_postprocess_cache_path.unlink()
|
||||
print(f'| export variance predictor => {variance_predictor_path}')
|
||||
|
||||
def export_attachments(self, path: Path):
|
||||
for spk in self.export_spk:
|
||||
self._export_spk_embed(
|
||||
path / f'{self.model_name}.{spk[0]}.emb',
|
||||
self._perform_spk_mix(spk[1])
|
||||
)
|
||||
self.export_dictionaries(path)
|
||||
self._export_phonemes(path)
|
||||
|
||||
model_name = self.model_name
|
||||
if self.freeze_spk is not None:
|
||||
model_name += '.' + self.freeze_spk[0]
|
||||
dsconfig = {
|
||||
# basic configs
|
||||
'phonemes': f'{self.model_name}.phonemes.json',
|
||||
'languages': f'{self.model_name}.languages.json',
|
||||
'use_lang_id': self.use_lang_id,
|
||||
'linguistic': f'{model_name}.linguistic.onnx',
|
||||
'hidden_size': self.model.hidden_size,
|
||||
'predict_dur': self.model.predict_dur,
|
||||
}
|
||||
# multi-speaker
|
||||
if len(self.export_spk) > 0:
|
||||
dsconfig['speakers'] = [f'{self.model_name}.{spk[0]}' for spk in self.export_spk]
|
||||
# functionalities
|
||||
if self.model.predict_dur:
|
||||
dsconfig['dur'] = f'{model_name}.dur.onnx'
|
||||
if self.model.predict_pitch:
|
||||
dsconfig['pitch'] = f'{model_name}.pitch.onnx'
|
||||
dsconfig['use_expr'] = self.expose_expr
|
||||
dsconfig['use_note_rest'] = self.model.use_melody_encoder
|
||||
if self.model.predict_variances:
|
||||
dsconfig['variance'] = f'{model_name}.variance.onnx'
|
||||
for variance in VARIANCE_CHECKLIST:
|
||||
dsconfig[f'predict_{variance}'] = (variance in self.model.variance_prediction_list)
|
||||
# sampling acceleration
|
||||
dsconfig['use_continuous_acceleration'] = True
|
||||
# frame specifications
|
||||
dsconfig['sample_rate'] = hparams['audio_sample_rate']
|
||||
dsconfig['hop_size'] = hparams['hop_size']
|
||||
config_path = path / 'dsconfig.yaml'
|
||||
with open(config_path, 'w', encoding='utf8') as fw:
|
||||
yaml.safe_dump(dsconfig, fw, sort_keys=False)
|
||||
print(f'| export configs => {config_path} **PLEASE EDIT BEFORE USE**')
|
||||
|
||||
@torch.no_grad()
|
||||
def _torch_export_model(self):
|
||||
# Prepare inputs for FastSpeech2 and dur predictor tracing
|
||||
tokens = torch.LongTensor([[1] * 5]).to(self.device)
|
||||
ph_dur = torch.LongTensor([[3, 5, 2, 1, 4]]).to(self.device)
|
||||
word_div = torch.LongTensor([[2, 2, 1]]).to(self.device)
|
||||
word_dur = torch.LongTensor([[8, 3, 4]]).to(self.device)
|
||||
languages = torch.LongTensor([[0] * 5]).to(self.device)
|
||||
encoder_out = torch.rand(1, 5, hparams['hidden_size'], dtype=torch.float32, device=self.device)
|
||||
x_masks = tokens == 0
|
||||
ph_midi = torch.LongTensor([[60] * 5]).to(self.device)
|
||||
encoder_output_names = ['encoder_out', 'x_masks']
|
||||
encoder_common_axes = {
|
||||
'encoder_out': {
|
||||
1: 'n_tokens'
|
||||
},
|
||||
'x_masks': {
|
||||
1: 'n_tokens'
|
||||
}
|
||||
}
|
||||
input_lang_id = self.use_lang_id
|
||||
input_spk_embed = hparams['use_spk_id'] and not self.freeze_spk
|
||||
|
||||
print(f'Exporting {self.fs2_class_name}...')
|
||||
if self.model.predict_dur:
|
||||
torch.onnx.export(
|
||||
self.model.view_as_linguistic_encoder(),
|
||||
(
|
||||
tokens,
|
||||
word_div,
|
||||
word_dur,
|
||||
*([languages] if input_lang_id else [])
|
||||
),
|
||||
self.linguistic_encoder_cache_path,
|
||||
input_names=[
|
||||
'tokens',
|
||||
'word_div',
|
||||
'word_dur',
|
||||
*(['languages'] if input_lang_id else [])
|
||||
],
|
||||
output_names=encoder_output_names,
|
||||
dynamic_axes={
|
||||
'tokens': {
|
||||
1: 'n_tokens'
|
||||
},
|
||||
'word_div': {
|
||||
1: 'n_words'
|
||||
},
|
||||
'word_dur': {
|
||||
1: 'n_words'
|
||||
},
|
||||
**encoder_common_axes,
|
||||
**({'languages': {1: 'n_tokens'}} if input_lang_id else {})
|
||||
},
|
||||
opset_version=17,
|
||||
**onnx_helper.TORCHSCRIPT_EXPORT_KWARGS
|
||||
)
|
||||
|
||||
print(f'Exporting {self.dur_predictor_class_name}...')
|
||||
torch.onnx.export(
|
||||
self.model.view_as_dur_predictor(),
|
||||
(
|
||||
encoder_out,
|
||||
x_masks,
|
||||
ph_midi,
|
||||
*([torch.rand(
|
||||
1, 5, hparams['hidden_size'],
|
||||
dtype=torch.float32, device=self.device
|
||||
)] if input_spk_embed else [])
|
||||
),
|
||||
self.dur_predictor_cache_path,
|
||||
input_names=[
|
||||
'encoder_out',
|
||||
'x_masks',
|
||||
'ph_midi',
|
||||
*(['spk_embed'] if input_spk_embed else [])
|
||||
],
|
||||
output_names=[
|
||||
'ph_dur_pred'
|
||||
],
|
||||
dynamic_axes={
|
||||
'ph_midi': {
|
||||
1: 'n_tokens'
|
||||
},
|
||||
'ph_dur_pred': {
|
||||
1: 'n_tokens'
|
||||
},
|
||||
**({'spk_embed': {1: 'n_tokens'}} if input_spk_embed else {}),
|
||||
**encoder_common_axes
|
||||
},
|
||||
opset_version=17,
|
||||
**onnx_helper.TORCHSCRIPT_EXPORT_KWARGS
|
||||
)
|
||||
else:
|
||||
torch.onnx.export(
|
||||
self.model.view_as_linguistic_encoder(),
|
||||
(
|
||||
tokens,
|
||||
ph_dur,
|
||||
*([languages] if input_lang_id else [])
|
||||
),
|
||||
self.linguistic_encoder_cache_path,
|
||||
input_names=[
|
||||
'tokens',
|
||||
'ph_dur',
|
||||
*(['languages'] if input_lang_id else [])
|
||||
],
|
||||
output_names=encoder_output_names,
|
||||
dynamic_axes={
|
||||
'tokens': {
|
||||
1: 'n_tokens'
|
||||
},
|
||||
'ph_dur': {
|
||||
1: 'n_tokens'
|
||||
},
|
||||
**encoder_common_axes,
|
||||
**({'languages': {1: 'n_tokens'}} if input_lang_id else {})
|
||||
},
|
||||
opset_version=17,
|
||||
**onnx_helper.TORCHSCRIPT_EXPORT_KWARGS
|
||||
)
|
||||
|
||||
# Common dummy inputs
|
||||
dummy_time = (torch.rand((1,), device=self.device) * hparams.get('time_scale_factor', 1.0)).float()
|
||||
dummy_steps = 5
|
||||
|
||||
if self.model.predict_pitch:
|
||||
use_melody_encoder = hparams.get('use_melody_encoder', False)
|
||||
use_glide_embed = use_melody_encoder and hparams['use_glide_embed'] and not self.freeze_glide
|
||||
# Prepare inputs for preprocessor of the pitch predictor
|
||||
note_midi = torch.FloatTensor([[60.] * 4]).to(self.device)
|
||||
note_dur = torch.LongTensor([[2, 6, 3, 4]]).to(self.device)
|
||||
pitch = torch.FloatTensor([[60.] * 15]).to(self.device)
|
||||
retake = torch.ones_like(pitch, dtype=torch.bool)
|
||||
pitch_input_args = (
|
||||
encoder_out,
|
||||
ph_dur,
|
||||
{
|
||||
'note_midi': note_midi,
|
||||
**({'note_rest': note_midi >= 0} if use_melody_encoder else {}),
|
||||
'note_dur': note_dur,
|
||||
**({'note_glide': torch.zeros_like(note_midi, dtype=torch.long)} if use_glide_embed else {}),
|
||||
'pitch': pitch,
|
||||
**({'expr': torch.ones_like(pitch)} if self.expose_expr else {}),
|
||||
'retake': retake,
|
||||
**({'spk_embed': torch.rand(
|
||||
1, 15, hparams['hidden_size'], dtype=torch.float32, device=self.device
|
||||
)} if input_spk_embed else {})
|
||||
}
|
||||
)
|
||||
torch.onnx.export(
|
||||
self.model.view_as_pitch_preprocess(),
|
||||
pitch_input_args,
|
||||
self.pitch_preprocess_cache_path,
|
||||
input_names=[
|
||||
'encoder_out', 'ph_dur', 'note_midi',
|
||||
*(['note_rest'] if use_melody_encoder else []),
|
||||
'note_dur',
|
||||
*(['note_glide'] if use_glide_embed else []),
|
||||
'pitch',
|
||||
*(['expr'] if self.expose_expr else []),
|
||||
'retake',
|
||||
*(['spk_embed'] if input_spk_embed else [])
|
||||
],
|
||||
output_names=[
|
||||
'pitch_cond', 'base_pitch'
|
||||
],
|
||||
dynamic_axes={
|
||||
'encoder_out': {
|
||||
1: 'n_tokens'
|
||||
},
|
||||
'ph_dur': {
|
||||
1: 'n_tokens'
|
||||
},
|
||||
'note_midi': {
|
||||
1: 'n_notes'
|
||||
},
|
||||
**({'note_rest': {1: 'n_notes'}} if use_melody_encoder else {}),
|
||||
'note_dur': {
|
||||
1: 'n_notes'
|
||||
},
|
||||
**({'note_glide': {1: 'n_notes'}} if use_glide_embed else {}),
|
||||
'pitch': {
|
||||
1: 'n_frames'
|
||||
},
|
||||
**({'expr': {1: 'n_frames'}} if self.expose_expr else {}),
|
||||
'retake': {
|
||||
1: 'n_frames'
|
||||
},
|
||||
'pitch_cond': {
|
||||
1: 'n_frames'
|
||||
},
|
||||
'base_pitch': {
|
||||
1: 'n_frames'
|
||||
},
|
||||
**({'spk_embed': {1: 'n_frames'}} if input_spk_embed else {})
|
||||
},
|
||||
opset_version=17,
|
||||
**onnx_helper.TORCHSCRIPT_EXPORT_KWARGS
|
||||
)
|
||||
|
||||
# Prepare inputs for backbone tracing and pitch predictor scripting
|
||||
shape = (1, 1, hparams['pitch_prediction_args']['repeat_bins'], 15)
|
||||
noise = torch.randn(shape, device=self.device)
|
||||
condition = torch.rand((1, hparams['hidden_size'], 15), device=self.device)
|
||||
|
||||
print(f'Tracing {self.pitch_backbone_class_name} backbone...')
|
||||
pitch_predictor = self.model.view_as_pitch_predictor()
|
||||
pitch_predictor.pitch_predictor.set_backbone(
|
||||
torch.jit.trace(
|
||||
pitch_predictor.pitch_predictor.backbone,
|
||||
(
|
||||
noise,
|
||||
dummy_time,
|
||||
condition
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
print(f'Scripting {self.pitch_predictor_class_name}...')
|
||||
pitch_predictor = torch.jit.script(
|
||||
pitch_predictor,
|
||||
example_inputs=[
|
||||
(
|
||||
condition.transpose(1, 2),
|
||||
1 # p_sample branch
|
||||
),
|
||||
(
|
||||
condition.transpose(1, 2),
|
||||
dummy_steps # p_sample_plms branch
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
print(f'Exporting {self.pitch_predictor_class_name}...')
|
||||
torch.onnx.export(
|
||||
pitch_predictor,
|
||||
(
|
||||
condition.transpose(1, 2),
|
||||
dummy_steps
|
||||
),
|
||||
self.pitch_predictor_cache_path,
|
||||
input_names=[
|
||||
'pitch_cond',
|
||||
'steps'
|
||||
],
|
||||
output_names=[
|
||||
'x_pred'
|
||||
],
|
||||
dynamic_axes={
|
||||
'pitch_cond': {
|
||||
1: 'n_frames'
|
||||
},
|
||||
'x_pred': {
|
||||
1: 'n_frames'
|
||||
}
|
||||
},
|
||||
opset_version=17,
|
||||
**onnx_helper.TORCHSCRIPT_EXPORT_KWARGS
|
||||
)
|
||||
|
||||
# Prepare inputs for postprocessor of the multi-variance predictor
|
||||
torch.onnx.export(
|
||||
self.model.view_as_pitch_postprocess(),
|
||||
(
|
||||
pitch,
|
||||
pitch
|
||||
),
|
||||
self.pitch_postprocess_cache_path,
|
||||
input_names=[
|
||||
'x_pred',
|
||||
'base_pitch'
|
||||
],
|
||||
output_names=[
|
||||
'pitch_pred'
|
||||
],
|
||||
dynamic_axes={
|
||||
'x_pred': {
|
||||
1: 'n_frames'
|
||||
},
|
||||
'base_pitch': {
|
||||
1: 'n_frames'
|
||||
},
|
||||
'pitch_pred': {
|
||||
1: 'n_frames'
|
||||
}
|
||||
},
|
||||
opset_version=17,
|
||||
**onnx_helper.TORCHSCRIPT_EXPORT_KWARGS
|
||||
)
|
||||
|
||||
if self.model.predict_variances:
|
||||
total_repeat_bins = hparams['variances_prediction_args']['total_repeat_bins']
|
||||
repeat_bins = total_repeat_bins // len(self.model.variance_prediction_list)
|
||||
|
||||
# Prepare inputs for preprocessor of the multi-variance predictor
|
||||
pitch = torch.FloatTensor([[60.] * 15]).to(self.device)
|
||||
variances = {
|
||||
v_name: torch.FloatTensor([[0.] * 15]).to(self.device)
|
||||
for v_name in self.model.variance_prediction_list
|
||||
}
|
||||
retake = torch.ones_like(pitch, dtype=torch.bool)[..., None].tile(len(self.model.variance_prediction_list))
|
||||
torch.onnx.export(
|
||||
self.model.view_as_variance_preprocess(),
|
||||
(
|
||||
encoder_out,
|
||||
ph_dur,
|
||||
pitch,
|
||||
variances,
|
||||
retake,
|
||||
*([torch.rand(
|
||||
1, 15, hparams['hidden_size'],
|
||||
dtype=torch.float32, device=self.device
|
||||
)] if input_spk_embed else [])
|
||||
),
|
||||
self.variance_preprocess_cache_path,
|
||||
input_names=[
|
||||
'encoder_out', 'ph_dur', 'pitch',
|
||||
*self.model.variance_prediction_list,
|
||||
'retake',
|
||||
*(['spk_embed'] if input_spk_embed else [])
|
||||
],
|
||||
output_names=[
|
||||
'variance_cond'
|
||||
],
|
||||
dynamic_axes={
|
||||
'encoder_out': {
|
||||
1: 'n_tokens'
|
||||
},
|
||||
'ph_dur': {
|
||||
1: 'n_tokens'
|
||||
},
|
||||
'pitch': {
|
||||
1: 'n_frames'
|
||||
},
|
||||
**{
|
||||
v_name: {
|
||||
1: 'n_frames'
|
||||
}
|
||||
for v_name in self.model.variance_prediction_list
|
||||
},
|
||||
'retake': {
|
||||
1: 'n_frames'
|
||||
},
|
||||
**({'spk_embed': {1: 'n_frames'}} if input_spk_embed else {})
|
||||
},
|
||||
opset_version=17,
|
||||
**onnx_helper.TORCHSCRIPT_EXPORT_KWARGS
|
||||
)
|
||||
|
||||
# Prepare inputs for backbone tracing and multi-variance predictor scripting
|
||||
shape = (1, len(self.model.variance_prediction_list), repeat_bins, 15)
|
||||
noise = torch.randn(shape, device=self.device)
|
||||
condition = torch.rand((1, hparams['hidden_size'], 15), device=self.device)
|
||||
step = (torch.rand((1,), device=self.device) * hparams.get('time_scale_factor', hparams['K_step']))
|
||||
|
||||
print(f'Tracing {self.variance_backbone_class_name} backbone...')
|
||||
multi_var_predictor = self.model.view_as_variance_predictor()
|
||||
multi_var_predictor.variance_predictor.set_backbone(
|
||||
torch.jit.trace(
|
||||
multi_var_predictor.variance_predictor.backbone,
|
||||
(
|
||||
noise,
|
||||
step,
|
||||
condition
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
print(f'Scripting {self.multi_var_predictor_class_name}...')
|
||||
multi_var_predictor = torch.jit.script(
|
||||
multi_var_predictor,
|
||||
example_inputs=[
|
||||
(
|
||||
condition.transpose(1, 2),
|
||||
1 # p_sample branch
|
||||
),
|
||||
(
|
||||
condition.transpose(1, 2),
|
||||
dummy_steps # p_sample_plms branch
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
print(f'Exporting {self.multi_var_predictor_class_name}...')
|
||||
torch.onnx.export(
|
||||
multi_var_predictor,
|
||||
(
|
||||
condition.transpose(1, 2),
|
||||
dummy_steps
|
||||
),
|
||||
self.multi_var_predictor_cache_path,
|
||||
input_names=[
|
||||
'variance_cond',
|
||||
'steps'
|
||||
],
|
||||
output_names=[
|
||||
'xs_pred'
|
||||
],
|
||||
dynamic_axes={
|
||||
'variance_cond': {
|
||||
1: 'n_frames'
|
||||
},
|
||||
'xs_pred': {
|
||||
(1 if len(self.model.variance_prediction_list) == 1 else 2): 'n_frames'
|
||||
}
|
||||
},
|
||||
opset_version=17,
|
||||
**onnx_helper.TORCHSCRIPT_EXPORT_KWARGS
|
||||
)
|
||||
|
||||
# Prepare inputs for postprocessor of the multi-variance predictor
|
||||
xs_shape = (1, 15) \
|
||||
if len(self.model.variance_prediction_list) == 1 \
|
||||
else (1, len(self.model.variance_prediction_list), 15)
|
||||
xs_pred = torch.randn(xs_shape, dtype=torch.float32, device=self.device)
|
||||
torch.onnx.export(
|
||||
self.model.view_as_variance_postprocess(),
|
||||
(
|
||||
xs_pred
|
||||
),
|
||||
self.variance_postprocess_cache_path,
|
||||
input_names=[
|
||||
'xs_pred'
|
||||
],
|
||||
output_names=[
|
||||
f'{v_name}_pred'
|
||||
for v_name in self.model.variance_prediction_list
|
||||
],
|
||||
dynamic_axes={
|
||||
'xs_pred': {
|
||||
(1 if len(self.model.variance_prediction_list) == 1 else 2): 'n_frames'
|
||||
},
|
||||
**{
|
||||
f'{v_name}_pred': {
|
||||
1: 'n_frames'
|
||||
}
|
||||
for v_name in self.model.variance_prediction_list
|
||||
}
|
||||
},
|
||||
opset_version=17,
|
||||
**onnx_helper.TORCHSCRIPT_EXPORT_KWARGS
|
||||
)
|
||||
|
||||
@torch.no_grad()
|
||||
def _perform_spk_mix(self, spk_mix: Dict[str, float]):
|
||||
spk_mix_ids = []
|
||||
spk_mix_values = []
|
||||
for name, value in spk_mix.items():
|
||||
spk_mix_ids.append(self.spk_map[name])
|
||||
assert value >= 0., f'Speaker mix checks failed.\n' \
|
||||
f'Proportion of speaker \'{name}\' is negative.'
|
||||
spk_mix_values.append(value)
|
||||
spk_mix_id_N = torch.LongTensor(spk_mix_ids).to(self.device)[None] # => [1, N]
|
||||
spk_mix_value_N = torch.FloatTensor(spk_mix_values).to(self.device)[None] # => [1, N]
|
||||
spk_mix_value_sum = spk_mix_value_N.sum()
|
||||
assert spk_mix_value_sum > 0., 'Speaker mix checks failed.\n' \
|
||||
'Proportions of speaker mix sum to zero.'
|
||||
spk_mix_value_N /= spk_mix_value_sum # normalize
|
||||
spk_mix_embed = torch.sum(
|
||||
self.model.spk_embed(spk_mix_id_N) * spk_mix_value_N.unsqueeze(2), # => [1, N, H]
|
||||
dim=1, keepdim=True
|
||||
) # => [1, 1, H]
|
||||
return spk_mix_embed
|
||||
|
||||
def _optimize_linguistic_graph(self, linguistic: onnx.ModelProto) -> onnx.ModelProto:
|
||||
onnx_helper.model_override_io_shapes(
|
||||
linguistic,
|
||||
output_shapes={
|
||||
'encoder_out': (1, 'n_tokens', hparams['hidden_size'])
|
||||
}
|
||||
)
|
||||
print(f'Running ONNX Simplifier on {self.fs2_class_name}...')
|
||||
linguistic, check = onnxsim.simplify(linguistic, include_subgraph=True)
|
||||
assert check, 'Simplified ONNX model could not be validated'
|
||||
onnx_helper.model_reorder_io_list(
|
||||
linguistic, 'input',
|
||||
target_name='languages', insert_after_name='tokens'
|
||||
)
|
||||
print(f'| optimize graph: {self.fs2_class_name}')
|
||||
return linguistic
|
||||
|
||||
def _optimize_dur_predictor_graph(self, dur_predictor: onnx.ModelProto) -> onnx.ModelProto:
|
||||
onnx_helper.model_override_io_shapes(
|
||||
dur_predictor,
|
||||
output_shapes={
|
||||
'ph_dur_pred': (1, 'n_tokens')
|
||||
}
|
||||
)
|
||||
print(f'Running ONNX Simplifier on {self.dur_predictor_class_name}...')
|
||||
dur_predictor, check = onnxsim.simplify(dur_predictor, include_subgraph=True)
|
||||
assert check, 'Simplified ONNX model could not be validated'
|
||||
print(f'| optimize graph: {self.dur_predictor_class_name}')
|
||||
return dur_predictor
|
||||
|
||||
def _optimize_merge_pitch_predictor_graph(
|
||||
self, pitch_pre: onnx.ModelProto, pitch_predictor: onnx.ModelProto, pitch_post: onnx.ModelProto
|
||||
) -> onnx.ModelProto:
|
||||
onnx_helper.model_override_io_shapes(
|
||||
pitch_pre, output_shapes={'pitch_cond': (1, 'n_frames', hparams['hidden_size'])}
|
||||
)
|
||||
pitch_pre, check = onnxsim.simplify(pitch_pre, include_subgraph=True)
|
||||
assert check, 'Simplified ONNX model could not be validated'
|
||||
|
||||
onnx_helper.model_override_io_shapes(
|
||||
pitch_predictor, output_shapes={'pitch_pred': (1, 'n_frames')}
|
||||
)
|
||||
print(f'Running ONNX Simplifier #1 on {self.pitch_predictor_class_name}...')
|
||||
pitch_predictor, check = onnxsim.simplify(pitch_predictor, include_subgraph=True)
|
||||
assert check, 'Simplified ONNX model could not be validated'
|
||||
onnx_helper.graph_fold_back_to_squeeze(pitch_predictor.graph)
|
||||
onnx_helper.graph_extract_conditioner_projections(
|
||||
graph=pitch_predictor.graph, op_type='Conv',
|
||||
weight_pattern=r'pitch_predictor\..*\.conditioner_projection\.weight',
|
||||
alias_prefix='/pitch_predictor/backbone/cache'
|
||||
)
|
||||
onnx_helper.graph_remove_unused_values(pitch_predictor.graph)
|
||||
print(f'Running ONNX Simplifier #2 on {self.pitch_predictor_class_name}...')
|
||||
pitch_predictor, check = onnxsim.simplify(pitch_predictor, include_subgraph=True)
|
||||
assert check, 'Simplified ONNX model could not be validated'
|
||||
|
||||
onnx_helper.model_add_prefixes(pitch_pre, node_prefix='/pre', ignored_pattern=r'.*embed.*')
|
||||
onnx_helper.model_add_prefixes(pitch_pre, dim_prefix='pre.', ignored_pattern='(n_tokens)|(n_notes)|(n_frames)')
|
||||
onnx_helper.model_add_prefixes(pitch_post, node_prefix='/post', ignored_pattern=None)
|
||||
onnx_helper.model_add_prefixes(pitch_post, dim_prefix='post.', ignored_pattern='n_frames')
|
||||
pitch_pre_diffusion = onnx.compose.merge_models(
|
||||
pitch_pre, pitch_predictor, io_map=[('pitch_cond', 'pitch_cond')],
|
||||
prefix1='', prefix2='', doc_string='',
|
||||
producer_name=pitch_pre.producer_name, producer_version=pitch_pre.producer_version,
|
||||
domain=pitch_pre.domain, model_version=pitch_pre.model_version
|
||||
)
|
||||
pitch_pre_diffusion.graph.name = pitch_pre.graph.name
|
||||
pitch_predictor = onnx.compose.merge_models(
|
||||
pitch_pre_diffusion, pitch_post, io_map=[
|
||||
('x_pred', 'x_pred'), ('base_pitch', 'base_pitch')
|
||||
], prefix1='', prefix2='', doc_string='',
|
||||
producer_name=pitch_pre.producer_name, producer_version=pitch_pre.producer_version,
|
||||
domain=pitch_pre.domain, model_version=pitch_pre.model_version
|
||||
)
|
||||
pitch_predictor.graph.name = pitch_pre.graph.name
|
||||
|
||||
print(f'| optimize graph: {self.pitch_predictor_class_name}')
|
||||
return pitch_predictor
|
||||
|
||||
def _optimize_merge_variance_predictor_graph(
|
||||
self, var_pre: onnx.ModelProto, var_diffusion: onnx.ModelProto, var_post: onnx.ModelProto
|
||||
):
|
||||
onnx_helper.model_override_io_shapes(
|
||||
var_pre, output_shapes={'variance_cond': (1, 'n_frames', hparams['hidden_size'])}
|
||||
)
|
||||
var_pre, check = onnxsim.simplify(var_pre, include_subgraph=True)
|
||||
assert check, 'Simplified ONNX model could not be validated'
|
||||
|
||||
onnx_helper.model_override_io_shapes(
|
||||
var_diffusion, output_shapes={
|
||||
'xs_pred': (1, 'n_frames')
|
||||
if len(self.model.variance_prediction_list) == 1
|
||||
else (1, len(self.model.variance_prediction_list), 'n_frames')
|
||||
}
|
||||
)
|
||||
print(f'Running ONNX Simplifier #1 on {self.multi_var_predictor_class_name}...')
|
||||
var_diffusion, check = onnxsim.simplify(var_diffusion, include_subgraph=True)
|
||||
assert check, 'Simplified ONNX model could not be validated'
|
||||
onnx_helper.graph_fold_back_to_squeeze(var_diffusion.graph)
|
||||
onnx_helper.graph_extract_conditioner_projections(
|
||||
graph=var_diffusion.graph, op_type='Conv',
|
||||
weight_pattern=r'variance_predictor\..*\.conditioner_projection\.weight',
|
||||
alias_prefix='/variance_predictor/backbone/cache'
|
||||
)
|
||||
onnx_helper.graph_remove_unused_values(var_diffusion.graph)
|
||||
print(f'Running ONNX Simplifier #2 on {self.multi_var_predictor_class_name}...')
|
||||
var_diffusion, check = onnxsim.simplify(var_diffusion, include_subgraph=True)
|
||||
assert check, 'Simplified ONNX model could not be validated'
|
||||
|
||||
var_post, check = onnxsim.simplify(var_post, include_subgraph=True)
|
||||
assert check, 'Simplified ONNX model could not be validated'
|
||||
|
||||
ignored_variance_names = '|'.join([f'({v_name})' for v_name in self.model.variance_prediction_list])
|
||||
onnx_helper.model_add_prefixes(
|
||||
var_pre, node_prefix='/pre', value_info_prefix='/pre', initializer_prefix='/pre',
|
||||
ignored_pattern=fr'.*((embed)|{ignored_variance_names}).*'
|
||||
)
|
||||
onnx_helper.model_add_prefixes(var_pre, dim_prefix='pre.', ignored_pattern='(n_tokens)|(n_frames)')
|
||||
onnx_helper.model_add_prefixes(
|
||||
var_post, node_prefix='/post', value_info_prefix='/post', initializer_prefix='/post',
|
||||
ignored_pattern=None
|
||||
)
|
||||
onnx_helper.model_add_prefixes(var_post, dim_prefix='post.', ignored_pattern='n_frames')
|
||||
|
||||
print(f'Merging {self.multi_var_predictor_class_name} subroutines...')
|
||||
var_pre_diffusion = onnx.compose.merge_models(
|
||||
var_pre, var_diffusion, io_map=[('variance_cond', 'variance_cond')],
|
||||
prefix1='', prefix2='', doc_string='',
|
||||
producer_name=var_pre.producer_name, producer_version=var_pre.producer_version,
|
||||
domain=var_pre.domain, model_version=var_pre.model_version
|
||||
)
|
||||
var_pre_diffusion.graph.name = var_pre.graph.name
|
||||
var_predictor = onnx.compose.merge_models(
|
||||
var_pre_diffusion, var_post, io_map=[('xs_pred', 'xs_pred')],
|
||||
prefix1='', prefix2='', doc_string='',
|
||||
producer_name=var_pre.producer_name, producer_version=var_pre.producer_version,
|
||||
domain=var_pre.domain, model_version=var_pre.model_version
|
||||
)
|
||||
var_predictor.graph.name = var_pre.graph.name
|
||||
return var_predictor
|
||||
|
||||
# noinspection PyMethodMayBeStatic
|
||||
def _export_spk_embed(self, path: Path, spk_embed: torch.Tensor):
|
||||
with open(path, 'wb') as f:
|
||||
f.write(spk_embed.cpu().numpy().tobytes())
|
||||
print(f'| export spk embed => {path}')
|
||||
|
||||
def _export_phonemes(self, path: Path):
|
||||
ph_path = path / f'{self.model_name}.phonemes.json'
|
||||
self.phoneme_dictionary.dump(ph_path)
|
||||
print(f'| export phonemes => {ph_path}')
|
||||
lang_path = path / f'{self.model_name}.languages.json'
|
||||
with open(lang_path, 'w', encoding='utf8') as fw:
|
||||
json.dump(self.lang_map, fw, ensure_ascii=False, indent=2)
|
||||
print(f'| export languages => {lang_path}')
|
||||
@@ -0,0 +1,220 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import List, Tuple
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
|
||||
from modules.core import (
|
||||
GaussianDiffusion, PitchDiffusion, MultiVarianceDiffusion
|
||||
)
|
||||
|
||||
|
||||
def extract(a, t):
|
||||
return a[t].reshape((1, 1, 1, 1))
|
||||
|
||||
|
||||
# noinspection PyMethodOverriding
|
||||
class GaussianDiffusionONNX(GaussianDiffusion):
|
||||
@property
|
||||
def backbone(self):
|
||||
return self.denoise_fn
|
||||
|
||||
# We give up the setter for the property `backbone` because this will cause TorchScript to fail
|
||||
# @backbone.setter
|
||||
@torch.jit.unused
|
||||
def set_backbone(self, value):
|
||||
self.denoise_fn = value
|
||||
|
||||
def q_sample(self, x_start, t, noise):
|
||||
return (
|
||||
extract(self.sqrt_alphas_cumprod, t) * x_start +
|
||||
extract(self.sqrt_one_minus_alphas_cumprod, t) * noise
|
||||
)
|
||||
|
||||
def p_sample(self, x, t, cond):
|
||||
x_pred = self.denoise_fn(x, t, cond)
|
||||
x_recon = (
|
||||
extract(self.sqrt_recip_alphas_cumprod, t) * x -
|
||||
extract(self.sqrt_recipm1_alphas_cumprod, t) * x_pred
|
||||
)
|
||||
# This is previously inherited from original DiffSinger repository
|
||||
# and disabled due to some loudness issues when speedup = 1.
|
||||
# x_recon = torch.clamp(x_recon, min=-1., max=1.)
|
||||
|
||||
model_mean = (
|
||||
extract(self.posterior_mean_coef1, t) * x_recon +
|
||||
extract(self.posterior_mean_coef2, t) * x
|
||||
)
|
||||
model_log_variance = extract(self.posterior_log_variance_clipped, t)
|
||||
noise = torch.randn_like(x)
|
||||
# no noise when t == 0
|
||||
nonzero_mask = ((t > 0).float()).reshape(1, 1, 1, 1)
|
||||
return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise
|
||||
|
||||
def p_sample_ddim(self, x, t, interval: int, cond):
|
||||
a_t = extract(self.alphas_cumprod, t)
|
||||
t_prev = t - interval
|
||||
a_prev = extract(self.alphas_cumprod, t_prev * (t_prev > 0))
|
||||
|
||||
noise_pred = self.denoise_fn(x, t, cond=cond)
|
||||
x_prev = a_prev.sqrt() * (
|
||||
x / a_t.sqrt() + (((1 - a_prev) / a_prev).sqrt() - ((1 - a_t) / a_t).sqrt()) * noise_pred
|
||||
)
|
||||
return x_prev
|
||||
|
||||
def plms_get_x_pred(self, x, noise_t, t, t_prev):
|
||||
a_t = extract(self.alphas_cumprod, t)
|
||||
a_prev = extract(self.alphas_cumprod, t_prev)
|
||||
a_t_sq, a_prev_sq = a_t.sqrt(), a_prev.sqrt()
|
||||
|
||||
x_delta = (a_prev - a_t) * ((1 / (a_t_sq * (a_t_sq + a_prev_sq))) * x - 1 / (
|
||||
a_t_sq * (((1 - a_prev) * a_t).sqrt() + ((1 - a_t) * a_prev).sqrt())) * noise_t)
|
||||
x_pred = x + x_delta
|
||||
|
||||
return x_pred
|
||||
|
||||
def p_sample_plms(self, x_prev, t, interval: int, cond, noise_list: List[Tensor], stage: int):
|
||||
noise_pred = self.denoise_fn(x_prev, t, cond)
|
||||
t_prev = t - interval
|
||||
t_prev = t_prev * (t_prev > 0)
|
||||
if stage == 0:
|
||||
x_pred = self.plms_get_x_pred(x_prev, noise_pred, t, t_prev)
|
||||
noise_pred_prev = self.denoise_fn(x_pred, t_prev, cond)
|
||||
noise_pred_prime = (noise_pred + noise_pred_prev) / 2.
|
||||
elif stage == 1:
|
||||
noise_pred_prime = (3. * noise_pred - noise_list[-1]) / 2.
|
||||
elif stage == 2:
|
||||
noise_pred_prime = (23. * noise_pred - 16. * noise_list[-1] + 5. * noise_list[-2]) / 12.
|
||||
else:
|
||||
noise_pred_prime = (55. * noise_pred - 59. * noise_list[-1] + 37.
|
||||
* noise_list[-2] - 9. * noise_list[-3]) / 24.
|
||||
x_prev = self.plms_get_x_pred(x_prev, noise_pred_prime, t, t_prev)
|
||||
return noise_pred, x_prev
|
||||
|
||||
def norm_spec(self, x):
|
||||
k = (self.spec_max - self.spec_min) / 2.
|
||||
b = (self.spec_max + self.spec_min) / 2.
|
||||
return (x - b) / k
|
||||
|
||||
def denorm_spec(self, x):
|
||||
k = (self.spec_max - self.spec_min) / 2.
|
||||
b = (self.spec_max + self.spec_min) / 2.
|
||||
return x * k + b
|
||||
|
||||
def forward(self, condition, x_start=None, depth=None, steps: int = 10):
|
||||
condition = condition.transpose(1, 2) # [1, T, H] => [1, H, T]
|
||||
device = condition.device
|
||||
n_frames = condition.shape[2]
|
||||
|
||||
noise = torch.randn((1, self.num_feats, self.out_dims, n_frames), device=device)
|
||||
if x_start is None:
|
||||
speedup = max(1, self.timesteps // steps)
|
||||
speedup = self.timestep_factors[torch.sum(self.timestep_factors <= speedup) - 1]
|
||||
step_range = torch.arange(0, self.k_step, speedup, dtype=torch.long, device=device).flip(0)[:, None]
|
||||
x = noise
|
||||
else:
|
||||
depth_int64 = min(torch.round(depth * self.timesteps).long(), self.k_step)
|
||||
speedup = max(1, depth_int64 // steps)
|
||||
depth_int64 = depth_int64 // speedup * speedup # make depth_int64 a multiple of speedup
|
||||
step_range = torch.arange(0, depth_int64, speedup, dtype=torch.long, device=device).flip(0)[:, None]
|
||||
x_start = self.norm_spec(x_start).transpose(-2, -1)
|
||||
if self.num_feats == 1:
|
||||
x_start = x_start[:, None, :, :]
|
||||
if depth_int64 >= self.timesteps:
|
||||
x = noise
|
||||
elif depth_int64 > 0:
|
||||
x = self.q_sample(
|
||||
x_start, torch.full((1,), depth_int64 - 1, device=device, dtype=torch.long), noise
|
||||
)
|
||||
else:
|
||||
x = x_start
|
||||
|
||||
if speedup > 1:
|
||||
for t in step_range:
|
||||
x = self.p_sample_ddim(x, t, interval=speedup, cond=condition)
|
||||
# plms_noise_stage: int = 0
|
||||
# noise_list: List[Tensor] = []
|
||||
# for t in step_range:
|
||||
# noise_pred, x = self.p_sample_plms(
|
||||
# x, t, interval=speedup, cond=condition,
|
||||
# noise_list=noise_list, stage=plms_noise_stage
|
||||
# )
|
||||
# if plms_noise_stage == 0:
|
||||
# noise_list = [noise_pred]
|
||||
# plms_noise_stage = plms_noise_stage + 1
|
||||
# else:
|
||||
# if plms_noise_stage >= 3:
|
||||
# noise_list.pop(0)
|
||||
# else:
|
||||
# plms_noise_stage = plms_noise_stage + 1
|
||||
# noise_list.append(noise_pred)
|
||||
else:
|
||||
for t in step_range:
|
||||
x = self.p_sample(x, t, cond=condition)
|
||||
|
||||
if self.num_feats == 1:
|
||||
x = x.squeeze(1).permute(0, 2, 1) # [B, 1, M, T] => [B, T, M]
|
||||
else:
|
||||
x = x.permute(0, 1, 3, 2) # [B, F, M, T] => [B, F, T, M]
|
||||
x = self.denorm_spec(x)
|
||||
return x
|
||||
|
||||
|
||||
class PitchDiffusionONNX(GaussianDiffusionONNX, PitchDiffusion):
|
||||
def __init__(self, vmin: float, vmax: float,
|
||||
cmin: float, cmax: float, repeat_bins,
|
||||
timesteps=1000, k_step=1000,
|
||||
backbone_type=None, backbone_args=None,
|
||||
betas=None):
|
||||
self.vmin = vmin
|
||||
self.vmax = vmax
|
||||
self.cmin = cmin
|
||||
self.cmax = cmax
|
||||
super(PitchDiffusion, self).__init__(
|
||||
vmin=vmin, vmax=vmax, repeat_bins=repeat_bins,
|
||||
timesteps=timesteps, k_step=k_step,
|
||||
backbone_type=backbone_type, backbone_args=backbone_args,
|
||||
betas=betas
|
||||
)
|
||||
|
||||
def clamp_spec(self, x):
|
||||
return x.clamp(min=self.cmin, max=self.cmax)
|
||||
|
||||
def denorm_spec(self, x):
|
||||
d = (self.spec_max - self.spec_min) / 2.
|
||||
m = (self.spec_max + self.spec_min) / 2.
|
||||
x = x * d + m
|
||||
x = x.mean(dim=-1)
|
||||
return x
|
||||
|
||||
|
||||
class MultiVarianceDiffusionONNX(GaussianDiffusionONNX, MultiVarianceDiffusion):
|
||||
def __init__(
|
||||
self, ranges: List[Tuple[float, float]],
|
||||
clamps: List[Tuple[float | None, float | None] | None],
|
||||
repeat_bins, timesteps=1000, k_step=1000,
|
||||
backbone_type=None, backbone_args=None,
|
||||
betas=None
|
||||
):
|
||||
assert len(ranges) == len(clamps)
|
||||
self.clamps = clamps
|
||||
vmin = [r[0] for r in ranges]
|
||||
vmax = [r[1] for r in ranges]
|
||||
if len(vmin) == 1:
|
||||
vmin = vmin[0]
|
||||
if len(vmax) == 1:
|
||||
vmax = vmax[0]
|
||||
super(MultiVarianceDiffusion, self).__init__(
|
||||
vmin=vmin, vmax=vmax, repeat_bins=repeat_bins,
|
||||
timesteps=timesteps, k_step=k_step,
|
||||
backbone_type=backbone_type, backbone_args=backbone_args,
|
||||
betas=betas
|
||||
)
|
||||
|
||||
def denorm_spec(self, x):
|
||||
d = (self.spec_max - self.spec_min) / 2.
|
||||
m = (self.spec_max + self.spec_min) / 2.
|
||||
x = x * d + m
|
||||
x = x.mean(dim=-1)
|
||||
return x
|
||||
@@ -0,0 +1,235 @@
|
||||
import copy
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from modules.commons.common_layers import NormalInitEmbedding as Embedding
|
||||
from modules.fastspeech.acoustic_encoder import FastSpeech2Acoustic
|
||||
from modules.fastspeech.variance_encoder import FastSpeech2Variance
|
||||
from utils.hparams import hparams
|
||||
from utils.phoneme_utils import PAD_INDEX
|
||||
|
||||
f0_bin = 256
|
||||
f0_max = 1100.0
|
||||
f0_min = 50.0
|
||||
f0_mel_min = 1127 * np.log(1 + f0_min / 700)
|
||||
f0_mel_max = 1127 * np.log(1 + f0_max / 700)
|
||||
|
||||
|
||||
def uniform_attention_pooling(spk_embed, durations):
|
||||
_, T_mel, _ = spk_embed.shape
|
||||
ph_starts = torch.cumsum(torch.cat([torch.zeros_like(durations[:, :1]), durations[:, :-1]], dim=1), dim=1)
|
||||
ph_ends = ph_starts + durations
|
||||
mel_indices = torch.arange(T_mel, device=spk_embed.device).view(1, 1, T_mel)
|
||||
phoneme_to_mel_mask = (mel_indices >= ph_starts.unsqueeze(-1)) & (mel_indices < ph_ends.unsqueeze(-1))
|
||||
uniform_scores = phoneme_to_mel_mask.float()
|
||||
sum_scores = uniform_scores.sum(dim=2, keepdim=True)
|
||||
attn_weights = uniform_scores / (sum_scores + (sum_scores == 0).float()) # [B, T_ph, T_mel]
|
||||
ph_spk_embed = torch.bmm(attn_weights, spk_embed)
|
||||
|
||||
return ph_spk_embed
|
||||
|
||||
|
||||
def f0_to_coarse(f0):
|
||||
f0_mel = 1127 * (1 + f0 / 700).log()
|
||||
a = (f0_bin - 2) / (f0_mel_max - f0_mel_min)
|
||||
b = f0_mel_min * a - 1.
|
||||
f0_mel = torch.where(f0_mel > 0, f0_mel * a - b, f0_mel)
|
||||
torch.clip_(f0_mel, min=1., max=float(f0_bin - 1))
|
||||
f0_coarse = torch.round(f0_mel).long()
|
||||
return f0_coarse
|
||||
|
||||
|
||||
class LengthRegulator(nn.Module):
|
||||
# noinspection PyMethodMayBeStatic
|
||||
def forward(self, dur):
|
||||
token_idx = torch.arange(1, dur.shape[1] + 1, device=dur.device)[None, :, None]
|
||||
dur_cumsum = torch.cumsum(dur, dim=1)
|
||||
dur_cumsum_prev = F.pad(dur_cumsum, (1, -1), mode='constant', value=0)
|
||||
pos_idx = torch.arange(dur.sum(dim=1).max(), device=dur.device)[None, None]
|
||||
token_mask = (pos_idx >= dur_cumsum_prev[:, :, None]) & (pos_idx < dur_cumsum[:, :, None])
|
||||
mel2ph = (token_idx * token_mask).sum(dim=1)
|
||||
return mel2ph
|
||||
|
||||
|
||||
class FastSpeech2AcousticONNX(FastSpeech2Acoustic):
|
||||
def __init__(self, vocab_size, cross_lingual_token_idx=None):
|
||||
super().__init__(vocab_size=vocab_size)
|
||||
self.register_buffer(
|
||||
'cross_lingual_token_idx',
|
||||
torch.LongTensor(cross_lingual_token_idx),
|
||||
persistent=False
|
||||
) # [N,]
|
||||
if len(cross_lingual_token_idx) == 0:
|
||||
self.use_lang_id = False
|
||||
|
||||
# for temporary compatibility; will be completely removed in the future
|
||||
self.f0_embed_type = hparams.get('f0_embed_type', 'continuous')
|
||||
if self.f0_embed_type == 'discrete':
|
||||
self.pitch_embed = Embedding(300, hparams['hidden_size'], PAD_INDEX)
|
||||
|
||||
self.lr = LengthRegulator()
|
||||
if hparams['use_key_shift_embed']:
|
||||
self.shift_min, self.shift_max = hparams['augmentation_args']['random_pitch_shifting']['range']
|
||||
if hparams['use_speed_embed']:
|
||||
self.speed_min, self.speed_max = hparams['augmentation_args']['random_time_stretching']['range']
|
||||
|
||||
# noinspection PyMethodOverriding
|
||||
def forward(
|
||||
self, tokens, durations,
|
||||
f0, variances: dict,
|
||||
gender=None, velocity=None,
|
||||
spk_embed=None,
|
||||
languages=None
|
||||
):
|
||||
txt_embed = self.txt_embed(tokens)
|
||||
durations = durations * (tokens > 0)
|
||||
mel2ph = self.lr(durations)
|
||||
_mel2ph = mel2ph
|
||||
f0 = f0 * (mel2ph > 0)
|
||||
mel2ph = mel2ph[..., None].repeat((1, 1, hparams['hidden_size']))
|
||||
if self.use_variance_scaling:
|
||||
dur_embed = self.dur_embed(torch.log(1 + durations.float())[:, :, None])
|
||||
else:
|
||||
dur_embed = self.dur_embed(durations.float()[:, :, None])
|
||||
if self.use_lang_id:
|
||||
lang_mask = torch.any(
|
||||
tokens[..., None] == self.cross_lingual_token_idx[None, None],
|
||||
dim=-1
|
||||
)
|
||||
lang_embed = self.lang_embed(languages * lang_mask)
|
||||
extra_embed = dur_embed + lang_embed
|
||||
else:
|
||||
extra_embed = dur_embed
|
||||
if hparams.get('use_mix_ln', False):
|
||||
if hasattr(self, 'frozen_spk_embed'):
|
||||
ph_spk_embed = self.frozen_spk_embed.repeat(1, tokens.shape[1], 1)
|
||||
else:
|
||||
ph_spk_embed = uniform_attention_pooling(spk_embed, durations)
|
||||
else:
|
||||
ph_spk_embed = None
|
||||
encoded = self.encoder(txt_embed, extra_embed, tokens == PAD_INDEX, spk_embed=ph_spk_embed)
|
||||
encoded = F.pad(encoded, (0, 0, 1, 0))
|
||||
condition = torch.gather(encoded, 1, mel2ph)
|
||||
|
||||
if self.use_stretch_embed:
|
||||
stretch = torch.round(1000 * self.sr(_mel2ph, durations))
|
||||
table = self.stretch_embed(torch.arange(0, 1001, device=stretch.device))
|
||||
stretch_embed = torch.index_select(table, 0, stretch.view(-1).long()).view_as(condition)
|
||||
condition += stretch_embed
|
||||
stretch_embed_rnn_out, _ = self.stretch_embed_rnn(condition)
|
||||
condition += stretch_embed_rnn_out
|
||||
|
||||
if self.f0_embed_type == 'discrete':
|
||||
pitch = f0_to_coarse(f0)
|
||||
pitch_embed = self.pitch_embed(pitch)
|
||||
else:
|
||||
f0_mel = (1 + f0 / 700).log()
|
||||
pitch_embed = self.pitch_embed(f0_mel[:, :, None])
|
||||
condition += pitch_embed
|
||||
|
||||
if self.use_variance_embeds:
|
||||
variance_embeds = torch.stack([
|
||||
self.variance_embeds[v_name](variances[v_name][:, :, None] * self.variance_scaling_factor[v_name])
|
||||
for v_name in self.variance_embed_list
|
||||
], dim=-1).sum(-1)
|
||||
condition += variance_embeds
|
||||
|
||||
if hparams['use_key_shift_embed']:
|
||||
if hasattr(self, 'frozen_key_shift'):
|
||||
key_shift_embed = self.key_shift_embed(self.frozen_key_shift[:, None, None] * self.variance_scaling_factor['key_shift'])
|
||||
else:
|
||||
gender = torch.clip(gender, min=-1., max=1.)
|
||||
gender_mask = (gender < 0.).float()
|
||||
key_shift = gender * ((1. - gender_mask) * self.shift_max + gender_mask * abs(self.shift_min))
|
||||
key_shift_embed = self.key_shift_embed(key_shift[:, :, None] * self.variance_scaling_factor['key_shift'])
|
||||
condition += key_shift_embed
|
||||
|
||||
if hparams['use_speed_embed']:
|
||||
if velocity is not None:
|
||||
velocity = torch.clip(velocity, min=self.speed_min, max=self.speed_max)
|
||||
speed_embed = self.speed_embed(velocity[:, :, None] * self.variance_scaling_factor['speed'])
|
||||
else:
|
||||
speed_embed = self.speed_embed(torch.FloatTensor([1.]).to(condition.device)[:, None, None] * self.variance_scaling_factor['speed'])
|
||||
condition += speed_embed
|
||||
|
||||
if hparams['use_spk_id']:
|
||||
if hasattr(self, 'frozen_spk_embed'):
|
||||
condition += self.frozen_spk_embed
|
||||
else:
|
||||
condition += spk_embed
|
||||
return condition
|
||||
|
||||
|
||||
class FastSpeech2VarianceONNX(FastSpeech2Variance):
|
||||
def __init__(self, vocab_size, cross_lingual_token_idx=None):
|
||||
super().__init__(vocab_size=vocab_size)
|
||||
self.register_buffer(
|
||||
'cross_lingual_token_idx',
|
||||
torch.LongTensor(cross_lingual_token_idx),
|
||||
persistent=False
|
||||
)
|
||||
if len(cross_lingual_token_idx) == 0:
|
||||
self.use_lang_id = False
|
||||
self.lr = LengthRegulator()
|
||||
|
||||
def forward_encoder_word(self, tokens, word_div, word_dur, languages=None):
|
||||
txt_embed = self.txt_embed(tokens)
|
||||
ph2word = self.lr(word_div)
|
||||
onset = ph2word > F.pad(ph2word, [1, -1])
|
||||
onset_embed = self.onset_embed(onset.long())
|
||||
ph_word_dur = torch.gather(F.pad(word_dur, [1, 0]), 1, ph2word)
|
||||
word_dur_embed = self.word_dur_embed(ph_word_dur.float()[:, :, None])
|
||||
extra_embed = onset_embed + word_dur_embed
|
||||
if self.use_lang_id:
|
||||
lang_mask = torch.any(
|
||||
tokens[..., None] == self.cross_lingual_token_idx[None, None],
|
||||
dim=-1
|
||||
)
|
||||
lang_embed = self.lang_embed(languages * lang_mask)
|
||||
extra_embed += lang_embed
|
||||
x_masks = tokens == PAD_INDEX
|
||||
return self.encoder(txt_embed, extra_embed, x_masks), x_masks
|
||||
|
||||
def forward_encoder_phoneme(self, tokens, ph_dur, languages=None):
|
||||
txt_embed = self.txt_embed(tokens)
|
||||
if self.use_variance_scaling:
|
||||
ph_dur_embed = self.ph_dur_embed(torch.log(1 + ph_dur.float())[:, :, None])
|
||||
else:
|
||||
ph_dur_embed = self.ph_dur_embed(ph_dur.float()[:, :, None])
|
||||
if self.use_lang_id:
|
||||
lang_mask = torch.any(
|
||||
tokens[..., None] == self.cross_lingual_token_idx[None, None],
|
||||
dim=-1
|
||||
)
|
||||
lang_embed = self.lang_embed(languages * lang_mask)
|
||||
extra_embed = ph_dur_embed + lang_embed
|
||||
else:
|
||||
extra_embed = ph_dur_embed
|
||||
x_masks = tokens == PAD_INDEX
|
||||
return self.encoder(txt_embed, extra_embed, x_masks), x_masks
|
||||
|
||||
def forward_dur_predictor(self, encoder_out, x_masks, ph_midi, spk_embed=None):
|
||||
midi_embed = self.midi_embed(ph_midi)
|
||||
dur_cond = encoder_out + midi_embed
|
||||
if hparams['use_spk_id'] and spk_embed is not None:
|
||||
dur_cond += spk_embed
|
||||
ph_dur = self.dur_predictor(dur_cond, x_masks=x_masks)
|
||||
return ph_dur
|
||||
|
||||
def view_as_encoder(self):
|
||||
model = copy.deepcopy(self)
|
||||
if self.predict_dur:
|
||||
del model.dur_predictor
|
||||
model.forward = model.forward_encoder_word
|
||||
else:
|
||||
model.forward = model.forward_encoder_phoneme
|
||||
return model
|
||||
|
||||
def view_as_dur_predictor(self):
|
||||
model = copy.deepcopy(self)
|
||||
del model.encoder
|
||||
model.forward = model.forward_dur_predictor
|
||||
return model
|
||||
@@ -0,0 +1,16 @@
|
||||
import torch
|
||||
|
||||
from modules.nsf_hifigan.env import AttrDict
|
||||
from modules.nsf_hifigan.models import Generator
|
||||
|
||||
|
||||
# noinspection SpellCheckingInspection
|
||||
class NSFHiFiGANONNX(torch.nn.Module):
|
||||
def __init__(self, attrs: dict):
|
||||
super().__init__()
|
||||
self.generator = Generator(AttrDict(attrs))
|
||||
|
||||
def forward(self, mel: torch.Tensor, f0: torch.Tensor):
|
||||
mel = mel.transpose(1, 2)
|
||||
wav = self.generator(mel, f0)
|
||||
return wav.squeeze(1)
|
||||
@@ -0,0 +1,123 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import List, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
from modules.core import (
|
||||
RectifiedFlow, PitchRectifiedFlow, MultiVarianceRectifiedFlow
|
||||
)
|
||||
|
||||
|
||||
class RectifiedFlowONNX(RectifiedFlow):
|
||||
@property
|
||||
def backbone(self):
|
||||
return self.velocity_fn
|
||||
|
||||
# We give up the setter for the property `backbone` because this will cause TorchScript to fail
|
||||
# @backbone.setter
|
||||
@torch.jit.unused
|
||||
def set_backbone(self, value):
|
||||
self.velocity_fn = value
|
||||
|
||||
def sample_euler(self, x, t, dt: float, cond):
|
||||
x += self.velocity_fn(x, t * self.time_scale_factor, cond) * dt
|
||||
return x
|
||||
|
||||
def norm_spec(self, x):
|
||||
k = (self.spec_max - self.spec_min) / 2.
|
||||
b = (self.spec_max + self.spec_min) / 2.
|
||||
return (x - b) / k
|
||||
|
||||
def denorm_spec(self, x):
|
||||
k = (self.spec_max - self.spec_min) / 2.
|
||||
b = (self.spec_max + self.spec_min) / 2.
|
||||
return x * k + b
|
||||
|
||||
def forward(self, condition, x_end=None, depth=None, steps: int = 10):
|
||||
condition = condition.transpose(1, 2) # [1, T, H] => [1, H, T]
|
||||
device = condition.device
|
||||
n_frames = condition.shape[2]
|
||||
noise = torch.randn((1, self.num_feats, self.out_dims, n_frames), device=device)
|
||||
if x_end is None:
|
||||
t_start = 0.
|
||||
x = noise
|
||||
else:
|
||||
t_start = torch.max(1 - depth, torch.tensor(self.t_start, dtype=torch.float32, device=device))
|
||||
x_end = self.norm_spec(x_end).transpose(-2, -1)
|
||||
if self.num_feats == 1:
|
||||
x_end = x_end[:, None, :, :]
|
||||
if t_start <= 0.:
|
||||
x = noise
|
||||
elif t_start >= 1.:
|
||||
x = x_end
|
||||
else:
|
||||
x = t_start * x_end + (1 - t_start) * noise
|
||||
|
||||
t_width = 1. - t_start
|
||||
if t_width >= 0.:
|
||||
dt = t_width / max(1, steps)
|
||||
for t in torch.arange(steps, dtype=torch.long, device=device)[:, None].float() * dt + t_start:
|
||||
x = self.sample_euler(x, t, dt, condition)
|
||||
|
||||
if self.num_feats == 1:
|
||||
x = x.squeeze(1).permute(0, 2, 1) # [B, 1, M, T] => [B, T, M]
|
||||
else:
|
||||
x = x.permute(0, 1, 3, 2) # [B, F, M, T] => [B, F, T, M]
|
||||
x = self.denorm_spec(x)
|
||||
return x
|
||||
|
||||
|
||||
class PitchRectifiedFlowONNX(RectifiedFlowONNX, PitchRectifiedFlow):
|
||||
def __init__(self, vmin: float, vmax: float,
|
||||
cmin: float, cmax: float, repeat_bins,
|
||||
time_scale_factor=1000,
|
||||
backbone_type=None, backbone_args=None):
|
||||
self.vmin = vmin
|
||||
self.vmax = vmax
|
||||
self.cmin = cmin
|
||||
self.cmax = cmax
|
||||
super(PitchRectifiedFlow, self).__init__(
|
||||
vmin=vmin, vmax=vmax, repeat_bins=repeat_bins,
|
||||
time_scale_factor=time_scale_factor,
|
||||
backbone_type=backbone_type, backbone_args=backbone_args
|
||||
)
|
||||
|
||||
def clamp_spec(self, x):
|
||||
return x.clamp(min=self.cmin, max=self.cmax)
|
||||
|
||||
def denorm_spec(self, x):
|
||||
d = (self.spec_max - self.spec_min) / 2.
|
||||
m = (self.spec_max + self.spec_min) / 2.
|
||||
x = x * d + m
|
||||
x = x.mean(dim=-1)
|
||||
return x
|
||||
|
||||
|
||||
class MultiVarianceRectifiedFlowONNX(RectifiedFlowONNX, MultiVarianceRectifiedFlow):
|
||||
def __init__(
|
||||
self, ranges: List[Tuple[float, float]],
|
||||
clamps: List[Tuple[float | None, float | None] | None],
|
||||
repeat_bins, time_scale_factor=1000,
|
||||
backbone_type=None, backbone_args=None
|
||||
):
|
||||
assert len(ranges) == len(clamps)
|
||||
self.clamps = clamps
|
||||
vmin = [r[0] for r in ranges]
|
||||
vmax = [r[1] for r in ranges]
|
||||
if len(vmin) == 1:
|
||||
vmin = vmin[0]
|
||||
if len(vmax) == 1:
|
||||
vmax = vmax[0]
|
||||
super(MultiVarianceRectifiedFlow, self).__init__(
|
||||
vmin=vmin, vmax=vmax, repeat_bins=repeat_bins,
|
||||
time_scale_factor=time_scale_factor,
|
||||
backbone_type=backbone_type, backbone_args=backbone_args
|
||||
)
|
||||
|
||||
def denorm_spec(self, x):
|
||||
d = (self.spec_max - self.spec_min) / 2.
|
||||
m = (self.spec_max + self.spec_min) / 2.
|
||||
x = x * d + m
|
||||
x = x.mean(dim=-1)
|
||||
return x
|
||||
@@ -0,0 +1,413 @@
|
||||
import copy
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from torch import Tensor
|
||||
|
||||
from deployment.modules.diffusion import (
|
||||
GaussianDiffusionONNX, PitchDiffusionONNX, MultiVarianceDiffusionONNX
|
||||
)
|
||||
from deployment.modules.rectified_flow import (
|
||||
RectifiedFlowONNX, PitchRectifiedFlowONNX, MultiVarianceRectifiedFlowONNX
|
||||
)
|
||||
from deployment.modules.fastspeech2 import FastSpeech2AcousticONNX, FastSpeech2VarianceONNX
|
||||
from modules.toplevel import DiffSingerAcoustic, DiffSingerVariance
|
||||
from utils.hparams import hparams
|
||||
|
||||
|
||||
class DiffSingerAcousticONNX(DiffSingerAcoustic):
|
||||
def __init__(self, vocab_size, out_dims, cross_lingual_token_idx=None):
|
||||
super().__init__(vocab_size, out_dims)
|
||||
del self.fs2
|
||||
del self.diffusion
|
||||
self.fs2 = FastSpeech2AcousticONNX(
|
||||
vocab_size=vocab_size,
|
||||
cross_lingual_token_idx=cross_lingual_token_idx
|
||||
)
|
||||
if self.diffusion_type == 'ddpm':
|
||||
self.diffusion = GaussianDiffusionONNX(
|
||||
out_dims=out_dims,
|
||||
num_feats=1,
|
||||
timesteps=hparams['timesteps'],
|
||||
k_step=hparams['K_step'],
|
||||
backbone_type=self.backbone_type,
|
||||
backbone_args=self.backbone_args,
|
||||
spec_min=hparams['spec_min'],
|
||||
spec_max=hparams['spec_max']
|
||||
)
|
||||
elif self.diffusion_type == 'reflow':
|
||||
self.diffusion = RectifiedFlowONNX(
|
||||
out_dims=out_dims,
|
||||
num_feats=1,
|
||||
t_start=hparams['T_start'],
|
||||
time_scale_factor=hparams['time_scale_factor'],
|
||||
backbone_type=self.backbone_type,
|
||||
backbone_args=self.backbone_args,
|
||||
spec_min=hparams['spec_min'],
|
||||
spec_max=hparams['spec_max']
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Invalid diffusion type: {self.diffusion_type}")
|
||||
self.mel_base = hparams.get('mel_base', '10')
|
||||
|
||||
def ensure_mel_base(self, mel):
|
||||
if self.mel_base != 'e':
|
||||
# log10 mel to log mel
|
||||
mel = mel * 2.30259
|
||||
return mel
|
||||
|
||||
def forward_fs2_aux(
|
||||
self,
|
||||
tokens: Tensor,
|
||||
durations: Tensor,
|
||||
f0: Tensor,
|
||||
variances: dict,
|
||||
gender: Tensor = None,
|
||||
velocity: Tensor = None,
|
||||
spk_embed: Tensor = None,
|
||||
languages: Tensor = None
|
||||
):
|
||||
condition = self.fs2(
|
||||
tokens, durations, f0, variances=variances,
|
||||
gender=gender, velocity=velocity, spk_embed=spk_embed,
|
||||
languages=languages
|
||||
)
|
||||
if self.use_shallow_diffusion:
|
||||
aux_mel_pred = self.aux_decoder(condition, infer=True)
|
||||
return condition, aux_mel_pred
|
||||
else:
|
||||
return condition
|
||||
|
||||
def forward_shallow_diffusion(
|
||||
self, condition: Tensor, x_start: Tensor,
|
||||
depth, steps: int
|
||||
) -> Tensor:
|
||||
mel_pred = self.diffusion(condition, x_start=x_start, depth=depth, steps=steps)
|
||||
return self.ensure_mel_base(mel_pred)
|
||||
|
||||
def forward_diffusion(self, condition: Tensor, steps: int):
|
||||
mel_pred = self.diffusion(condition, steps=steps)
|
||||
return self.ensure_mel_base(mel_pred)
|
||||
|
||||
def forward_shallow_reflow(
|
||||
self, condition: Tensor, x_end: Tensor,
|
||||
depth, steps: int
|
||||
):
|
||||
mel_pred = self.diffusion(condition, x_end=x_end, depth=depth, steps=steps)
|
||||
return self.ensure_mel_base(mel_pred)
|
||||
|
||||
def forward_reflow(self, condition: Tensor, steps: int):
|
||||
mel_pred = self.diffusion(condition, steps=steps)
|
||||
return self.ensure_mel_base(mel_pred)
|
||||
|
||||
def view_as_fs2_aux(self) -> nn.Module:
|
||||
model = copy.deepcopy(self)
|
||||
del model.diffusion
|
||||
model.forward = model.forward_fs2_aux
|
||||
return model
|
||||
|
||||
def view_as_diffusion(self) -> nn.Module:
|
||||
model = copy.deepcopy(self)
|
||||
del model.fs2
|
||||
if self.use_shallow_diffusion:
|
||||
del model.aux_decoder
|
||||
model.forward = model.forward_shallow_diffusion
|
||||
else:
|
||||
model.forward = model.forward_diffusion
|
||||
return model
|
||||
|
||||
def view_as_reflow(self) -> nn.Module:
|
||||
model = copy.deepcopy(self)
|
||||
del model.fs2
|
||||
if self.use_shallow_diffusion:
|
||||
del model.aux_decoder
|
||||
model.forward = model.forward_shallow_reflow
|
||||
else:
|
||||
model.forward = model.forward_reflow
|
||||
return model
|
||||
|
||||
|
||||
class DiffSingerVarianceONNX(DiffSingerVariance):
|
||||
def __init__(self, vocab_size, cross_lingual_token_idx=None):
|
||||
super().__init__(vocab_size=vocab_size)
|
||||
del self.fs2
|
||||
self.fs2 = FastSpeech2VarianceONNX(
|
||||
vocab_size=vocab_size,
|
||||
cross_lingual_token_idx=cross_lingual_token_idx
|
||||
)
|
||||
self.hidden_size = hparams['hidden_size']
|
||||
if self.predict_pitch:
|
||||
del self.pitch_predictor
|
||||
self.smooth: nn.Conv1d = None
|
||||
pitch_hparams = hparams['pitch_prediction_args']
|
||||
if self.diffusion_type == 'ddpm':
|
||||
self.pitch_predictor = PitchDiffusionONNX(
|
||||
vmin=pitch_hparams['pitd_norm_min'],
|
||||
vmax=pitch_hparams['pitd_norm_max'],
|
||||
cmin=pitch_hparams['pitd_clip_min'],
|
||||
cmax=pitch_hparams['pitd_clip_max'],
|
||||
repeat_bins=pitch_hparams['repeat_bins'],
|
||||
timesteps=hparams['timesteps'],
|
||||
k_step=hparams['K_step'],
|
||||
backbone_type=self.pitch_backbone_type,
|
||||
backbone_args=self.pitch_backbone_args
|
||||
)
|
||||
elif self.diffusion_type == 'reflow':
|
||||
self.pitch_predictor = PitchRectifiedFlowONNX(
|
||||
vmin=pitch_hparams['pitd_norm_min'],
|
||||
vmax=pitch_hparams['pitd_norm_max'],
|
||||
cmin=pitch_hparams['pitd_clip_min'],
|
||||
cmax=pitch_hparams['pitd_clip_max'],
|
||||
repeat_bins=pitch_hparams['repeat_bins'],
|
||||
time_scale_factor=hparams['time_scale_factor'],
|
||||
backbone_type=self.pitch_backbone_type,
|
||||
backbone_args=self.pitch_backbone_args
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Invalid diffusion type: {self.diffusion_type}")
|
||||
if self.predict_variances:
|
||||
del self.variance_predictor
|
||||
if self.diffusion_type == 'ddpm':
|
||||
self.variance_predictor = self.build_adaptor(cls=MultiVarianceDiffusionONNX)
|
||||
elif self.diffusion_type == 'reflow':
|
||||
self.variance_predictor = self.build_adaptor(cls=MultiVarianceRectifiedFlowONNX)
|
||||
else:
|
||||
raise NotImplementedError(self.diffusion_type)
|
||||
|
||||
def build_smooth_op(self, device):
|
||||
smooth_kernel_size = round(hparams['midi_smooth_width'] * hparams['audio_sample_rate'] / hparams['hop_size'])
|
||||
smooth = nn.Conv1d(
|
||||
in_channels=1,
|
||||
out_channels=1,
|
||||
kernel_size=smooth_kernel_size,
|
||||
bias=False,
|
||||
padding='same',
|
||||
padding_mode='replicate'
|
||||
).eval()
|
||||
smooth_kernel = torch.sin(torch.from_numpy(
|
||||
np.linspace(0, 1, smooth_kernel_size).astype(np.float32) * np.pi
|
||||
))
|
||||
smooth_kernel /= smooth_kernel.sum()
|
||||
smooth.weight.data = smooth_kernel[None, None]
|
||||
self.smooth = smooth.to(device)
|
||||
|
||||
def embed_frozen_spk(self, encoder_out):
|
||||
if hparams['use_spk_id'] and hasattr(self, 'frozen_spk_embed'):
|
||||
encoder_out += self.frozen_spk_embed
|
||||
return encoder_out
|
||||
|
||||
def forward_linguistic_encoder_word(self, tokens, word_div, word_dur, languages=None):
|
||||
encoder_out, x_masks = self.fs2.forward_encoder_word(tokens, word_div, word_dur, languages=languages)
|
||||
encoder_out = self.embed_frozen_spk(encoder_out)
|
||||
return encoder_out, x_masks
|
||||
|
||||
def forward_linguistic_encoder_phoneme(self, tokens, ph_dur, languages=None):
|
||||
encoder_out, x_masks = self.fs2.forward_encoder_phoneme(tokens, ph_dur, languages=languages)
|
||||
encoder_out = self.embed_frozen_spk(encoder_out)
|
||||
return encoder_out, x_masks
|
||||
|
||||
def forward_dur_predictor(self, encoder_out, x_masks, ph_midi, spk_embed=None):
|
||||
return self.fs2.forward_dur_predictor(encoder_out, x_masks, ph_midi, spk_embed=spk_embed)
|
||||
|
||||
def forward_mel2x_gather(self, x_src, x_dur, x_dim=None, check_stretch_embed=False):
|
||||
mel2x = self.lr(x_dur)
|
||||
_mel2x = mel2x
|
||||
if x_dim is not None:
|
||||
x_src = F.pad(x_src, [0, 0, 1, 0])
|
||||
mel2x = mel2x[..., None].repeat([1, 1, x_dim])
|
||||
else:
|
||||
x_src = F.pad(x_src, [1, 0])
|
||||
x_cond = torch.gather(x_src, 1, mel2x)
|
||||
if self.use_stretch_embed and check_stretch_embed:
|
||||
stretch = torch.round(1000 * self.sr(_mel2x, x_dur))
|
||||
table = self.stretch_embed(torch.arange(0, 1001, device=stretch.device))
|
||||
stretch_embed = torch.index_select(table, 0, stretch.view(-1).long()).view_as(x_cond)
|
||||
x_cond += stretch_embed
|
||||
stretch_embed_rnn_out, _ = self.stretch_embed_rnn(x_cond)
|
||||
x_cond += stretch_embed_rnn_out
|
||||
return x_cond
|
||||
|
||||
def forward_pitch_preprocess(
|
||||
self, encoder_out, ph_dur,
|
||||
note_midi=None, note_rest=None, note_dur=None, note_glide=None,
|
||||
pitch=None, expr=None, retake=None, spk_embed=None
|
||||
):
|
||||
condition = self.forward_mel2x_gather(encoder_out, ph_dur, x_dim=self.hidden_size, check_stretch_embed=True)
|
||||
if self.use_melody_encoder:
|
||||
if self.melody_encoder.use_glide_embed and note_glide is None:
|
||||
note_glide = torch.LongTensor([[0]]).to(encoder_out.device)
|
||||
melody_encoder_out = self.melody_encoder(
|
||||
note_midi, note_rest, note_dur,
|
||||
glide=note_glide
|
||||
)
|
||||
melody_encoder_out = self.forward_mel2x_gather(melody_encoder_out, note_dur, x_dim=self.hidden_size)
|
||||
condition += melody_encoder_out
|
||||
if expr is None:
|
||||
retake_embed = self.pitch_retake_embed(retake.long())
|
||||
else:
|
||||
retake_true_embed = self.pitch_retake_embed(
|
||||
torch.ones(1, 1, dtype=torch.long, device=encoder_out.device)
|
||||
) # [B=1, T=1] => [B=1, T=1, H]
|
||||
retake_false_embed = self.pitch_retake_embed(
|
||||
torch.zeros(1, 1, dtype=torch.long, device=encoder_out.device)
|
||||
) # [B=1, T=1] => [B=1, T=1, H]
|
||||
expr = (expr * retake)[:, :, None] # [B, T, 1]
|
||||
retake_embed = expr * retake_true_embed + (1. - expr) * retake_false_embed
|
||||
pitch_cond = condition + retake_embed
|
||||
frame_midi_pitch = self.forward_mel2x_gather(note_midi, note_dur, x_dim=None)
|
||||
base_pitch = self.smooth(frame_midi_pitch)
|
||||
if self.use_melody_encoder:
|
||||
delta_pitch = (pitch - base_pitch) * ~retake
|
||||
if self.use_variance_scaling:
|
||||
pitch_cond += self.delta_pitch_embed(delta_pitch[:, :, None] / 12)
|
||||
else:
|
||||
pitch_cond += self.delta_pitch_embed(delta_pitch[:, :, None])
|
||||
else:
|
||||
base_pitch = base_pitch * retake + pitch * ~retake
|
||||
if self.use_variance_scaling:
|
||||
pitch_cond += self.base_pitch_embed(base_pitch[:, :, None] / 128)
|
||||
else:
|
||||
pitch_cond += self.base_pitch_embed(base_pitch[:, :, None])
|
||||
if hparams['use_spk_id'] and spk_embed is not None:
|
||||
pitch_cond += spk_embed
|
||||
return pitch_cond, base_pitch
|
||||
|
||||
def forward_pitch_reflow(
|
||||
self, pitch_cond, steps: int = 10
|
||||
):
|
||||
x_pred = self.pitch_predictor(pitch_cond, steps=steps)
|
||||
return x_pred
|
||||
|
||||
def forward_pitch_postprocess(self, x_pred, base_pitch):
|
||||
pitch_pred = self.pitch_predictor.clamp_spec(x_pred) + base_pitch
|
||||
return pitch_pred
|
||||
|
||||
def forward_variance_preprocess(
|
||||
self, encoder_out, ph_dur, pitch,
|
||||
variances: dict = None, retake=None, spk_embed=None
|
||||
):
|
||||
condition = self.forward_mel2x_gather(encoder_out, ph_dur, x_dim=self.hidden_size, check_stretch_embed=True)
|
||||
if self.use_variance_scaling:
|
||||
variance_cond = condition + self.pitch_embed(pitch[:, :, None] / 12)
|
||||
else:
|
||||
variance_cond = condition + self.pitch_embed(pitch[:, :, None])
|
||||
non_retake_masks = [
|
||||
v_retake.float() # [B, T, 1]
|
||||
for v_retake in (~retake).split(1, dim=2)
|
||||
]
|
||||
variance_embeds = [
|
||||
self.variance_embeds[v_name](variances[v_name][:, :, None] * self.variance_retake_scaling[v_name]) * v_masks
|
||||
for v_name, v_masks in zip(self.variance_prediction_list, non_retake_masks)
|
||||
]
|
||||
variance_cond += torch.stack(variance_embeds, dim=-1).sum(-1)
|
||||
if hparams['use_spk_id'] and spk_embed is not None:
|
||||
variance_cond += spk_embed
|
||||
return variance_cond
|
||||
|
||||
def forward_variance_reflow(self, variance_cond, steps: int = 10):
|
||||
xs_pred = self.variance_predictor(variance_cond, steps=steps)
|
||||
return xs_pred
|
||||
|
||||
def forward_variance_postprocess(self, xs_pred):
|
||||
if self.variance_predictor.num_feats == 1:
|
||||
xs_pred = [xs_pred]
|
||||
else:
|
||||
xs_pred = xs_pred.unbind(dim=1)
|
||||
variance_pred = self.variance_predictor.clamp_spec(xs_pred)
|
||||
return tuple(variance_pred)
|
||||
|
||||
def view_as_linguistic_encoder(self):
|
||||
model = copy.deepcopy(self)
|
||||
if self.predict_pitch:
|
||||
del model.pitch_predictor
|
||||
if self.use_melody_encoder:
|
||||
del model.melody_encoder
|
||||
if self.predict_variances:
|
||||
del model.variance_predictor
|
||||
model.fs2 = model.fs2.view_as_encoder()
|
||||
if self.predict_dur:
|
||||
model.forward = model.forward_linguistic_encoder_word
|
||||
else:
|
||||
model.forward = model.forward_linguistic_encoder_phoneme
|
||||
return model
|
||||
|
||||
def view_as_dur_predictor(self):
|
||||
assert self.predict_dur
|
||||
model = copy.deepcopy(self)
|
||||
if self.predict_pitch:
|
||||
del model.pitch_predictor
|
||||
if self.use_melody_encoder:
|
||||
del model.melody_encoder
|
||||
if self.predict_variances:
|
||||
del model.variance_predictor
|
||||
model.fs2 = model.fs2.view_as_dur_predictor()
|
||||
model.forward = model.forward_dur_predictor
|
||||
return model
|
||||
|
||||
def view_as_pitch_preprocess(self):
|
||||
model = copy.deepcopy(self)
|
||||
del model.fs2
|
||||
if self.predict_pitch:
|
||||
del model.pitch_predictor
|
||||
if self.predict_variances:
|
||||
del model.variance_predictor
|
||||
model.forward = model.forward_pitch_preprocess
|
||||
return model
|
||||
|
||||
def view_as_pitch_predictor(self):
|
||||
assert self.predict_pitch
|
||||
model = copy.deepcopy(self)
|
||||
del model.fs2
|
||||
del model.lr
|
||||
if self.use_melody_encoder:
|
||||
del model.melody_encoder
|
||||
if self.predict_variances:
|
||||
del model.variance_predictor
|
||||
model.forward = model.forward_pitch_reflow
|
||||
return model
|
||||
|
||||
def view_as_pitch_postprocess(self):
|
||||
model = copy.deepcopy(self)
|
||||
del model.fs2
|
||||
if self.use_melody_encoder:
|
||||
del model.melody_encoder
|
||||
if self.predict_variances:
|
||||
del model.variance_predictor
|
||||
model.forward = model.forward_pitch_postprocess
|
||||
return model
|
||||
|
||||
def view_as_variance_preprocess(self):
|
||||
model = copy.deepcopy(self)
|
||||
del model.fs2
|
||||
if self.predict_pitch:
|
||||
del model.pitch_predictor
|
||||
if self.use_melody_encoder:
|
||||
del model.melody_encoder
|
||||
if self.predict_variances:
|
||||
del model.variance_predictor
|
||||
model.forward = model.forward_variance_preprocess
|
||||
return model
|
||||
|
||||
def view_as_variance_predictor(self):
|
||||
assert self.predict_variances
|
||||
model = copy.deepcopy(self)
|
||||
del model.fs2
|
||||
del model.lr
|
||||
if self.predict_pitch:
|
||||
del model.pitch_predictor
|
||||
if self.use_melody_encoder:
|
||||
del model.melody_encoder
|
||||
model.forward = model.forward_variance_reflow
|
||||
return model
|
||||
|
||||
def view_as_variance_postprocess(self):
|
||||
model = copy.deepcopy(self)
|
||||
del model.fs2
|
||||
if self.predict_pitch:
|
||||
del model.pitch_predictor
|
||||
if self.use_melody_encoder:
|
||||
del model.melody_encoder
|
||||
model.forward = model.forward_variance_postprocess
|
||||
return model
|
||||
@@ -0,0 +1,3 @@
|
||||
*.py
|
||||
*.txt
|
||||
!opencpop*
|
||||
@@ -0,0 +1,601 @@
|
||||
a a
|
||||
ai ai
|
||||
an an
|
||||
ang ang
|
||||
ao ao
|
||||
ba b a
|
||||
bai b ai
|
||||
ban b an
|
||||
bang b ang
|
||||
bao b ao
|
||||
be b e
|
||||
bei b ei
|
||||
ben b en
|
||||
beng b eng
|
||||
ber b er
|
||||
bi b i
|
||||
bia b ia
|
||||
bian b ian
|
||||
biang b iang
|
||||
biao b iao
|
||||
bie b ie
|
||||
bin b in
|
||||
bing b ing
|
||||
biong b iong
|
||||
biu b iu
|
||||
bo b o
|
||||
bong b ong
|
||||
bou b ou
|
||||
bu b u
|
||||
bua b ua
|
||||
buai b uai
|
||||
buan b uan
|
||||
buang b uang
|
||||
bui b ui
|
||||
bun b un
|
||||
bv b v
|
||||
bve b ve
|
||||
ca c a
|
||||
cai c ai
|
||||
can c an
|
||||
cang c ang
|
||||
cao c ao
|
||||
ce c e
|
||||
cei c ei
|
||||
cen c en
|
||||
ceng c eng
|
||||
cer c er
|
||||
cha ch a
|
||||
chai ch ai
|
||||
chan ch an
|
||||
chang ch ang
|
||||
chao ch ao
|
||||
che ch e
|
||||
chei ch ei
|
||||
chen ch en
|
||||
cheng ch eng
|
||||
cher ch er
|
||||
chi ch ir
|
||||
chong ch ong
|
||||
chou ch ou
|
||||
chu ch u
|
||||
chua ch ua
|
||||
chuai ch uai
|
||||
chuan ch uan
|
||||
chuang ch uang
|
||||
chui ch ui
|
||||
chun ch un
|
||||
chuo ch uo
|
||||
chv ch v
|
||||
chyi ch i
|
||||
ci c i0
|
||||
cong c ong
|
||||
cou c ou
|
||||
cu c u
|
||||
cua c ua
|
||||
cuai c uai
|
||||
cuan c uan
|
||||
cuang c uang
|
||||
cui c ui
|
||||
cun c un
|
||||
cuo c uo
|
||||
cv c v
|
||||
cyi c i
|
||||
da d a
|
||||
dai d ai
|
||||
dan d an
|
||||
dang d ang
|
||||
dao d ao
|
||||
de d e
|
||||
dei d ei
|
||||
den d en
|
||||
deng d eng
|
||||
der d er
|
||||
di d i
|
||||
dia d ia
|
||||
dian d ian
|
||||
diang d iang
|
||||
diao d iao
|
||||
die d ie
|
||||
din d in
|
||||
ding d ing
|
||||
diong d iong
|
||||
diu d iu
|
||||
dong d ong
|
||||
dou d ou
|
||||
du d u
|
||||
dua d ua
|
||||
duai d uai
|
||||
duan d uan
|
||||
duang d uang
|
||||
dui d ui
|
||||
dun d un
|
||||
duo d uo
|
||||
dv d v
|
||||
dve d ve
|
||||
e e
|
||||
ei ei
|
||||
en en
|
||||
eng eng
|
||||
er er
|
||||
fa f a
|
||||
fai f ai
|
||||
fan f an
|
||||
fang f ang
|
||||
fao f ao
|
||||
fe f e
|
||||
fei f ei
|
||||
fen f en
|
||||
feng f eng
|
||||
fer f er
|
||||
fi f i
|
||||
fia f ia
|
||||
fian f ian
|
||||
fiang f iang
|
||||
fiao f iao
|
||||
fie f ie
|
||||
fin f in
|
||||
fing f ing
|
||||
fiong f iong
|
||||
fiu f iu
|
||||
fo f o
|
||||
fong f ong
|
||||
fou f ou
|
||||
fu f u
|
||||
fua f ua
|
||||
fuai f uai
|
||||
fuan f uan
|
||||
fuang f uang
|
||||
fui f ui
|
||||
fun f un
|
||||
fv f v
|
||||
fve f ve
|
||||
ga g a
|
||||
gai g ai
|
||||
gan g an
|
||||
gang g ang
|
||||
gao g ao
|
||||
ge g e
|
||||
gei g ei
|
||||
gen g en
|
||||
geng g eng
|
||||
ger g er
|
||||
gi g i
|
||||
gia g ia
|
||||
gian g ian
|
||||
giang g iang
|
||||
giao g iao
|
||||
gie g ie
|
||||
gin g in
|
||||
ging g ing
|
||||
giong g iong
|
||||
giu g iu
|
||||
gong g ong
|
||||
gou g ou
|
||||
gu g u
|
||||
gua g ua
|
||||
guai g uai
|
||||
guan g uan
|
||||
guang g uang
|
||||
gui g ui
|
||||
gun g un
|
||||
guo g uo
|
||||
gv g v
|
||||
gve g ve
|
||||
ha h a
|
||||
hai h ai
|
||||
han h an
|
||||
hang h ang
|
||||
hao h ao
|
||||
he h e
|
||||
hei h ei
|
||||
hen h en
|
||||
heng h eng
|
||||
her h er
|
||||
hi h i
|
||||
hia h ia
|
||||
hian h ian
|
||||
hiang h iang
|
||||
hiao h iao
|
||||
hie h ie
|
||||
hin h in
|
||||
hing h ing
|
||||
hiong h iong
|
||||
hiu h iu
|
||||
hong h ong
|
||||
hou h ou
|
||||
hu h u
|
||||
hua h ua
|
||||
huai h uai
|
||||
huan h uan
|
||||
huang h uang
|
||||
hui h ui
|
||||
hun h un
|
||||
huo h uo
|
||||
hv h v
|
||||
hve h ve
|
||||
ji j i
|
||||
jia j ia
|
||||
jian j ian
|
||||
jiang j iang
|
||||
jiao j iao
|
||||
jie j ie
|
||||
jin j in
|
||||
jing j ing
|
||||
jiong j iong
|
||||
jiu j iu
|
||||
ju j v
|
||||
juan j van
|
||||
jue j ve
|
||||
jun j vn
|
||||
ka k a
|
||||
kai k ai
|
||||
kan k an
|
||||
kang k ang
|
||||
kao k ao
|
||||
ke k e
|
||||
kei k ei
|
||||
ken k en
|
||||
keng k eng
|
||||
ker k er
|
||||
ki k i
|
||||
kia k ia
|
||||
kian k ian
|
||||
kiang k iang
|
||||
kiao k iao
|
||||
kie k ie
|
||||
kin k in
|
||||
king k ing
|
||||
kiong k iong
|
||||
kiu k iu
|
||||
kong k ong
|
||||
kou k ou
|
||||
ku k u
|
||||
kua k ua
|
||||
kuai k uai
|
||||
kuan k uan
|
||||
kuang k uang
|
||||
kui k ui
|
||||
kun k un
|
||||
kuo k uo
|
||||
kv k v
|
||||
kve k ve
|
||||
la l a
|
||||
lai l ai
|
||||
lan l an
|
||||
lang l ang
|
||||
lao l ao
|
||||
le l e
|
||||
lei l ei
|
||||
len l en
|
||||
leng l eng
|
||||
ler l er
|
||||
li l i
|
||||
lia l ia
|
||||
lian l ian
|
||||
liang l iang
|
||||
liao l iao
|
||||
lie l ie
|
||||
lin l in
|
||||
ling l ing
|
||||
liong l iong
|
||||
liu l iu
|
||||
lo l o
|
||||
long l ong
|
||||
lou l ou
|
||||
lu l u
|
||||
lua l ua
|
||||
luai l uai
|
||||
luan l uan
|
||||
luang l uang
|
||||
lui l ui
|
||||
lun l un
|
||||
luo l uo
|
||||
lv l v
|
||||
lve l ve
|
||||
ma m a
|
||||
mai m ai
|
||||
man m an
|
||||
mang m ang
|
||||
mao m ao
|
||||
me m e
|
||||
mei m ei
|
||||
men m en
|
||||
meng m eng
|
||||
mer m er
|
||||
mi m i
|
||||
mia m ia
|
||||
mian m ian
|
||||
miang m iang
|
||||
miao m iao
|
||||
mie m ie
|
||||
min m in
|
||||
ming m ing
|
||||
miong m iong
|
||||
miu m iu
|
||||
mo m o
|
||||
mong m ong
|
||||
mou m ou
|
||||
mu m u
|
||||
mua m ua
|
||||
muai m uai
|
||||
muan m uan
|
||||
muang m uang
|
||||
mui m ui
|
||||
mun m un
|
||||
mv m v
|
||||
mve m ve
|
||||
na n a
|
||||
nai n ai
|
||||
nan n an
|
||||
nang n ang
|
||||
nao n ao
|
||||
ne n e
|
||||
nei n ei
|
||||
nen n en
|
||||
neng n eng
|
||||
ner n er
|
||||
ni n i
|
||||
nia n ia
|
||||
nian n ian
|
||||
niang n iang
|
||||
niao n iao
|
||||
nie n ie
|
||||
nin n in
|
||||
ning n ing
|
||||
niong n iong
|
||||
niu n iu
|
||||
nong n ong
|
||||
nou n ou
|
||||
nu n u
|
||||
nua n ua
|
||||
nuai n uai
|
||||
nuan n uan
|
||||
nuang n uang
|
||||
nui n ui
|
||||
nun n un
|
||||
nuo n uo
|
||||
nv n v
|
||||
nve n ve
|
||||
o o
|
||||
ong ong
|
||||
ou ou
|
||||
pa p a
|
||||
pai p ai
|
||||
pan p an
|
||||
pang p ang
|
||||
pao p ao
|
||||
pe p e
|
||||
pei p ei
|
||||
pen p en
|
||||
peng p eng
|
||||
per p er
|
||||
pi p i
|
||||
pia p ia
|
||||
pian p ian
|
||||
piang p iang
|
||||
piao p iao
|
||||
pie p ie
|
||||
pin p in
|
||||
ping p ing
|
||||
piong p iong
|
||||
piu p iu
|
||||
po p o
|
||||
pong p ong
|
||||
pou p ou
|
||||
pu p u
|
||||
pua p ua
|
||||
puai p uai
|
||||
puan p uan
|
||||
puang p uang
|
||||
pui p ui
|
||||
pun p un
|
||||
pv p v
|
||||
pve p ve
|
||||
qi q i
|
||||
qia q ia
|
||||
qian q ian
|
||||
qiang q iang
|
||||
qiao q iao
|
||||
qie q ie
|
||||
qin q in
|
||||
qing q ing
|
||||
qiong q iong
|
||||
qiu q iu
|
||||
qu q v
|
||||
quan q van
|
||||
que q ve
|
||||
qun q vn
|
||||
ra r a
|
||||
rai r ai
|
||||
ran r an
|
||||
rang r ang
|
||||
rao r ao
|
||||
re r e
|
||||
rei r ei
|
||||
ren r en
|
||||
reng r eng
|
||||
rer r er
|
||||
ri r ir
|
||||
rong r ong
|
||||
rou r ou
|
||||
ru r u
|
||||
rua r ua
|
||||
ruai r uai
|
||||
ruan r uan
|
||||
ruang r uang
|
||||
rui r ui
|
||||
run r un
|
||||
ruo r uo
|
||||
rv r v
|
||||
ryi r i
|
||||
sa s a
|
||||
sai s ai
|
||||
san s an
|
||||
sang s ang
|
||||
sao s ao
|
||||
se s e
|
||||
sei s ei
|
||||
sen s en
|
||||
seng s eng
|
||||
ser s er
|
||||
sha sh a
|
||||
shai sh ai
|
||||
shan sh an
|
||||
shang sh ang
|
||||
shao sh ao
|
||||
she sh e
|
||||
shei sh ei
|
||||
shen sh en
|
||||
sheng sh eng
|
||||
sher sh er
|
||||
shi sh ir
|
||||
shong sh ong
|
||||
shou sh ou
|
||||
shu sh u
|
||||
shua sh ua
|
||||
shuai sh uai
|
||||
shuan sh uan
|
||||
shuang sh uang
|
||||
shui sh ui
|
||||
shun sh un
|
||||
shuo sh uo
|
||||
shv sh v
|
||||
shyi sh i
|
||||
si s i0
|
||||
song s ong
|
||||
sou s ou
|
||||
su s u
|
||||
sua s ua
|
||||
suai s uai
|
||||
suan s uan
|
||||
suang s uang
|
||||
sui s ui
|
||||
sun s un
|
||||
suo s uo
|
||||
sv s v
|
||||
syi s i
|
||||
ta t a
|
||||
tai t ai
|
||||
tan t an
|
||||
tang t ang
|
||||
tao t ao
|
||||
te t e
|
||||
tei t ei
|
||||
ten t en
|
||||
teng t eng
|
||||
ter t er
|
||||
ti t i
|
||||
tia t ia
|
||||
tian t ian
|
||||
tiang t iang
|
||||
tiao t iao
|
||||
tie t ie
|
||||
tin t in
|
||||
ting t ing
|
||||
tiong t iong
|
||||
tong t ong
|
||||
tou t ou
|
||||
tu t u
|
||||
tua t ua
|
||||
tuai t uai
|
||||
tuan t uan
|
||||
tuang t uang
|
||||
tui t ui
|
||||
tun t un
|
||||
tuo t uo
|
||||
tv t v
|
||||
tve t ve
|
||||
wa w a
|
||||
wai w ai
|
||||
wan w an
|
||||
wang w ang
|
||||
wao w ao
|
||||
we w e
|
||||
wei w ei
|
||||
wen w en
|
||||
weng w eng
|
||||
wer w er
|
||||
wi w i
|
||||
wo w o
|
||||
wong w ong
|
||||
wou w ou
|
||||
wu w u
|
||||
xi x i
|
||||
xia x ia
|
||||
xian x ian
|
||||
xiang x iang
|
||||
xiao x iao
|
||||
xie x ie
|
||||
xin x in
|
||||
xing x ing
|
||||
xiong x iong
|
||||
xiu x iu
|
||||
xu x v
|
||||
xuan x van
|
||||
xue x ve
|
||||
xun x vn
|
||||
ya y a
|
||||
yai y ai
|
||||
yan y En
|
||||
yang y ang
|
||||
yao y ao
|
||||
ye y E
|
||||
yei y ei
|
||||
yi y i
|
||||
yin y in
|
||||
ying y ing
|
||||
yo y o
|
||||
yong y ong
|
||||
you y ou
|
||||
yu y v
|
||||
yuan y van
|
||||
yue y ve
|
||||
yun y vn
|
||||
ywu y u
|
||||
za z a
|
||||
zai z ai
|
||||
zan z an
|
||||
zang z ang
|
||||
zao z ao
|
||||
ze z e
|
||||
zei z ei
|
||||
zen z en
|
||||
zeng z eng
|
||||
zer z er
|
||||
zha zh a
|
||||
zhai zh ai
|
||||
zhan zh an
|
||||
zhang zh ang
|
||||
zhao zh ao
|
||||
zhe zh e
|
||||
zhei zh ei
|
||||
zhen zh en
|
||||
zheng zh eng
|
||||
zher zh er
|
||||
zhi zh ir
|
||||
zhong zh ong
|
||||
zhou zh ou
|
||||
zhu zh u
|
||||
zhua zh ua
|
||||
zhuai zh uai
|
||||
zhuan zh uan
|
||||
zhuang zh uang
|
||||
zhui zh ui
|
||||
zhun zh un
|
||||
zhuo zh uo
|
||||
zhv zh v
|
||||
zhyi zh i
|
||||
zi z i0
|
||||
zong z ong
|
||||
zou z ou
|
||||
zu z u
|
||||
zua z ua
|
||||
zuai z uai
|
||||
zuan z uan
|
||||
zuang z uang
|
||||
zui z ui
|
||||
zun z un
|
||||
zuo z uo
|
||||
zv z v
|
||||
zyi z i
|
||||
@@ -0,0 +1,651 @@
|
||||
# Best Practices
|
||||
|
||||
## Fundamental concepts and materials
|
||||
|
||||
### Configuration files
|
||||
|
||||
A configuration file is a YAML file that defines enabled features, model hyperparameters and controls the behavior of the binarizer, trainer and inference. Almost all settings and controls in this repository, including the practices in this guidance, are achieved through configuration files.
|
||||
|
||||
For more information of the configuration system and configurable attributes, see [Configuration Schemas](ConfigurationSchemas.md).
|
||||
|
||||
### Languages
|
||||
|
||||
Each language you are dealing with should have a unique tag in the configuration file. **We highly recommend using ISO 639 language codes as language tags.** For example, `zh` and `zho` stands for Chinese (`cmn` specifically for Mandarin Chinese), `ja` and `jpn` for Japanese, `en` and `eng` for English, `yue` for Cantonese (Yue). You can download a complete language code table from https://iso639-3.sil.org/code_tables/download_tables.
|
||||
|
||||
### Phonemes
|
||||
|
||||
Phonemes are the fundamental part of dictionaries and labels. There are two types of phonemes: language-specific phonemes and global phonemes.
|
||||
|
||||
**Language-specific phonemes:** If there are multiple languages, all language-specific phonemes will be prefixed with its language name. For example: `zh/a`, `ja/o`, `en/eh`. These are called the **full name** of the phonemes, while `a`, `o`, `eh` are called the **short name** which has definite meaning only in a specific language context. If there is only one language, the short names can be used to determine each phoneme.
|
||||
|
||||
**Global phonemes:** Some phonemes do not belong to any language. There are two reserved global phoneme tags: `SP` for space, and `AP` for aspiration. There can also be other user-defined tags (`EP`, `GS`, `VF`, etc.). These tags will not be prefixed with language, and are prior when identifying phoneme names.
|
||||
|
||||
Extra phonemes, including user-defined global phonemes and additional language-specific phonemes that are not present in the dictionaries, can be defined in a list in the configuration file (full names should be used):
|
||||
|
||||
```yaml
|
||||
extra_phonemes: ['EP', 'ja/cl']
|
||||
```
|
||||
|
||||
The phoneme set expands rapidly with the number of languages. There are actually many similar phonemes that can be merged. Define the merging groups in your configuration file (full names should be used):
|
||||
|
||||
```yaml
|
||||
merged_phoneme_groups:
|
||||
- [zh/i, ja/i, en/iy]
|
||||
- [zh/s, ja/s, en/s]
|
||||
- [ja/cl, SP] # global phonemes can also be merged
|
||||
# ... (other groups omitted for brevity)
|
||||
use_lang_id: true # whether to use language embedding; only take effects if there are cross-lingual phonemes
|
||||
```
|
||||
|
||||
Merging phonemes does not mean that they are exactly the same for the dictionary. For those cross-lingual merged phonemes, Setting `use_lang_id` to true will still distinguish them by language IDs.
|
||||
|
||||
#### Phoneme naming principles
|
||||
|
||||
- Short names of language-specific phonemes should not conflict with global phoneme names, including reserved ones.
|
||||
- `/` cannot be used because it is already used for splitting the language tag and the short name.
|
||||
- `-` and `+` cannot be used because they are defined as slur tags in most singing voice synthesis editors.
|
||||
- Other special characters, including but not limited to `@`, `#`, `&`, `|`, `<`, `>`, is not recommended because they may be used as special tags in the future format changes.
|
||||
- ASCII characters are preferred for the best encoding compatibility, but all UTF-8 characters are acceptable.
|
||||
|
||||
### Dictionaries
|
||||
|
||||
Each language should have a corresponding dictionary. Define languages and dictionaries in your configuration file:
|
||||
|
||||
```yaml
|
||||
dictionaries:
|
||||
zh: dictionaries/opencpop-extension.txt
|
||||
ja: dictionaries/japanese_dict_full.txt
|
||||
en: dictionaries/ds_cmudict-07b.txt
|
||||
num_lang: 3 # number of languages; should be >= number of defined languages
|
||||
```
|
||||
|
||||
Each dictionary is a *.txt* file, in which each line represents a mapping rule from one syllable to its phoneme sequence. The syllable and the phonemes are split by `tab`, and the phonemes are split by `space`:
|
||||
|
||||
```
|
||||
<syllable> <phoneme1> <phoneme2> ...
|
||||
```
|
||||
|
||||
#### Syllable naming principles
|
||||
|
||||
- Try to use a standard writing or pronouncing system. For example, pinyin for Mandarin Chinese, romaji for Japanese and English words for English.
|
||||
- `AP` and `SP` cannot be used because they are reserved tags when using DiffSinger in editors.
|
||||
- `/` cannot be used because it is already used for splitting the language tag and the short name.
|
||||
- `-` and `+` cannot be used because they are defined as slur tags in most singing voice synthesis editors.
|
||||
- Syllable names is not recommended to start with `.` because this may have special meanings in the future editors.
|
||||
- Other special characters, including but not limited to `@`, `#`, `&`, `|`, `<`, `>`, is not recommended because they may be used as special tags in the future format changes.
|
||||
- ASCII characters are preferred for the best encoding compatibility, but all UTF-8 characters are acceptable.
|
||||
|
||||
There are some example dictionaries in the [dictionaries/](../dictionaries) folder.
|
||||
|
||||
### Datasets
|
||||
|
||||
A dataset mainly includes recordings and transcriptions, which is called a _raw dataset_. Raw datasets should be organized as the following folder structure:
|
||||
|
||||
- my_raw_data/
|
||||
- wavs/
|
||||
- 001.wav
|
||||
- 002.wav
|
||||
- ... (more recording files)
|
||||
- transcriptions.csv
|
||||
|
||||
In the example above, the _my_raw_data_ directory is the root directory of a raw dataset.
|
||||
|
||||
The _transcriptions.csv_ file contains all labels of the recordings. The common column of the CSV file is `name`, which represents all recording items by their filenames **without extension**. Elements of sequence attributes should be split by `space`. Other required columns may vary according to the category of the model you are training, and will be introduced in the following sections.
|
||||
|
||||
Each dataset should have a main language. If you have many recordings in multiple languages, it is recommended to separate them by language (you can merge their speaker IDs in the configuration). In each dataset, the main language is set as the language context, and phoneme labels in transcriptions.csv do not need a prefix (short name). It is also valid if there are phonemes from other languages, but all of them should be prefixed with their actual language (full name). Global phonemes should not be prefixed in any datasets.
|
||||
|
||||
You can define your datasets in the configuration file like this:
|
||||
|
||||
```yaml
|
||||
datasets: # define all raw datasets
|
||||
- raw_data_dir: data/spk1-zh/raw # path to the root of a raw dataset
|
||||
speaker: speaker1 # speaker name
|
||||
spk_id: 0 # optional; use this to merge two datasets; otherwise automatically assigned
|
||||
language: zh # language tag (main language) of this dataset
|
||||
test_prefixes: # optional; validation samples from this dataset
|
||||
- wav1
|
||||
- wav2
|
||||
- raw_data_dir: data/spk1-en/raw
|
||||
speaker: speaker1
|
||||
spk_id: 0 # specify the same speaker ID to merge into the previous one
|
||||
language: en
|
||||
test_prefixes:
|
||||
- wav1
|
||||
- wav2
|
||||
- raw_data_dir: data/spk2/raw
|
||||
speaker: speaker2
|
||||
language: ja
|
||||
test_prefixes:
|
||||
- wav1
|
||||
- wav2
|
||||
# ... (other datasets omitted for brevity)
|
||||
num_spk: 2 # number of languages; should be > maximum speaker ID
|
||||
```
|
||||
|
||||
### DS files
|
||||
|
||||
DS files are JSON files with _.ds_ suffix that contains phoneme sequence, phoneme durations, music scores or curve parameters. They are mainly used to run inference on models for test and evaluation purposes, and they can be used as training data in some cases. There are some example DS files in the [samples/](../samples) folder.
|
||||
|
||||
The current recommended way of using a model for production purposes is to use [OpenUTAU for DiffSinger](https://github.com/xunmengshe/OpenUtau). It can export DS files as well.
|
||||
|
||||
### Other fundamental assets
|
||||
|
||||
#### Vocoders
|
||||
|
||||
A vocoder is a model that can reconstruct the audio waveform given the low-dimensional mel-spectrogram. The vocoder is the essential dependency if you want to train an acoustic model and hear the voice on the TensorBoard.
|
||||
|
||||
The [DiffSinger Community Vocoders Project](https://openvpi.github.io/vocoders) provides a universal pre-trained NSF-HiFiGAN vocoder that can be used for starters of this repository. To use it, download the model (~50 MB size) from its releases and unzip it into the `checkpoints/` folder.
|
||||
|
||||
The pre-trained vocoder can be fine-tuned on your target dataset. It is highly recommended to do so because fine-tuned vocoder can generate much better results on specific (seen) datasets while does not need much computing resources. See the [vocoder training and fine-tuning repository](https://github.com/openvpi/SingingVocoders) for detailed instructions. After you get the fine-tuned vocoder checkpoint, you can configure it by `vocoder_ckpt` key in your configuration file. The fine-tuned NSF-HiFiGAN vocoder checkpoints can be exported to ONNX format like other DiffSinger user models for further production purposes.
|
||||
|
||||
Another unrecommended option: train an ultra-lightweight [DDSP vocoder](https://github.com/yxlllc/pc-ddsp) first by yourself, then configure it according to the relevant [instructions](https://github.com/yxlllc/pc-ddsp/blob/master/DiffSinger.md).
|
||||
|
||||
#### Feature extractors or auxiliary models
|
||||
|
||||
RMVPE is the recommended pitch extractor of this repository, which is an NN-based algorithm and requires a pre-trained model. For more information about pitch extractors and how to configure them, see [feature extraction](#pitch-extraction).
|
||||
|
||||
Vocal Remover (VR) is the recommended harmonic-noise separator of this repository, which is an NN-based algorithm and requires a pre-trained model. For more information about harmonic-noise separators and how to configure them, see [feature extraction](#harmonic-noise-separation).
|
||||
|
||||
## Overview: training acoustic models
|
||||
|
||||
An acoustic model takes low-level singing information as input, including (but not limited to) phoneme sequence, phoneme durations and F0 sequence. The only output of an acoustic model is the mel-spectrogram, which can be converted to waveform (the final audio) through the vocoder. Briefly speaking, an acoustic model takes in all features that are explicitly given, and produces the singing voice.
|
||||
|
||||
### Datasets
|
||||
|
||||
To train an acoustic model, you must have three columns in your transcriptions.csv: `name`, `ph_seq` and `ph_dur`, where `ph_seq` is the phoneme sequence and `ph_dur` is the phoneme duration sequence in seconds. You must have all corresponding recordings declared by the `name` column in mono, WAV format.
|
||||
|
||||
Training from multiple datasets in one model (so that the model is a multi-speaker model) is supported. See `speakers`, `spk_ids` and `use_spk_id` in the configuration schemas.
|
||||
|
||||
### Functionalities
|
||||
|
||||
Functionalities of acoustic models are defined by their inputs. Acoustic models have three basic and fixed inputs: phoneme sequence, phoneme duration sequence and F0 (pitch) sequence. There are three categories of additional inputs (control parameters):
|
||||
|
||||
- speaker IDs: if your acoustic model is a multi-speaker model, you can use different speaker in the same model, or mix their timbre and style.
|
||||
- variance parameters: these curve parameters are features extracted from the recordings, and can control the timbre and style of the singing voice. See `use_energy_embed` and `use_breathiness_embed` in the configuration schemas. Please note that variance parameters **do not have default values**, so they are usually obtained from the variance model at inference time.
|
||||
- transition parameters: these values represent the transition of the mel-spectrogram, and are obtained by enabling data augmentation. They are scalars at training time and sequences at inference time. See `augmentation_args`, `use_key_shift_embed` and `use_speed_embed` in the configuration schemas.
|
||||
|
||||
## Overview: training variance models
|
||||
|
||||
A variance model takes high-level music information as input, including phoneme sequence, word division, word durations and music scores. The outputs of a variance model may include phoneme durations, pitch curve and other control parameters that will be consumed by acoustic models. Briefly speaking, a variance model works as an auxiliary tool (so-called _automatic parameter generator_) for the acoustic models.
|
||||
|
||||
### Datasets
|
||||
|
||||
To train a variance model, you must have all the required attributes listed in the following table in your transcriptions.csv according to the functionalities enabled.
|
||||
|
||||
| | name | ph_seq | ph_dur | ph_num | note_seq | note_dur |
|
||||
|:------------------------------:|:----:|:------:|:------:|:------:|:--------:|:--------:|
|
||||
| phoneme duration prediction | ✓ | ✓ | ✓ | ✓ | | |
|
||||
| pitch prediction | ✓ | ✓ | ✓ | | ✓ | ✓ |
|
||||
| variance parameters prediction | ✓ | ✓ | ✓ | | | |
|
||||
|
||||
The recommended way of building a variance dataset is to extend an acoustic dataset. You may have all the recordings prepared like the acoustic dataset as well, or [use DS files in your variance datasets](#build-variance-datasets-with-ds-files).
|
||||
|
||||
Variance models support multi-speaker settings like acoustic models do.
|
||||
|
||||
### Functionalities
|
||||
|
||||
Functionalities of variance models are defined by their outputs. There are three main prediction modules that can be enabled/disable independently:
|
||||
|
||||
- Duration Predictor: predicts the phoneme durations. See `predict_dur` in the configuration schemas.
|
||||
- Pitch Predictor: predicts the pitch curve. See `predict_pitch` in the configuration schemas.
|
||||
- Multi-Variance Predictor: jointly predicts other variance parameters. See `predict_energy` and `predict_breathiness` in the configuration schemas.
|
||||
|
||||
There may be some mutual influence between the modules above when they are enabled together. See [mutual influence between variance modules](#mutual-influence-between-variance-modules) for more details.
|
||||
|
||||
## Build variance datasets with DS files
|
||||
|
||||
By default, the variance binarizer loads attributes from transcriptions.csv and searches for recording files (*.wav) to extract features and parameters. These attributes and parameters also exist in DS files, which are normally used for inference. This section introduces the required settings and important notes to build a variance dataset from DS files.
|
||||
|
||||
First of all, you should edit your configuration file to enable loading from DS files:
|
||||
|
||||
```yaml
|
||||
binarization_args:
|
||||
prefer_ds: true # prefer loading from DS files
|
||||
```
|
||||
|
||||
Then you should prepare some DS files which are properly segmented. If you export DS files with OpenUTAU for DiffSinger, the DS files are already segmented according to the spaces between notes. You should put these DS files in a folder named `ds` in your raw dataset directory (besides the `wavs` folder).
|
||||
|
||||
The DS files should also use the same dictionary as that of your target model. The attributes required vary from your target functionalities, as listed below:
|
||||
|
||||
| attribute name | required by duration prediction | required by pitch prediction | required by variance parameters prediction | previous source | current source |
|
||||
|:----------------------------:|:-------------------------------:|:----------------------------:|:------------------------------------------:|:---------------:|:--------------:|
|
||||
| `name` | ✓ | ✓ | ✓ | CSV | CSV |
|
||||
| `ph_seq` | ✓ | ✓ | ✓ | CSV | DS/CSV |
|
||||
| `ph_dur` | ✓ | ✓ | ✓ | CSV | DS/CSV |
|
||||
| `ph_num` | ✓ | | | CSV | DS/CSV |
|
||||
| `note_seq` | | ✓ | | CSV | DS/CSV |
|
||||
| `note_dur` | | ✓ | | CSV | DS/CSV |
|
||||
| `f0_seq` | ✓ | ✓ | ✓ | WAV | DS/WAV |
|
||||
| `energy`, `breathiness`, ... | | | ✓ | WAV | DS/WAV |
|
||||
|
||||
This means you only need one column in transcriptions.csv, the `name` column, to declare all DS files included in the dataset. The name pattern can be:
|
||||
|
||||
- Full name: `some-name` will firstly match the first segment in `some-name.ds`.
|
||||
- Name with index: `some-name#0` and `some-name#1` will match segment 0 and segment 1 in `some-name.ds` if there are no match with full name.
|
||||
|
||||
Though not recommended, the binarizer will still try to load attributes from transcriptions.csv or extract parameters from recordings if there are no matching DS files. In this case the full name matching logic is applied (the same as the normal binarization process).
|
||||
|
||||
## Choosing variance parameters
|
||||
|
||||
Variance parameters are a type of parameters that are significantly related to singing styles and emotions, have no default values and need to be predicted by the variance models. Choosing the proper variance parameters can obtain more controllability and expressiveness for your singing models. In this section, we are only talking about **narrowly defined variance parameters**, which are variance parameters except the pitch.
|
||||
|
||||
### Supported variance parameters
|
||||
|
||||
#### Energy
|
||||
|
||||
> WARNING
|
||||
>
|
||||
> This parameter is no longer recommended in favor of the new voicing parameter. The latter are less coupled with breathiness than energy.
|
||||
|
||||
Energy is defined as the RMS curve of the singing, in dB, which can control the strength of voice to a certain extent.
|
||||
|
||||
#### Breathiness
|
||||
|
||||
Breathiness is defined as the RMS curve of the aperiodic part of the singing, in dB, which can control the power of the air and unvoiced consonants in the voice.
|
||||
|
||||
#### Voicing
|
||||
|
||||
Voicing is defined as the RMS curve of the harmonic part of the singing, in dB, which can control the power of the harmonics in vowels and voiced consonants in the voice.
|
||||
|
||||
#### Tension
|
||||
|
||||
Tension is mostly related to the ratio of the base harmonic to the full harmonics, which can be used to control the strength and timbre of the voice. The ratio is calculated as
|
||||
$$
|
||||
r = \frac{\text{RMS}(H_{full}-H_{base})}{\text{RMS}(H_{full})}
|
||||
$$
|
||||
where $H_{full}$ is the full harmonics and $H_{base}$ is the base harmonic. The ratio is then mapped to the final domain via the inverse function of Sigmoid, that
|
||||
$$
|
||||
T = \log{\frac{r}{1-r}}
|
||||
$$
|
||||
where $T$ is the tension value.
|
||||
|
||||
### Principles of choosing multiple parameters
|
||||
|
||||
#### Energy, breathiness and voicing
|
||||
|
||||
These three parameters should **NOT** be enabled together. Energy is the RMS of the full waveform, which is the composition of the harmonic part and the aperiodic part. Therefore, these three parameters are coupled with each other.
|
||||
|
||||
#### Energy, voicing and tension
|
||||
|
||||
When voicing (or energy) is enabled, it almost fixes the loudness. However, tension sometimes rely on the implicitly predicted loudness for more expressiveness, because when a person sings with higher tension, he/she always produces louder voice. For this reason, some people may find their models or datasets _less natural_ with tension control. To be specific, changing tension will change the timbre but keep the loudness, and changing voicing (or energy) will change the loudness but keep the timbre. This behavior can be suitable for some, but not all datasets and users. Therefore, it is highly recommended for everyone to conduct some experiments on the actual datasets used to train the model.
|
||||
|
||||
## Mutual influence between variance modules
|
||||
|
||||
In some recent experiments and researches, some mutual influence between the modules of variance models has been found. In practice, being aware of the influence and making use of it can improve accuracy and avoid instability of the model.
|
||||
|
||||
### Influence on the duration predictor
|
||||
|
||||
The duration predictor benefits from its downstream modules, like the pitch predictor and the variance predictor.
|
||||
|
||||
The experiments were conducted on both manually refined datasets and automatically labeled datasets, and with pitch predictors driven by both base pitch and melody encoder. All the results have shown that when either of the pitch predictor and the variance predictor is enabled together with the duration predictor, its rhythm correctness and duration accuracy significantly outperforms those of a solely trained duration predictor.
|
||||
|
||||
Possible reason for this difference can be the lack of information carried by pure phoneme duration sequences, which may not fully represent the phoneme features in the real world. With the help of frame-level feature predictors, the encoder learns more knowledge about the voice features related to the phoneme types and durations, thus making the duration predictor produce better results.
|
||||
|
||||
### Influence on frame-level feature predictors
|
||||
|
||||
Frame-level feature predictors, including the pitch predictor and the variance predictor, have better performance when trained without enabling the duration predictor.
|
||||
|
||||
The experiments found that when the duration predictor is enabled, the pitch accuracy drops and the dynamics of variance parameters sometimes become unstable. And it has nothing to do with the gradients from the duration predictor, because applying a scale factor on the gradients does not make any difference even if the gradients are completely cut off.
|
||||
|
||||
Possible reason for this phenomenon can be the lack of direct phoneme duration input. When the duration predictor is enabled, the model takes in word durations instead of phoneme durations; when there is no duration predictor together, the phoneme duration sequence is directly taken in and passed through the attention-based linguistic encoder. With direct modeling on the phoneme duration, the frame-level predictors can have a better understanding of the context, thus producing better results.
|
||||
|
||||
Another set of experiments showed that there is no significant influence between the pitch predictor and the variance predictor. When they are enabled together without the duration predictor, both can converge well and produce satisfactory results. No conclusion can be drawn on this issue, and it can depend on the dataset.
|
||||
|
||||
### Suggested procedures of training variance models
|
||||
|
||||
According to the experiment results and the analysis above, the suggested procedures of training a set of variance models are listed below:
|
||||
|
||||
1. Train the duration predictor together with the variance predictor, and discard the variance predictor part.
|
||||
2. Train the pitch predictor and the variance predictor separately or together.
|
||||
3. If interested, compare across different combinations in step 2 and choose the best.
|
||||
|
||||
## Feature extraction
|
||||
|
||||
Feature extraction is the process of extracting low-level features from the recordings, which are needed as inputs for the acoustic models, or as outputs for the variance models.
|
||||
|
||||
### Pitch extraction
|
||||
|
||||
A pitch extractor estimates pitch (F0 sequence) from given recordings. F0 (fundamental frequency) is one of the most important components of singing voice that is needed by both acoustic models and variance models.
|
||||
|
||||
```yaml
|
||||
pe: parselmouth # pitch extractor type
|
||||
pe_ckpt: checkpoints/xxx/model.pt # pitch extractor model path (if it requires any)
|
||||
```
|
||||
|
||||
#### Parselmouth
|
||||
|
||||
[Parselmouth](https://github.com/YannickJadoul/Parselmouth) is the default pitch extractor in this repository. It is based on DSP algorithms, runs fast on CPU and can get accurate F0 on clean and normal recordings.
|
||||
|
||||
To use parselmouth, simply include the following line in your configuration file:
|
||||
|
||||
```yaml
|
||||
pe: parselmouth
|
||||
```
|
||||
|
||||
#### RMVPE (recommended)
|
||||
|
||||
[RMVPE](https://github.com/Dream-High/RMVPE) (Robust Model for Vocal Pitch Estimation) is the state-of-the-art NN-based pitch estimation model for singing voice. It runs slower than parselmouth, consumes more memory, however uses CUDA to accelerate computation (if available) and produce better results on noisy recordings and edge cases.
|
||||
|
||||
To enable RMVPE, download its pre-trained checkpoint from [here](https://github.com/yxlllc/RMVPE/releases), extract it into the `checkpoints/` folder and edit the configuration file:
|
||||
|
||||
```yaml
|
||||
pe: rmvpe
|
||||
pe_ckpt: checkpoints/rmvpe/model.pt
|
||||
```
|
||||
|
||||
#### Harvest
|
||||
|
||||
Harvest (Harvest: A high-performance fundamental frequency estimator from speech signals) is the recommended pitch extractor from Masanori Morise's [WORLD](https://github.com/mmorise/World), a free software for high-quality speech analysis, manipulation and synthesis. It is a state-of-the-art algorithmic pitch estimator designed for speech, but has seen use in singing voice synthesis. It runs the slowest compared to the others, but provides very accurate F0 on clean and normal recordings compared to parselmouth.
|
||||
|
||||
To use Harvest, simply include the following line in your configuration file:
|
||||
|
||||
```yaml
|
||||
pe: harvest
|
||||
```
|
||||
|
||||
**Note:** It is also recommended to change the F0 detection range for Harvest with accordance to your dataset, as they are hard boundaries for this algorithm and the defaults might not suffice for most use cases. To change the F0 detection range, you may include or edit this part in the configuration file:
|
||||
|
||||
```yaml
|
||||
f0_min: 65 # Minimum F0 to detect
|
||||
f0_max: 800 # Maximum F0 to detect
|
||||
```
|
||||
|
||||
### Harmonic-noise separation
|
||||
|
||||
Harmonic-noise separation is the process of separating the harmonic part and the aperiodic part of the singing voice. These parts are the fundamental components for variance parameters including breathiness, voicing and tension to be calculated from.
|
||||
|
||||
#### WORLD
|
||||
|
||||
This algorithm uses Masanori Morise's [WORLD](https://github.com/mmorise/World), a free software for high-quality speech analysis, manipulation and synthesis. It uses CPU (no CUDA required) but runs relatively slow.
|
||||
|
||||
To use WORLD, simply include the following line in your configuration file:
|
||||
|
||||
```yaml
|
||||
hnsep: world
|
||||
```
|
||||
|
||||
#### Vocal Remover (recommended)
|
||||
|
||||
Vocal Remover (VR) is originally a popular NN-based algorithm for music source separation that removes the vocal part from the music. This repository uses a specially trained model for harmonic-noise separation. VR extracts much cleaner harmonic parts, utilizes CUDA to accelerate computation (if available) and runs much faster than WORLD. However, it consumes more memory and should not be used with too many parallel workers.
|
||||
|
||||
To enable VR, download its pre-trained checkpoint from [here](https://github.com/yxlllc/vocal-remover/releases), extract it into the `checkpoints/` folder and edit the configuration file:
|
||||
|
||||
```yaml
|
||||
hnsep: vr
|
||||
hnsep_ckpt: checkpoints/vr/model.pt
|
||||
```
|
||||
|
||||
## Shallow diffusion
|
||||
|
||||
Shallow diffusion is a mechanism that can improve quality and save inference time for diffusion models that was first introduced in the original DiffSinger [paper](https://arxiv.org/abs/2105.02446). Instead of starting the diffusion process from purely gaussian noise as classic diffusion does, shallow diffusion adds a shallow gaussian noise on a low-quality results generated by a simple network (which is called the auxiliary decoder) to skip many unnecessary steps from the beginning. With the combination of shallow diffusion and sampling acceleration algorithms, we can get better results under the same inference speed as before, or achieve higher inference speed without quality deterioration.
|
||||
|
||||
Currently, acoustic models in this repository support shallow diffusion. The main switch of shallow diffusion is `use_shallow_diffusion` in the configuration file, and most arguments of shallow diffusion can be adjusted under `shallow_diffusion_args`. See [Configuration Schemas](ConfigurationSchemas.md) for more details.
|
||||
|
||||
### Train full shallow diffusion models from scratch
|
||||
|
||||
To train a full shallow diffusion model from scratch, simply introduce the following settings in your configuration file:
|
||||
|
||||
```yaml
|
||||
use_shallow_diffusion: true
|
||||
K_step: 400 # adjust according to your needs
|
||||
K_step_infer: 400 # should be <= K_step
|
||||
```
|
||||
|
||||
Please note that when shallow diffusion is enabled, only the last $K$ diffusion steps will be trained. Unlike classic diffusion models which are trained on full steps, the limit of `K_step` can make the training more efficient. However, `K_step` should not be set too small because without enough diffusion depth (steps), the low-quality auxiliary decoder results cannot be well refined. 200 ~ 400 should be the proper range of `K_step`.
|
||||
|
||||
The auxiliary decoder and the diffusion decoder shares the same linguistic encoder, which receives gradients from both the decoders. In some experiments, it was found that gradients from the auxiliary decoder will cause mismatching between the encoder and the diffusion decoder, resulting in the latter being unable to produce reasonable results. To prevent this case, a configuration item called `aux_decoder_grad` is introduced to apply a scale factor on the gradients from the auxiliary decoder during training. To adjust this factor, introduce the following in the configuration file:
|
||||
|
||||
```yaml
|
||||
shallow_diffusion_args:
|
||||
aux_decoder_grad: 0.1 # should not be too high
|
||||
```
|
||||
|
||||
### Train auxiliary decoder and diffusion decoder separately
|
||||
|
||||
Training a full shallow diffusion model can consume more memory because the auxiliary decoder is also in the training graph. In limited situations, the two decoders can be trained separately, i.e. train one decoder after another.
|
||||
|
||||
**STEP 1: train the diffusion decoder**
|
||||
|
||||
In the first stage, the linguistic encoder and the diffusion decoder is trained together, while the auxiliary decoder is left unchanged. Edit your configuration file like this:
|
||||
|
||||
```yaml
|
||||
use_shallow_diffusion: true # make sure the main option is turned on
|
||||
shallow_diffusion_args:
|
||||
train_aux_decoder: false # exclude the auxiliary decoder from the training graph
|
||||
train_diffusion: true # train diffusion decoder as normal
|
||||
val_gt_start: true # should be true because the auxiliary decoder is not trained yet
|
||||
```
|
||||
|
||||
Start training until `max_updates` is reached, or until you get satisfactory results on the TensorBoard.
|
||||
|
||||
**STEP 2: train the auxiliary decoder**
|
||||
|
||||
In the second stage, the auxiliary decoder is trained besides the linguistic encoder and the diffusion decoder. Edit your configuration file like this:
|
||||
|
||||
```yaml
|
||||
shallow_diffusion_args:
|
||||
train_aux_decoder: true
|
||||
train_diffusion: false # exclude the diffusion decoder from the training graph
|
||||
lambda_aux_mel_loss: 1.0 # no more need to limit the auxiliary loss
|
||||
```
|
||||
|
||||
Then you should freeze the encoder to prevent it from getting updates. This is because if the encoder changes, it no longer matches with the diffusion decoder, thus making the latter unable to produce correct results again. Edit your configuration file:
|
||||
|
||||
```yaml
|
||||
freezing_enabled: true
|
||||
frozen_params:
|
||||
- model.fs2 # the linguistic encoder
|
||||
```
|
||||
|
||||
You should also manually reset your learning rate scheduler because this is a new training process for the auxiliary decoder. Possible ways are:
|
||||
|
||||
1. Rename the latest checkpoint to `model_ckpt_steps_0.ckpt` and remove the other checkpoints from the directory.
|
||||
2. Increase the initial learning rate (if you use a scheduler that decreases the LR over training steps) so that the auxiliary decoder gets proper learning rate.
|
||||
|
||||
Additionally, `max_updates` should be adjusted to ensure enough training steps for the auxiliary decoder.
|
||||
|
||||
Once you finished the configurations above, you can resume the training. The auxiliary decoder normally does not need many steps to train, and you can stop training when you get stable results on the TensorBoard. Because this step is much more complicated than the previous step, it is recommended to run some inference to verify if the model is trained properly after everything is finished.
|
||||
|
||||
### Add shallow diffusion to classic diffusion models
|
||||
|
||||
Actually, all classic DDPMs have the ability to be "shallow". If you want to add shallow diffusion functionality to a former classic diffusion model, the only thing you need to do is to train an auxiliary decoder for it.
|
||||
|
||||
Before you start, you should edit the configuration file to ensure that you use the same datasets, and that you do not remove or add any of the functionalities of the old model. Then you can configure the old checkpoint in your configuration file:
|
||||
|
||||
```yaml
|
||||
finetune_enabled: true
|
||||
finetune_ckpt_path: xxx.ckpt # path to your old checkpoint
|
||||
finetune_ignored_params: [] # do not ignore any parameters
|
||||
```
|
||||
|
||||
Then you can follow the instructions in STEP 2 of the [previous section](#add-shallow-diffusion-to-classic-diffusion-models) to finish your training.
|
||||
|
||||
## Performance tuning
|
||||
|
||||
This section is about accelerating training and utilizing hardware.
|
||||
|
||||
### Data loader and batch sampler
|
||||
|
||||
The data loader loads data pieces from the binary dataset, and the batch sampler forms batches according to data lengths.
|
||||
|
||||
To configure the data loader, edit your configuration file:
|
||||
|
||||
```yaml
|
||||
ds_workers: 4 # number of DataLoader workers
|
||||
dataloader_prefetch_factor: 2 # load data in advance
|
||||
```
|
||||
|
||||
To configure the batch sampler, edit your configuration file:
|
||||
|
||||
```yaml
|
||||
sampler_frame_count_grid: 6 # lower value means higher speed but less randomness
|
||||
```
|
||||
|
||||
For more details of the batch sampler algorithm and this configuration key, see [sampler_frame_count_grid](ConfigurationSchemas.md#sampler_frame_count_grid).
|
||||
|
||||
### Automatic mixed precision
|
||||
|
||||
Enabling automatic mixed precision (AMP) can accelerate training and save GPU memory. DiffSinger have adapted the latest version of PyTorch Lightning for AMP functionalities.
|
||||
|
||||
By default, the training runs in FP32 precision. To enable AMP, edit your configuration file:
|
||||
|
||||
```yaml
|
||||
pl_trainer_precision: 16-mixed # FP16 precision
|
||||
```
|
||||
|
||||
or
|
||||
|
||||
```yaml
|
||||
pl_trainer_precision: bf16-mixed # BF16 precision
|
||||
```
|
||||
|
||||
For more precision options, please check out the [official documentation](https://lightning.ai/docs/pytorch/stable/common/trainer.html#precision).
|
||||
|
||||
### Training on multiple GPUs
|
||||
|
||||
Using distributed data parallel (DDP) can divide training tasks to multiple GPUs and synchronize gradients and weights between them. DiffSinger have adapted the latest version of PyTorch Lightning for DDP functionalities.
|
||||
|
||||
By default, the trainer will utilize all CUDA devices defined in the `CUDA_VISIBLE_DEVICES` environment variable (empty means using all available devices). If you want to specify which GPUs to use, edit your configuration file:
|
||||
|
||||
```yaml
|
||||
pl_trainer_devices: [0, 1, 2, 3] # use the first 4 GPUs defined in CUDA_VISIBLE_DEVICES
|
||||
```
|
||||
|
||||
Please note that `max_batch_size` and `max_batch_frames` are values for **each** GPU.
|
||||
|
||||
By default, the trainer uses NCCL as the DDP backend. If this gets stuck on your machine, try disabling P2P first via
|
||||
|
||||
```yaml
|
||||
nccl_p2p: false # disable P2P in NCCL
|
||||
```
|
||||
|
||||
Or if your machine does not support NCCL, you can switch to Gloo instead:
|
||||
|
||||
```yaml
|
||||
pl_trainer_strategy:
|
||||
name: ddp # must manually choose a strategy instead of 'auto'
|
||||
process_group_backend: gloo # however, it has a lower performance than NCCL
|
||||
```
|
||||
|
||||
### Gradient accumulation
|
||||
|
||||
Gradient accumulation means accumulating losses for several batches before each time the weights are updated. This can simulate a larger batch size with a lower GPU memory cost.
|
||||
|
||||
By default, the trainer calls `backward()` each time the losses are calculated through one batch of data. To enable gradient accumulation, edit your configuration file:
|
||||
|
||||
```yaml
|
||||
accumulate_grad_batches: 4 # the actual batch size will be 4x.
|
||||
```
|
||||
|
||||
Please note that enabling gradient accumulation will slow down training because the losses must be calculated for several times before the weights are updated (1 update to the weights = 1 actual training step).
|
||||
|
||||
## Optimizers and learning rate schedulers
|
||||
|
||||
The optimizer and the learning rate scheduler can take an important role in the training process. DiffSinger uses a flexible configuration logic for these two modules.
|
||||
|
||||
### Basic configurations
|
||||
|
||||
The optimizer and learning rate scheduler used during training can be configured by their full class name and keyword arguments in the configuration file. Take the following as an example for the optimizer:
|
||||
|
||||
```yaml
|
||||
optimizer_args:
|
||||
optimizer_cls: torch.optim.AdamW # class name of optimizer
|
||||
lr: 0.0004
|
||||
beta1: 0.9
|
||||
beta2: 0.98
|
||||
weight_decay: 0
|
||||
```
|
||||
|
||||
and for the learning rate scheduler:
|
||||
|
||||
```yaml
|
||||
lr_scheduler_args:
|
||||
scheduler_cls: torch.optim.lr_scheduler.StepLR # class name of learning rate schedule
|
||||
warmup_steps: 2000
|
||||
step_size: 50000
|
||||
gamma: 0.5
|
||||
```
|
||||
|
||||
Note that `optimizer_args` and `lr_scheduler_args` will be filtered by needed parameters and passed to `__init__` as keyword arguments (`kwargs`) when constructing the optimizer and scheduler. Therefore, you could specify all arguments according to your need in the configuration file to directly control the behavior of optimization and LR scheduling. It will also tolerate parameters existing in the configuration but not needed in `__init__`.
|
||||
|
||||
Also, note that the LR scheduler performs scheduling on the granularity of steps, not epochs.
|
||||
|
||||
The special case applies when a tuple is needed in `__init__`: `beta1` and `beta2` are treated separately and form a tuple in the code. You could try to pass in an array instead. (And as an experiment, AdamW does accept `[beta1, beta2]`). If there is another special treatment required, please submit an issue.
|
||||
|
||||
For PyTorch built-in optimizers and LR schedulers, see official [documentation](https://pytorch.org/docs/stable/optim.html) of the `torch.optim` package. If you found other optimizer and learning rate scheduler useful, you can raise a topic in [Discussions](https://github.com/openvpi/DiffSinger/discussions), raise [Issues](https://github.com/openvpi/DiffSinger/issues) or submit [PRs](https://github.com/openvpi/DiffSinger/pulls) if it introduces new codes or dependencies.
|
||||
|
||||
### Composite LR schedulers
|
||||
|
||||
Some LR schedulers like `SequentialLR` and `ChainedScheduler` may use other schedulers as arguments. Besides built-in types, there is a special design to configure these scheduler objects. See the following example.
|
||||
|
||||
```yaml
|
||||
lr_scheduler_args:
|
||||
scheduler_cls: torch.optim.lr_scheduler.SequentialLR
|
||||
schedulers:
|
||||
- cls: torch.optim.lr_scheduler.ExponentialLR
|
||||
gamma: 0.5
|
||||
- cls: torch.optim.lr_scheduler.LinearLR
|
||||
- cls: torch.optim.lr_scheduler.MultiStepLR
|
||||
milestones:
|
||||
- 10
|
||||
- 20
|
||||
milestones:
|
||||
- 10
|
||||
- 20
|
||||
```
|
||||
|
||||
The LR scheduler objects will be recursively construct objects if `cls` is present in sub-arguments. Please note that `cls` must be a scheduler class because this is a special design.
|
||||
|
||||
**WARNING:** Nested `SequentialLR` and `ChainedScheduler` have unexpected behavior. **DO NOT** nest them. Also, make sure the scheduler is _chainable_ before using it in `ChainedScheduler`.
|
||||
|
||||
## Fine-tuning and parameter freezing
|
||||
|
||||
### Fine-tuning from existing checkpoints
|
||||
|
||||
By default, the training starts from a model from scratch with randomly initialized parameters. However, if you already have some pre-trained checkpoints, and you need to adapt them to other datasets with their functionalities unchanged, fine-tuning may save training steps and time. In general, you need to add the following structure into the configuration file:
|
||||
|
||||
```yaml
|
||||
# take acoustic models as an example
|
||||
finetune_enabled: true # the main switch to enable fine-tuning
|
||||
finetune_ckpt_path: checkpoints/pretrained/model_ckpt_steps_320000.ckpt # path to your pre-trained checkpoint
|
||||
finetune_ignored_params: # prefix rules to exclude specific parameters when loading the checkpoints
|
||||
- model.fs2.encoder.embed_tokens # in case when the phoneme set is changed
|
||||
- model.fs2.txt_embed # same as above
|
||||
- model.fs2.spk_embed # in case when the speaker set is changed
|
||||
finetune_strict_shapes: true # whether to raise an error when parameter shapes mismatch
|
||||
```
|
||||
|
||||
For the pre-trained checkpoint, it must be a file saved with `torch.save`, containing a `dict` object and a `state_dict` key, like the following example:
|
||||
|
||||
```json5
|
||||
{
|
||||
"state_dict": {
|
||||
"model.fs2.txt_embed": null, // torch.Tensor
|
||||
"model.fs2.pitch_embed.weight": null, // torch.Tensor
|
||||
"model.fs2.pitch_embed.bias": null, // torch.Tensor
|
||||
// ... (other parameters)
|
||||
}
|
||||
// ... (other possible keys
|
||||
}
|
||||
```
|
||||
|
||||
**IMPORTANT NOTES**:
|
||||
|
||||
- The pre-trained checkpoint is **loaded only once** at the beginning of the training experiment. You may interrupt the training at any time, but after this new experiment has saved its own checkpoint, the pre-trained checkpoint will not be loaded again when the training is resumed.
|
||||
- Only the state dict of the checkpoint will be loaded. The optimizer state in the pre-trained checkpoint will be ignored.
|
||||
- The parameter name matching is **not strict** when loading the pre-trained checkpoint. This means that missing parameters in the state dict will still be left as randomly initialized, and redundant parameters will be ignored without any warnings and errors. There are cases where the tensor shapes mismatch between the pre-trained state dict and the model - edit `finetune_strict_shapes` to change the behavior when dealing with this.
|
||||
- Be careful if you want to change the functionalities when fine-tuning. Starting from a checkpoint trained under different functionalities may be even slower than training from scratch.
|
||||
|
||||
### Freezing model parameters
|
||||
|
||||
Sometimes you want to freeze part of the model during training or fine-tuning to save GPU memory, accelerate the training process or avoid catastrophic forgetting. Parameter freezing may also be useful if you want to add/remove functionalities from pre-trained checkpoints. In general, you need to add the following structure into the configuration file:
|
||||
|
||||
```yaml
|
||||
# take acoustic models as an example
|
||||
freezing_enabled: true # main switch to enable parameter freezing
|
||||
frozen_params: # prefix rules to freeze specific parameters during training
|
||||
- model.fs2.encoder
|
||||
- model.fs2.pitch_embed
|
||||
```
|
||||
|
||||
You may interrupt the training and change the settings above at any time. Sometimes this will cause mismatching optimizer state - and it will be discarded silently.
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,152 @@
|
||||
# Getting Started
|
||||
|
||||
## Installation
|
||||
|
||||
### Environments and dependencies
|
||||
|
||||
DiffSinger requires Python 3.10 or later. We strongly recommend you create a virtual environment via Conda, venv or uv before installing dependencies.
|
||||
|
||||
1. Install The latest PyTorch following the [official instructions](https://pytorch.org/get-started/locally/) according to your OS and hardware. We recommend using the latest stable release that is >= 2.4.0.
|
||||
|
||||
2. Install other dependencies via the following command:
|
||||
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
### Concepts and materials
|
||||
|
||||
Before you proceed, it is necessary to understand some fundamental concepts in this repository and prepare some materials and assets. See [fundamental concepts and materials](BestPractices.md#fundamental-concepts-and-materials) for detailed information.
|
||||
|
||||
## Configuration
|
||||
|
||||
Every model needs a configuration file to run preprocessing, training, inference and deployment. Templates of configurations files are in [configs/templates](../configs/templates). Please **copy** the templates to your own data directory before you edit them.
|
||||
|
||||
Before you continue, it is highly recommended to read through [Best Practices](BestPractices.md), which is a more detailed tutorial on how to configure your experiments.
|
||||
|
||||
For more details about configurable parameters, see [Configuration Schemas](ConfigurationSchemas.md).
|
||||
|
||||
> Tips: to see which parameters are required or recommended to be edited, you can search by _customizability_ in the configuration schemas.
|
||||
|
||||
## Preprocessing
|
||||
|
||||
Raw data pieces and transcriptions should be binarized into dataset files before training. Before doing this step, please ensure all required configurations like `raw_data_dir` and `binary_data_dir` are set properly, and all your desired functionalities and features are enabled and configured.
|
||||
|
||||
Assume that you have a configuration file called `my_config.yaml`. Run:
|
||||
|
||||
```bash
|
||||
python scripts/binarize.py --config my_config.yaml
|
||||
```
|
||||
|
||||
Preprocessing can be accelerated through multiprocessing. See [binarization_args.num_workers](ConfigurationSchemas.md#binarization_args.num_workers) for more explanations.
|
||||
|
||||
## Training
|
||||
|
||||
Assume that you have a configuration file called `my_config.yaml` and the name of your model is `my_experiment`. Run:
|
||||
|
||||
```bash
|
||||
python scripts/train.py --config my_config.yaml --exp_name my_experiment --reset
|
||||
```
|
||||
|
||||
Checkpoints will be saved at the `checkpoints/my_experiment/` directory. When interrupting the program and running the above command again, the training resumes automatically from the latest checkpoint.
|
||||
|
||||
For more suggestions related to training performance, see [performance tuning](BestPractices.md#performance-tuning).
|
||||
|
||||
### TensorBoard
|
||||
|
||||
Run the following command to start the TensorBoard:
|
||||
|
||||
```bash
|
||||
tensorboard --logdir checkpoints/
|
||||
```
|
||||
|
||||
> NOTICE
|
||||
>
|
||||
> If you are training a model with multiple GPUs (DDP), please add `--reload_multifile=true` option when launching TensorBoard, otherwise it may not update properly.
|
||||
|
||||
## Inference
|
||||
|
||||
Inference of DiffSinger is based on DS files. Assume that you have a DS file named `my_song.ds` and your model is named `my_experiment`.
|
||||
|
||||
If your model is a variance model, run:
|
||||
|
||||
```bash
|
||||
python scripts/infer.py variance my_song.ds --exp my_experiment
|
||||
```
|
||||
|
||||
or run
|
||||
|
||||
```bash
|
||||
python scripts/infer.py variance --help
|
||||
```
|
||||
|
||||
for more configurable options.
|
||||
|
||||
If your model is an acoustic model, run:
|
||||
|
||||
```bash
|
||||
python scripts/infer.py acoustic my_song.ds --exp my_experiment
|
||||
```
|
||||
|
||||
or run
|
||||
|
||||
```bash
|
||||
python scripts/infer.py acoustic --help
|
||||
```
|
||||
|
||||
for more configurable options.
|
||||
|
||||
## Deployment
|
||||
|
||||
DiffSinger uses [ONNX](https://onnx.ai/) as the deployment format.
|
||||
|
||||
Assume that you have a model named `my_experiment`.
|
||||
|
||||
If your model is a variance model, run:
|
||||
|
||||
```bash
|
||||
python scripts/export.py variance --exp my_experiment
|
||||
```
|
||||
|
||||
or run
|
||||
|
||||
```bash
|
||||
python scripts/export.py variance --help
|
||||
```
|
||||
|
||||
for more configurable options.
|
||||
|
||||
If your model is an acoustic model, run:
|
||||
|
||||
```bash
|
||||
python scripts/export.py acoustic --exp my_experiment
|
||||
```
|
||||
|
||||
or run
|
||||
|
||||
```bash
|
||||
python scripts/export.py acoustic --help
|
||||
```
|
||||
|
||||
for more configurable options.
|
||||
|
||||
To export an NSF-HiFiGAN vocoder checkpoint, run:
|
||||
|
||||
```bash
|
||||
python scripts/export.py nsf-hifigan --config CONFIG --ckpt CKPT
|
||||
```
|
||||
|
||||
where `CONFIG` is a configuration file that has configured the same mel parameters as the vocoder (can be configs/acoustic.yaml for most cases) and `CKPT` is the path of the checkpoint to be exported.
|
||||
|
||||
For more configurable options, run
|
||||
|
||||
```bash
|
||||
python scripts/export.py nsf-hifigan --help
|
||||
```
|
||||
|
||||
## Other utilities
|
||||
|
||||
There are other useful CLI tools in the [scripts/](../scripts) directory not mentioned above:
|
||||
|
||||
- drop_spk.py - delete speaker embeddings from checkpoints (for data security reasons when distributing models)
|
||||
- vocoder.py - bypass the acoustic model and only run the vocoder on given mel-spectrograms
|
||||
File diff suppressed because one or more lines are too long
Binary file not shown.
|
After Width: | Height: | Size: 112 KiB |
File diff suppressed because one or more lines are too long
Binary file not shown.
|
After Width: | Height: | Size: 73 KiB |
File diff suppressed because one or more lines are too long
Binary file not shown.
|
After Width: | Height: | Size: 226 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 517 KiB |
File diff suppressed because it is too large
Load Diff
@@ -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'])
|
||||
@@ -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)
|
||||
@@ -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)]
|
||||
@@ -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'])
|
||||
@@ -0,0 +1,70 @@
|
||||
import torch.nn
|
||||
from torch import nn
|
||||
|
||||
from .convnext import ConvNeXtDecoder
|
||||
from utils import filter_kwargs
|
||||
|
||||
AUX_DECODERS = {
|
||||
'convnext': ConvNeXtDecoder
|
||||
}
|
||||
AUX_LOSSES = {
|
||||
'convnext': nn.L1Loss
|
||||
}
|
||||
|
||||
|
||||
def build_aux_decoder(
|
||||
in_dims: int, out_dims: int,
|
||||
aux_decoder_arch: str, aux_decoder_args: dict
|
||||
) -> torch.nn.Module:
|
||||
decoder_cls = AUX_DECODERS[aux_decoder_arch]
|
||||
kwargs = filter_kwargs(aux_decoder_args, decoder_cls)
|
||||
return AUX_DECODERS[aux_decoder_arch](in_dims, out_dims, **kwargs)
|
||||
|
||||
|
||||
def build_aux_loss(aux_decoder_arch):
|
||||
return AUX_LOSSES[aux_decoder_arch]()
|
||||
|
||||
|
||||
class AuxDecoderAdaptor(nn.Module):
|
||||
def __init__(self, in_dims: int, out_dims: int, num_feats: int,
|
||||
spec_min: list, spec_max: list,
|
||||
aux_decoder_arch: str, aux_decoder_args: dict):
|
||||
super().__init__()
|
||||
self.decoder = build_aux_decoder(
|
||||
in_dims=in_dims, out_dims=out_dims * num_feats,
|
||||
aux_decoder_arch=aux_decoder_arch,
|
||||
aux_decoder_args=aux_decoder_args
|
||||
)
|
||||
self.out_dims = out_dims
|
||||
self.n_feats = num_feats
|
||||
if spec_min is not None and spec_max is not None:
|
||||
# spec: [B, T, M] or [B, F, T, M]
|
||||
# spec_min and spec_max: [1, 1, M] or [1, 1, F, M] => transpose(-3, -2) => [1, 1, M] or [1, F, 1, M]
|
||||
spec_min = torch.FloatTensor(spec_min)[None, None, :].transpose(-3, -2)
|
||||
spec_max = torch.FloatTensor(spec_max)[None, None, :].transpose(-3, -2)
|
||||
self.register_buffer('spec_min', spec_min, persistent=False)
|
||||
self.register_buffer('spec_max', spec_max, persistent=False)
|
||||
|
||||
def norm_spec(self, x):
|
||||
k = (self.spec_max - self.spec_min) / 2.
|
||||
b = (self.spec_max + self.spec_min) / 2.
|
||||
return (x - b) / k
|
||||
|
||||
def denorm_spec(self, x):
|
||||
k = (self.spec_max - self.spec_min) / 2.
|
||||
b = (self.spec_max + self.spec_min) / 2.
|
||||
return x * k + b
|
||||
|
||||
def forward(self, condition, infer=False):
|
||||
x = self.decoder(condition, infer=infer) # [B, T, F x C]
|
||||
|
||||
if self.n_feats > 1:
|
||||
# This is the temporary solution since PyTorch 1.13
|
||||
# does not support exporting aten::unflatten to ONNX
|
||||
# x = x.unflatten(dim=2, sizes=(self.n_feats, self.in_dims))
|
||||
x = x.reshape(-1, x.shape[1], self.n_feats, self.out_dims) # [B, T, F, C]
|
||||
x = x.transpose(1, 2) # [B, F, T, C]
|
||||
if infer:
|
||||
x = self.denorm_spec(x)
|
||||
|
||||
return x # [B, T, C] or [B, F, T, C]
|
||||
@@ -0,0 +1,89 @@
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from modules.commons.common_layers import AdamWConv1d
|
||||
|
||||
|
||||
class ConvNeXtBlock(nn.Module):
|
||||
"""ConvNeXt Block adapted from https://github.com/facebookresearch/ConvNeXt to 1D audio signal.
|
||||
|
||||
Args:
|
||||
dim (int): Number of input channels.
|
||||
intermediate_dim (int): Dimensionality of the intermediate layer.
|
||||
layer_scale_init_value (float, optional): Initial value for the layer scale. None means no scaling.
|
||||
Defaults to None.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
intermediate_dim: int,
|
||||
layer_scale_init_value: Optional[float] = None, drop_out: float = 0.0
|
||||
|
||||
):
|
||||
super().__init__()
|
||||
self.dwconv = nn.Conv1d(dim, dim, kernel_size=7, padding=3, groups=dim) # depthwise conv
|
||||
|
||||
self.norm = nn.LayerNorm(dim, eps=1e-6)
|
||||
self.pwconv1 = nn.Linear(dim, intermediate_dim) # pointwise/1x1 convs, implemented with linear layers
|
||||
self.act = nn.GELU()
|
||||
self.pwconv2 = nn.Linear(intermediate_dim, dim)
|
||||
self.gamma = (
|
||||
nn.Parameter(layer_scale_init_value * torch.ones(dim), requires_grad=True)
|
||||
if layer_scale_init_value > 0
|
||||
else None
|
||||
)
|
||||
# self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
|
||||
self.drop_path = nn.Identity()
|
||||
self.dropout = nn.Dropout(drop_out) if drop_out > 0. else nn.Identity()
|
||||
|
||||
def forward(self, x: torch.Tensor, ) -> torch.Tensor:
|
||||
residual = x
|
||||
x = self.dwconv(x)
|
||||
x = x.transpose(1, 2) # (B, C, T) -> (B, T, C)
|
||||
|
||||
x = self.norm(x)
|
||||
x = self.pwconv1(x)
|
||||
x = self.act(x)
|
||||
x = self.pwconv2(x)
|
||||
if self.gamma is not None:
|
||||
x = self.gamma * x
|
||||
x = x.transpose(1, 2) # (B, T, C) -> (B, C, T)
|
||||
x = self.dropout(x)
|
||||
|
||||
x = residual + self.drop_path(x)
|
||||
return x
|
||||
|
||||
|
||||
class ConvNeXtDecoder(nn.Module):
|
||||
def __init__(
|
||||
self, in_dims, out_dims, /, *,
|
||||
num_channels=512, num_layers=6, kernel_size=7, dropout_rate=0.1
|
||||
):
|
||||
super().__init__()
|
||||
self.inconv = nn.Conv1d(
|
||||
in_dims, num_channels, kernel_size,
|
||||
stride=1, padding=(kernel_size - 1) // 2
|
||||
)
|
||||
self.conv = nn.ModuleList(
|
||||
ConvNeXtBlock(
|
||||
dim=num_channels, intermediate_dim=num_channels * 4,
|
||||
layer_scale_init_value=1e-6, drop_out=dropout_rate
|
||||
) for _ in range(num_layers)
|
||||
)
|
||||
self.outconv = AdamWConv1d(
|
||||
num_channels, out_dims, kernel_size,
|
||||
stride=1, padding=(kernel_size - 1) // 2
|
||||
)
|
||||
|
||||
# noinspection PyUnusedLocal
|
||||
def forward(self, x, infer=False):
|
||||
x = x.transpose(1, 2)
|
||||
x = self.inconv(x)
|
||||
for conv in self.conv:
|
||||
x = conv(x)
|
||||
x = self.outconv(x)
|
||||
x = x.transpose(1, 2)
|
||||
return x
|
||||
@@ -0,0 +1,20 @@
|
||||
import torch.nn
|
||||
from modules.backbones.wavenet import WaveNet
|
||||
from modules.backbones.lynxnet import LYNXNet
|
||||
from modules.backbones.lynxnet2 import LYNXNet2
|
||||
from utils import filter_kwargs
|
||||
|
||||
BACKBONES = {
|
||||
'wavenet': WaveNet,
|
||||
'lynxnet': LYNXNet,
|
||||
'lynxnet2': LYNXNet2,
|
||||
}
|
||||
|
||||
|
||||
def build_backbone(
|
||||
out_dims: int, num_feats: int,
|
||||
backbone_type: str, backbone_args: dict
|
||||
) -> torch.nn.Module:
|
||||
backbone = BACKBONES[backbone_type]
|
||||
kwargs = filter_kwargs(backbone_args, backbone)
|
||||
return BACKBONES[backbone_type](out_dims, num_feats, **kwargs)
|
||||
@@ -0,0 +1,147 @@
|
||||
# refer to：
|
||||
# https://github.com/CNChTu/Diffusion-SVC/blob/v2.0_dev/diffusion/naive_v2/model_conformer_naive.py
|
||||
# https://github.com/CNChTu/Diffusion-SVC/blob/v2.0_dev/diffusion/naive_v2/naive_v2_diff.py
|
||||
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from modules.commons.common_layers import SinusoidalPosEmb, SwiGLU, Transpose, AdamWConv1d
|
||||
from modules.commons.common_layers import KaimingNormalConv1d as Conv1d
|
||||
from utils.hparams import hparams
|
||||
|
||||
|
||||
class LYNXConvModule(nn.Module):
|
||||
@staticmethod
|
||||
def calc_same_padding(kernel_size):
|
||||
pad = kernel_size // 2
|
||||
return pad, pad - (kernel_size + 1) % 2
|
||||
|
||||
def __init__(self, dim, expansion_factor, kernel_size=31, activation='PReLU', dropout=0.0):
|
||||
super().__init__()
|
||||
inner_dim = dim * expansion_factor
|
||||
activation_classes = {
|
||||
'SiLU': nn.SiLU,
|
||||
'ReLU': nn.ReLU,
|
||||
'PReLU': lambda: nn.PReLU(inner_dim)
|
||||
}
|
||||
activation = activation if activation is not None else 'PReLU'
|
||||
if activation not in activation_classes:
|
||||
raise ValueError(f'{activation} is not a valid activation')
|
||||
_activation = activation_classes[activation]()
|
||||
padding = self.calc_same_padding(kernel_size)
|
||||
if float(dropout) > 0.:
|
||||
_dropout = nn.Dropout(dropout)
|
||||
else:
|
||||
_dropout = nn.Identity()
|
||||
self.net = nn.Sequential(
|
||||
nn.LayerNorm(dim),
|
||||
Transpose((1, 2)),
|
||||
nn.Conv1d(dim, inner_dim * 2, 1),
|
||||
SwiGLU(dim=1),
|
||||
nn.Conv1d(inner_dim, inner_dim, kernel_size=kernel_size, padding=padding[0], groups=inner_dim),
|
||||
_activation,
|
||||
nn.Conv1d(inner_dim, dim, 1),
|
||||
Transpose((1, 2)),
|
||||
_dropout
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
return self.net(x)
|
||||
|
||||
|
||||
class LYNXNetResidualLayer(nn.Module):
|
||||
def __init__(self, dim_cond, dim, expansion_factor, kernel_size=31, activation='PReLU', dropout=0.0):
|
||||
super().__init__()
|
||||
self.diffusion_projection = nn.Conv1d(dim, dim, 1)
|
||||
self.conditioner_projection = nn.Conv1d(dim_cond, dim, 1)
|
||||
self.convmodule = LYNXConvModule(dim=dim, expansion_factor=expansion_factor, kernel_size=kernel_size,
|
||||
activation=activation, dropout=dropout)
|
||||
|
||||
def forward(self, x, conditioner, diffusion_step, front_cond_inject=False):
|
||||
if front_cond_inject:
|
||||
x = x + self.conditioner_projection(conditioner)
|
||||
res_x = x
|
||||
else:
|
||||
res_x = x
|
||||
x = x + self.conditioner_projection(conditioner)
|
||||
x = x + self.diffusion_projection(diffusion_step)
|
||||
x = x.transpose(1, 2)
|
||||
x = self.convmodule(x) # (#batch, dim, length)
|
||||
x = x.transpose(1, 2) + res_x
|
||||
return x # (#batch, length, dim)
|
||||
|
||||
|
||||
class LYNXNet(nn.Module):
|
||||
def __init__(self, in_dims, n_feats, *, num_layers=6, num_channels=512, expansion_factor=2, kernel_size=31,
|
||||
activation='PReLU', dropout_rate=0.0, strong_cond=False):
|
||||
"""
|
||||
LYNXNet(Linear Gated Depthwise Separable Convolution Network)
|
||||
TIPS:You can control the style of the generated results by modifying the 'activation',
|
||||
- 'PReLU'(default) : Similar to WaveNet
|
||||
- 'SiLU' : Voice will be more pronounced, not recommended for use under DDPM
|
||||
- 'ReLU' : Contrary to 'SiLU', Voice will be weakened
|
||||
"""
|
||||
super().__init__()
|
||||
self.in_dims = in_dims
|
||||
self.n_feats = n_feats
|
||||
self.input_projection = Conv1d(in_dims * n_feats, num_channels, 1)
|
||||
self.diffusion_embedding = nn.Sequential(
|
||||
SinusoidalPosEmb(num_channels),
|
||||
nn.Linear(num_channels, num_channels * 4),
|
||||
nn.GELU(),
|
||||
nn.Linear(num_channels * 4, num_channels),
|
||||
)
|
||||
self.residual_layers = nn.ModuleList(
|
||||
[
|
||||
LYNXNetResidualLayer(
|
||||
dim_cond=hparams['hidden_size'],
|
||||
dim=num_channels,
|
||||
expansion_factor=expansion_factor,
|
||||
kernel_size=kernel_size,
|
||||
activation=activation,
|
||||
dropout=dropout_rate
|
||||
)
|
||||
for _ in range(num_layers)
|
||||
]
|
||||
)
|
||||
self.norm = nn.LayerNorm(num_channels)
|
||||
self.output_projection = AdamWConv1d(num_channels, in_dims * n_feats, kernel_size=1)
|
||||
self.strong_cond = strong_cond
|
||||
nn.init.zeros_(self.output_projection.weight)
|
||||
|
||||
def forward(self, spec, diffusion_step, cond):
|
||||
"""
|
||||
:param spec: [B, F, M, T]
|
||||
:param diffusion_step: [B, 1]
|
||||
:param cond: [B, H, T]
|
||||
:return:
|
||||
"""
|
||||
|
||||
if self.n_feats == 1:
|
||||
x = spec[:, 0] # [B, M, T]
|
||||
else:
|
||||
x = spec.flatten(start_dim=1, end_dim=2) # [B, F x M, T]
|
||||
|
||||
x = self.input_projection(x) # x [B, residual_channel, T]
|
||||
if not self.strong_cond:
|
||||
x = F.gelu(x)
|
||||
|
||||
diffusion_step = self.diffusion_embedding(diffusion_step).unsqueeze(-1)
|
||||
|
||||
for layer in self.residual_layers:
|
||||
x = layer(x, cond, diffusion_step, front_cond_inject=self.strong_cond)
|
||||
|
||||
# post-norm
|
||||
x = self.norm(x.transpose(1, 2)).transpose(1, 2)
|
||||
|
||||
# output_projection
|
||||
x = self.output_projection(x) # [B, 128, T]
|
||||
|
||||
if self.n_feats == 1:
|
||||
x = x[:, None, :, :]
|
||||
else:
|
||||
# This is the temporary solution since PyTorch 1.13
|
||||
# does not support exporting aten::unflatten to ONNX
|
||||
# x = x.unflatten(dim=1, sizes=(self.n_feats, self.in_dims))
|
||||
x = x.reshape(-1, self.n_feats, self.in_dims, x.shape[2])
|
||||
return x
|
||||
@@ -0,0 +1,115 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from modules.commons.common_layers import SinusoidalPosEmb, SwiGLU, ATanGLU, Transpose, AdamWLinear
|
||||
from utils.hparams import hparams
|
||||
|
||||
|
||||
class LYNXNet2Block(nn.Module):
|
||||
def __init__(self, dim, expansion_factor, kernel_size=31, dropout=0., glu_type='swiglu'):
|
||||
super().__init__()
|
||||
inner_dim = int(dim * expansion_factor)
|
||||
if glu_type == 'swiglu':
|
||||
_glu = SwiGLU()
|
||||
elif glu_type == 'atanglu':
|
||||
_glu = ATanGLU()
|
||||
else:
|
||||
raise ValueError(f'{glu_type} is not a valid activation')
|
||||
if float(dropout) > 0.:
|
||||
_dropout = nn.Dropout(dropout)
|
||||
else:
|
||||
_dropout = nn.Identity()
|
||||
self.net = nn.Sequential(
|
||||
nn.LayerNorm(dim),
|
||||
Transpose((1, 2)),
|
||||
nn.Conv1d(dim, dim, kernel_size=kernel_size, padding=kernel_size // 2, groups=dim),
|
||||
Transpose((1, 2)),
|
||||
nn.Linear(dim, inner_dim * 2),
|
||||
_glu,
|
||||
nn.Linear(inner_dim, inner_dim * 2),
|
||||
_glu,
|
||||
nn.Linear(inner_dim, dim),
|
||||
_dropout
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
return x + self.net(x)
|
||||
|
||||
|
||||
class LYNXNet2(nn.Module):
|
||||
def __init__(self, in_dims, n_feats, *, num_layers=6, num_channels=512, expansion_factor=1, kernel_size=31,
|
||||
dropout_rate=0.0, use_conditioner_cache=False, glu_type='swiglu'):
|
||||
"""
|
||||
LYNXNet2(Linear Gated Depthwise Separable Convolution Network Version 2)
|
||||
"""
|
||||
super().__init__()
|
||||
self.in_dims = in_dims
|
||||
self.n_feats = n_feats
|
||||
self.input_projection = nn.Linear(in_dims * n_feats, num_channels)
|
||||
self.use_conditioner_cache = use_conditioner_cache
|
||||
if self.use_conditioner_cache:
|
||||
# Conv1d is used for condition cache compatibility
|
||||
self.conditioner_projection = nn.Conv1d(hparams['hidden_size'], num_channels, 1)
|
||||
else:
|
||||
self.conditioner_projection = nn.Linear(hparams['hidden_size'], num_channels)
|
||||
self.diffusion_embedding = nn.Sequential(
|
||||
SinusoidalPosEmb(num_channels),
|
||||
nn.Linear(num_channels, num_channels * 4),
|
||||
nn.GELU(),
|
||||
nn.Linear(num_channels * 4, num_channels),
|
||||
)
|
||||
self.residual_layers = nn.ModuleList(
|
||||
[
|
||||
LYNXNet2Block(
|
||||
dim=num_channels,
|
||||
expansion_factor=expansion_factor,
|
||||
kernel_size=kernel_size,
|
||||
dropout=dropout_rate,
|
||||
glu_type=glu_type
|
||||
)
|
||||
for _ in range(num_layers)
|
||||
]
|
||||
)
|
||||
self.norm = nn.LayerNorm(num_channels)
|
||||
self.output_projection = AdamWLinear(num_channels, in_dims * n_feats)
|
||||
nn.init.kaiming_normal_(self.input_projection.weight)
|
||||
nn.init.kaiming_normal_(self.conditioner_projection.weight)
|
||||
nn.init.zeros_(self.output_projection.weight)
|
||||
|
||||
def forward(self, spec, diffusion_step, cond):
|
||||
"""
|
||||
:param spec: [B, F, M, T]
|
||||
:param diffusion_step: [B, 1]
|
||||
:param cond: [B, H, T]
|
||||
:return:
|
||||
"""
|
||||
|
||||
if self.n_feats == 1:
|
||||
x = spec[:, 0] # [B, M, T]
|
||||
else:
|
||||
x = spec.flatten(start_dim=1, end_dim=2) # [B, F x M, T]
|
||||
|
||||
x = self.input_projection(x.transpose(1, 2)) # [B, T, F x M]
|
||||
if self.use_conditioner_cache:
|
||||
x = x + self.conditioner_projection(cond).transpose(1, 2)
|
||||
else:
|
||||
x = x + self.conditioner_projection(cond.transpose(1, 2))
|
||||
x = x + self.diffusion_embedding(diffusion_step).unsqueeze(1)
|
||||
|
||||
for layer in self.residual_layers:
|
||||
x = layer(x)
|
||||
|
||||
# post-norm
|
||||
x = self.norm(x)
|
||||
|
||||
# output projection
|
||||
x = self.output_projection(x).transpose(1, 2) # [B, 128, T]
|
||||
|
||||
if self.n_feats == 1:
|
||||
x = x[:, None, :, :]
|
||||
else:
|
||||
# Using reshape instead of unflatten for ONNX export compatibility
|
||||
# x = x.unflatten(dim=1, sizes=(self.n_feats, self.in_dims))
|
||||
x = x.reshape(-1, self.n_feats, self.in_dims, x.shape[2])
|
||||
return x
|
||||
@@ -0,0 +1,104 @@
|
||||
import math
|
||||
from math import sqrt
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from modules.commons.common_layers import SinusoidalPosEmb, AdamWConv1d
|
||||
from modules.commons.common_layers import KaimingNormalConv1d as Conv1d
|
||||
from utils.hparams import hparams
|
||||
|
||||
|
||||
class ResidualBlock(nn.Module):
|
||||
def __init__(self, encoder_hidden, residual_channels, dilation):
|
||||
super().__init__()
|
||||
self.residual_channels = residual_channels
|
||||
self.dilated_conv = nn.Conv1d(
|
||||
residual_channels,
|
||||
2 * residual_channels,
|
||||
kernel_size=3,
|
||||
padding=dilation,
|
||||
dilation=dilation
|
||||
)
|
||||
self.diffusion_projection = nn.Linear(residual_channels, residual_channels)
|
||||
self.conditioner_projection = nn.Conv1d(encoder_hidden, 2 * residual_channels, 1)
|
||||
self.output_projection = nn.Conv1d(residual_channels, 2 * residual_channels, 1)
|
||||
|
||||
def forward(self, x, conditioner, diffusion_step):
|
||||
diffusion_step = self.diffusion_projection(diffusion_step).unsqueeze(-1)
|
||||
conditioner = self.conditioner_projection(conditioner)
|
||||
y = x + diffusion_step
|
||||
|
||||
y = self.dilated_conv(y) + conditioner
|
||||
|
||||
# Using torch.split instead of torch.chunk to avoid using onnx::Slice
|
||||
gate, filter = torch.split(y, [self.residual_channels, self.residual_channels], dim=1)
|
||||
y = torch.sigmoid(gate) * torch.tanh(filter)
|
||||
|
||||
y = self.output_projection(y)
|
||||
|
||||
# Using torch.split instead of torch.chunk to avoid using onnx::Slice
|
||||
residual, skip = torch.split(y, [self.residual_channels, self.residual_channels], dim=1)
|
||||
return (x + residual) / math.sqrt(2.0), skip
|
||||
|
||||
|
||||
class WaveNet(nn.Module):
|
||||
def __init__(self, in_dims, n_feats, *, num_layers=20, num_channels=256, dilation_cycle_length=4):
|
||||
super().__init__()
|
||||
self.in_dims = in_dims
|
||||
self.n_feats = n_feats
|
||||
self.input_projection = Conv1d(in_dims * n_feats, num_channels, 1)
|
||||
self.diffusion_embedding = SinusoidalPosEmb(num_channels)
|
||||
self.mlp = nn.Sequential(
|
||||
nn.Linear(num_channels, num_channels * 4),
|
||||
nn.Mish(),
|
||||
nn.Linear(num_channels * 4, num_channels)
|
||||
)
|
||||
self.residual_layers = nn.ModuleList([
|
||||
ResidualBlock(
|
||||
encoder_hidden=hparams['hidden_size'],
|
||||
residual_channels=num_channels,
|
||||
dilation=2 ** (i % dilation_cycle_length)
|
||||
)
|
||||
for i in range(num_layers)
|
||||
])
|
||||
self.skip_projection = Conv1d(num_channels, num_channels, 1)
|
||||
self.output_projection = AdamWConv1d(num_channels, in_dims * n_feats, 1)
|
||||
nn.init.zeros_(self.output_projection.weight)
|
||||
|
||||
def forward(self, spec, diffusion_step, cond):
|
||||
"""
|
||||
:param spec: [B, F, M, T]
|
||||
:param diffusion_step: [B, 1]
|
||||
:param cond: [B, H, T]
|
||||
:return:
|
||||
"""
|
||||
if self.n_feats == 1:
|
||||
# Use indexing instead of squeeze to avoid emitting an onnx::If
|
||||
# whose branches have different rank, which breaks shape inference
|
||||
# for the downstream Conv on PyTorch >= 2.0.
|
||||
x = spec[:, 0] # [B, M, T]
|
||||
else:
|
||||
x = spec.flatten(start_dim=1, end_dim=2) # [B, F x M, T]
|
||||
x = self.input_projection(x) # [B, C, T]
|
||||
|
||||
x = F.relu(x)
|
||||
diffusion_step = self.diffusion_embedding(diffusion_step)
|
||||
diffusion_step = self.mlp(diffusion_step)
|
||||
skip = []
|
||||
for layer in self.residual_layers:
|
||||
x, skip_connection = layer(x, cond, diffusion_step)
|
||||
skip.append(skip_connection)
|
||||
|
||||
x = torch.sum(torch.stack(skip), dim=0) / sqrt(len(self.residual_layers))
|
||||
x = self.skip_projection(x)
|
||||
x = F.relu(x)
|
||||
x = self.output_projection(x) # [B, M, T]
|
||||
if self.n_feats == 1:
|
||||
x = x[:, None, :, :]
|
||||
else:
|
||||
# Using reshape instead of unflatten for ONNX export compatibility
|
||||
# x = x.unflatten(dim=1, sizes=(self.n_feats, self.in_dims))
|
||||
x = x.reshape(-1, self.n_feats, self.in_dims, x.shape[2])
|
||||
return x
|
||||
@@ -0,0 +1,427 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import torch.onnx.operators
|
||||
from torch import nn
|
||||
from torch.nn import LayerNorm, ReLU, GELU, SiLU
|
||||
|
||||
import utils
|
||||
|
||||
|
||||
class NormalInitEmbedding(torch.nn.Embedding):
|
||||
def __init__(
|
||||
self,
|
||||
num_embeddings: int,
|
||||
embedding_dim: int,
|
||||
padding_idx: int | None = None,
|
||||
*args,
|
||||
**kwargs
|
||||
):
|
||||
super().__init__(num_embeddings, embedding_dim, *args, padding_idx=padding_idx, **kwargs)
|
||||
nn.init.normal_(self.weight, mean=0, std=self.embedding_dim ** -0.5)
|
||||
if padding_idx is not None:
|
||||
nn.init.constant_(self.weight[padding_idx], 0)
|
||||
|
||||
|
||||
class AdamWLinear(torch.nn.Linear):
|
||||
def __init__(
|
||||
self,
|
||||
in_features: int,
|
||||
out_features: int,
|
||||
*args,
|
||||
bias: bool = True,
|
||||
**kwargs
|
||||
):
|
||||
super().__init__(in_features, out_features, *args, bias=bias, **kwargs)
|
||||
nn.init.xavier_uniform_(self.weight)
|
||||
if bias:
|
||||
nn.init.constant_(self.bias, 0.)
|
||||
|
||||
|
||||
class XavierUniformInitLinear(torch.nn.Linear):
|
||||
def __init__(
|
||||
self,
|
||||
in_features: int,
|
||||
out_features: int,
|
||||
*args,
|
||||
bias: bool = True,
|
||||
**kwargs
|
||||
):
|
||||
super().__init__(in_features, out_features, *args, bias=bias, **kwargs)
|
||||
nn.init.xavier_uniform_(self.weight)
|
||||
if bias:
|
||||
nn.init.constant_(self.bias, 0.)
|
||||
|
||||
|
||||
class SinusoidalPositionalEmbedding(nn.Module):
|
||||
"""This module produces sinusoidal positional embeddings of any length.
|
||||
|
||||
Padding symbols are ignored.
|
||||
"""
|
||||
|
||||
def __init__(self, embedding_dim, padding_idx, init_size=1024):
|
||||
super().__init__()
|
||||
self.embedding_dim = embedding_dim
|
||||
self.padding_idx = padding_idx
|
||||
self.weights = SinusoidalPositionalEmbedding.get_embedding(
|
||||
init_size,
|
||||
embedding_dim,
|
||||
padding_idx,
|
||||
)
|
||||
self.register_buffer('_float_tensor', torch.FloatTensor(1))
|
||||
|
||||
@staticmethod
|
||||
def get_embedding(num_embeddings, embedding_dim, padding_idx=None):
|
||||
"""Build sinusoidal embeddings.
|
||||
|
||||
This matches the implementation in tensor2tensor, but differs slightly
|
||||
from the description in Section 3.5 of "Attention Is All You Need".
|
||||
"""
|
||||
half_dim = embedding_dim // 2
|
||||
emb = math.log(10000) / (half_dim - 1)
|
||||
emb = torch.exp(torch.arange(half_dim, dtype=torch.float) * -emb)
|
||||
emb = torch.arange(num_embeddings, dtype=torch.float).unsqueeze(1) * emb.unsqueeze(0)
|
||||
emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1).view(num_embeddings, -1)
|
||||
if embedding_dim % 2 == 1:
|
||||
# zero pad
|
||||
emb = torch.cat([emb, torch.zeros(num_embeddings, 1)], dim=1)
|
||||
if padding_idx is not None:
|
||||
emb[padding_idx, :] = 0
|
||||
return emb
|
||||
|
||||
def forward(self, x, incremental_state=None, timestep=None, positions=None):
|
||||
"""Input is expected to be of size [bsz x seqlen]."""
|
||||
bsz, seq_len = x.shape[:2]
|
||||
max_pos = self.padding_idx + 1 + seq_len
|
||||
if self.weights is None or max_pos > self.weights.size(0):
|
||||
# recompute/expand embeddings if needed
|
||||
self.weights = SinusoidalPositionalEmbedding.get_embedding(
|
||||
max_pos,
|
||||
self.embedding_dim,
|
||||
self.padding_idx,
|
||||
)
|
||||
self.weights = self.weights.to(self._float_tensor)
|
||||
|
||||
if incremental_state is not None:
|
||||
# positions is the same for every token when decoding a single step
|
||||
pos = timestep.view(-1)[0] + 1 if timestep is not None else seq_len
|
||||
return self.weights[self.padding_idx + pos, :].expand(bsz, 1, -1)
|
||||
|
||||
positions = utils.make_positions(x, self.padding_idx) if positions is None else positions
|
||||
return self.weights.index_select(0, positions.view(-1)).view(bsz, seq_len, -1).detach()
|
||||
|
||||
@staticmethod
|
||||
def max_positions():
|
||||
"""Maximum number of supported positions."""
|
||||
return int(1e5) # an arbitrary large number
|
||||
|
||||
|
||||
class SwiGLU(nn.Module):
|
||||
# Swish-Applies the gated linear unit function.
|
||||
def __init__(self, dim=-1):
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
|
||||
def forward(self, x):
|
||||
# out, gate = x.chunk(2, dim=self.dim)
|
||||
# Using torch.split instead of chunk for ONNX export compatibility.
|
||||
out, gate = torch.split(x, x.size(self.dim) // 2, dim=self.dim)
|
||||
gate = F.silu(gate)
|
||||
if x.dtype == torch.float16:
|
||||
out_min, out_max = torch.aminmax(out.detach())
|
||||
gate_min, gate_max = torch.aminmax(gate.detach())
|
||||
max_abs_out = torch.max(-out_min, out_max).float()
|
||||
max_abs_gate = torch.max(-gate_min, gate_max).float()
|
||||
max_abs_value = max_abs_out * max_abs_gate
|
||||
if max_abs_value > 1000:
|
||||
ratio = (1000 / max_abs_value).half()
|
||||
gate = gate * ratio
|
||||
return (out * gate).clamp(-1000 * ratio, 1000 * ratio) / ratio
|
||||
return out * gate
|
||||
|
||||
|
||||
class ATanGLUFunction(torch.autograd.Function):
|
||||
@staticmethod
|
||||
def forward(ctx, out, gate):
|
||||
atan_gate = torch.atan(gate)
|
||||
decay_out = out / gate.square().add(1.0)
|
||||
ctx.save_for_backward(decay_out, atan_gate)
|
||||
return out * atan_gate
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_output):
|
||||
decay_out, atan_gate = ctx.saved_tensors
|
||||
grad_out_part = grad_output * atan_gate
|
||||
grad_gate_part = grad_output * decay_out
|
||||
return grad_out_part, grad_gate_part
|
||||
|
||||
|
||||
class ATanGLU(nn.Module):
|
||||
# ArcTan-Applies the gated linear unit function.
|
||||
def __init__(self, dim=-1):
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
|
||||
def forward(self, x):
|
||||
# out, gate = x.chunk(2, dim=self.dim)
|
||||
# Using torch.split instead of chunk for ONNX export compatibility.
|
||||
out, gate = torch.split(x, x.size(self.dim) // 2, dim=self.dim)
|
||||
if self.training:
|
||||
return ATanGLUFunction.apply(out, gate)
|
||||
else:
|
||||
return out * torch.atan(gate)
|
||||
|
||||
|
||||
class AdamWConv1d(torch.nn.Conv1d):
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
nn.init.kaiming_normal_(self.weight)
|
||||
|
||||
|
||||
class KaimingNormalConv1d(torch.nn.Conv1d):
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
nn.init.kaiming_normal_(self.weight)
|
||||
|
||||
|
||||
class Transpose(nn.Module):
|
||||
def __init__(self, dims):
|
||||
super().__init__()
|
||||
assert len(dims) == 2, 'dims must be a tuple of two dimensions'
|
||||
self.dims = dims
|
||||
|
||||
def forward(self, x):
|
||||
return x.transpose(*self.dims)
|
||||
|
||||
|
||||
class Mixed_LayerNorm(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
channels: int,
|
||||
condition_channels: int,
|
||||
beta_distribution_concentration: float = 0.2,
|
||||
eps: float = 1e-5,
|
||||
bias: bool = True
|
||||
):
|
||||
super().__init__()
|
||||
self.channels = channels
|
||||
self.eps = eps
|
||||
|
||||
self.beta_distribution = torch.distributions.Beta(
|
||||
beta_distribution_concentration,
|
||||
beta_distribution_concentration
|
||||
)
|
||||
|
||||
self.affine = XavierUniformInitLinear(condition_channels, channels * 2, bias=bias)
|
||||
if self.affine.bias is not None:
|
||||
self.affine.bias.data[:channels] = 0 # betas (shift)
|
||||
self.affine.bias.data[channels:] = 1 # gammas (scale)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.FloatTensor,
|
||||
condition: torch.FloatTensor # -> shape [Batch, Cond_d]
|
||||
) -> torch.FloatTensor:
|
||||
x = F.layer_norm(x, normalized_shape=(self.channels,), weight=None, bias=None, eps=self.eps)
|
||||
|
||||
affine_params = self.affine(condition)
|
||||
if affine_params.ndim == 2:
|
||||
affine_params = affine_params.unsqueeze(1)
|
||||
betas, gammas = torch.split(affine_params, self.channels, dim=-1)
|
||||
|
||||
if not self.training or x.size(0) == 1:
|
||||
return gammas * x + betas
|
||||
|
||||
shuffle_indices = torch.randperm(x.size(0), device=x.device)
|
||||
shuffled_betas = betas[shuffle_indices]
|
||||
shuffled_gammas = gammas[shuffle_indices]
|
||||
|
||||
beta_samples = self.beta_distribution.sample((x.size(0), 1, 1)).to(x.device)
|
||||
mixed_betas = beta_samples * betas + (1 - beta_samples) * shuffled_betas
|
||||
mixed_gammas = beta_samples * gammas + (1 - beta_samples) * shuffled_gammas
|
||||
|
||||
return mixed_gammas * x + mixed_betas
|
||||
|
||||
|
||||
class TransformerFFNLayer(nn.Module):
|
||||
def __init__(self, hidden_size, filter_size, kernel_size=1, dropout=0., act='gelu'):
|
||||
super().__init__()
|
||||
self.kernel_size = kernel_size
|
||||
self.dropout = dropout
|
||||
self.act = act
|
||||
filter_size_1 = filter_size
|
||||
if self.act == 'relu':
|
||||
self.act_fn = ReLU()
|
||||
elif self.act == 'gelu':
|
||||
self.act_fn = GELU()
|
||||
elif self.act == 'swish':
|
||||
self.act_fn = SiLU()
|
||||
elif self.act == 'swiglu':
|
||||
self.act_fn = SwiGLU()
|
||||
filter_size_1 = filter_size * 2
|
||||
elif self.act == 'atanglu':
|
||||
self.act_fn = ATanGLU()
|
||||
filter_size_1 = filter_size * 2
|
||||
else:
|
||||
raise ValueError(f'{act} is not a valid activation')
|
||||
self.ffn_1 = nn.Conv1d(hidden_size, filter_size_1, kernel_size, padding=kernel_size // 2)
|
||||
self.ffn_2 = XavierUniformInitLinear(filter_size, hidden_size)
|
||||
|
||||
def forward(self, x):
|
||||
# x: B x T x C
|
||||
x = self.ffn_1(x.transpose(1, 2)).transpose(1, 2)
|
||||
x = x * self.kernel_size ** -0.5
|
||||
|
||||
x = self.act_fn(x)
|
||||
x = F.dropout(x, self.dropout, training=self.training)
|
||||
x = self.ffn_2(x)
|
||||
return x
|
||||
|
||||
|
||||
class MultiheadSelfAttentionWithRoPE(nn.Module):
|
||||
def __init__(self, embed_dim, num_heads, dropout=0.1, bias=False, rotary_embed=None):
|
||||
super().__init__()
|
||||
assert embed_dim % num_heads == 0, "Embedding dimension must be divisible by number of heads"
|
||||
|
||||
self.embed_dim = embed_dim
|
||||
self.num_heads = num_heads
|
||||
self.head_dim = embed_dim // num_heads
|
||||
|
||||
# Linear layers for Q, K, V projections
|
||||
self.in_proj = nn.Linear(embed_dim, embed_dim * 3, bias=bias)
|
||||
|
||||
# Final linear layer after concatenation
|
||||
self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
|
||||
|
||||
# Dropout layer
|
||||
self.dropout = nn.Dropout(dropout)
|
||||
|
||||
# Rotary Embeddings
|
||||
self.rotary_embed = rotary_embed
|
||||
|
||||
# Initialization parameters
|
||||
nn.init.xavier_uniform_(self.in_proj.weight)
|
||||
nn.init.xavier_uniform_(self.out_proj.weight)
|
||||
if bias:
|
||||
nn.init.constant_(self.in_proj.bias, 0.0)
|
||||
nn.init.constant_(self.out_proj.bias, 0.0)
|
||||
|
||||
def forward(self, x, key_padding_mask=None):
|
||||
# x: (B, L, C)
|
||||
# key_padding_mask: (B, L)
|
||||
batch_size, seq_len, embed_dim = x.size()
|
||||
|
||||
# Project inputs to Q, K, V
|
||||
Q, K, V = torch.split(self.in_proj(x), self.embed_dim, dim=-1)
|
||||
|
||||
# Reshape Q, K, V for multi-head attention
|
||||
Q = Q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2) # (B, H, L, D)
|
||||
K = K.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2) # (B, H, L, D)
|
||||
V = V.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2) # (B, H, L, D)
|
||||
|
||||
# Apply RoPE
|
||||
if self.rotary_embed is not None:
|
||||
Q = self.rotary_embed.rotate_queries_or_keys(Q)
|
||||
K = self.rotary_embed.rotate_queries_or_keys(K)
|
||||
|
||||
# Compute attention scores
|
||||
scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.head_dim) # (B, H, L, L)
|
||||
|
||||
# Apply key padding mask if provided
|
||||
if key_padding_mask is not None:
|
||||
# Expand mask to match attention scores shape
|
||||
mask = key_padding_mask.unsqueeze(1).unsqueeze(1) # (B, 1, 1, L)
|
||||
scores = scores.masked_fill(mask == 1, -np.inf) # Masked positions are set to -inf
|
||||
|
||||
# Compute attention weights
|
||||
attn_weights = F.softmax(scores, dim=-1) # (B, H, L, L)
|
||||
attn_weights = self.dropout(attn_weights)
|
||||
|
||||
# Apply attention weights to V
|
||||
attn_output = torch.matmul(attn_weights, V) # (B, H, L, D)
|
||||
|
||||
# Reshape and concatenate heads
|
||||
attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, embed_dim) # (B, L, C)
|
||||
|
||||
# Final linear projection
|
||||
output = self.out_proj(attn_output) # (B, L, C)
|
||||
|
||||
return output
|
||||
|
||||
|
||||
class EncSALayer(nn.Module):
|
||||
def __init__(self, c, num_heads, dropout, attention_dropout=0.1,
|
||||
relu_dropout=0.1, kernel_size=9, act='gelu', rotary_embed=None,
|
||||
layer_idx=None, mix_ln_layer=None
|
||||
):
|
||||
super().__init__()
|
||||
self.dropout = dropout
|
||||
self.use_mix_ln = (
|
||||
layer_idx is not None
|
||||
and mix_ln_layer is not None
|
||||
and layer_idx in mix_ln_layer
|
||||
)
|
||||
if self.use_mix_ln:
|
||||
self.layer_norm1 = Mixed_LayerNorm(c, c)
|
||||
else:
|
||||
self.layer_norm1 = LayerNorm(c)
|
||||
# Always use the in-house manual attention. With rotary_embed=None this
|
||||
# is a plain multi-head self-attention that is ONNX-export safe across
|
||||
# dynamic sequence lengths. Using torch.nn.MultiheadAttention here was
|
||||
# the source of the "Reshape baked tgt_len" bug on PyTorch >= 2.0
|
||||
# because its SDPA-branched implementation specializes tgt_len to a
|
||||
# Python int and re-injects it into the output Reshape.
|
||||
self.self_attn = MultiheadSelfAttentionWithRoPE(
|
||||
c, num_heads, dropout=attention_dropout, bias=False, rotary_embed=rotary_embed
|
||||
)
|
||||
if self.use_mix_ln:
|
||||
self.layer_norm2 = Mixed_LayerNorm(c, c)
|
||||
else:
|
||||
self.layer_norm2 = LayerNorm(c)
|
||||
self.ffn = TransformerFFNLayer(
|
||||
c, 4 * c, kernel_size=kernel_size, dropout=relu_dropout, act=act
|
||||
)
|
||||
|
||||
def forward(self, x, encoder_padding_mask=None, cond=None, **kwargs):
|
||||
layer_norm_training = kwargs.get('layer_norm_training', None)
|
||||
if layer_norm_training is not None:
|
||||
self.layer_norm1.training = layer_norm_training
|
||||
self.layer_norm2.training = layer_norm_training
|
||||
residual = x
|
||||
if self.use_mix_ln:
|
||||
x = self.layer_norm1(x, cond)
|
||||
else:
|
||||
x = self.layer_norm1(x)
|
||||
x = self.self_attn(x, key_padding_mask=encoder_padding_mask)
|
||||
x = F.dropout(x, self.dropout, training=self.training)
|
||||
x = residual + x
|
||||
x = x * (1 - encoder_padding_mask.float())[..., None]
|
||||
|
||||
residual = x
|
||||
if self.use_mix_ln:
|
||||
x = self.layer_norm2(x, cond)
|
||||
else:
|
||||
x = self.layer_norm2(x)
|
||||
x = self.ffn(x)
|
||||
x = F.dropout(x, self.dropout, training=self.training)
|
||||
x = residual + x
|
||||
x = x * (1 - encoder_padding_mask.float())[..., None]
|
||||
return x
|
||||
|
||||
|
||||
class SinusoidalPosEmb(nn.Module):
|
||||
def __init__(self, dim):
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
|
||||
def forward(self, x):
|
||||
device = x.device
|
||||
half_dim = self.dim // 2
|
||||
emb = math.log(10000) / (half_dim - 1)
|
||||
emb = torch.exp(torch.arange(half_dim, device=device) * -emb)
|
||||
emb = x.unsqueeze(-1) * emb.unsqueeze(0)
|
||||
emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
|
||||
return emb
|
||||
@@ -0,0 +1,113 @@
|
||||
import math
|
||||
import torch
|
||||
|
||||
|
||||
class PositionalEncoding(torch.nn.Module):
|
||||
"""Positional encoding.
|
||||
Args:
|
||||
d_model (int): Embedding dimension.
|
||||
dropout_rate (float): Dropout rate.
|
||||
max_len (int): Maximum input length.
|
||||
reverse (bool): Whether to reverse the input position.
|
||||
"""
|
||||
|
||||
def __init__(self, d_model, dropout_rate, max_len=5000, reverse=False):
|
||||
"""Construct an PositionalEncoding object."""
|
||||
super(PositionalEncoding, self).__init__()
|
||||
self.d_model = d_model
|
||||
self.reverse = reverse
|
||||
self.xscale = math.sqrt(self.d_model)
|
||||
self.dropout = torch.nn.Dropout(p=dropout_rate)
|
||||
self.pe = None
|
||||
self.extend_pe(torch.tensor(0.0).expand(1, max_len))
|
||||
|
||||
def extend_pe(self, x):
|
||||
"""Reset the positional encodings."""
|
||||
if self.pe is not None:
|
||||
if self.pe.size(1) >= x.size(1):
|
||||
if self.pe.dtype != x.dtype or self.pe.device != x.device:
|
||||
self.pe = self.pe.to(dtype=x.dtype, device=x.device)
|
||||
return
|
||||
if self.reverse:
|
||||
position = torch.arange(
|
||||
x.size(1) - 1, -1, -1.0, dtype=torch.float32
|
||||
).unsqueeze(1)
|
||||
else:
|
||||
position = torch.arange(0, x.size(1), dtype=torch.float32).unsqueeze(1)
|
||||
div_term = torch.exp(
|
||||
torch.arange(0, self.d_model, 2, dtype=torch.float32)
|
||||
* -(math.log(10000.0) / self.d_model)
|
||||
)
|
||||
pe = torch.stack([
|
||||
torch.sin(position * div_term),
|
||||
torch.cos(position * div_term)
|
||||
], dim=2).view(-1, self.d_model).unsqueeze(0)
|
||||
self.pe = pe.to(device=x.device, dtype=x.dtype)
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
"""Add positional encoding.
|
||||
Args:
|
||||
x (torch.Tensor): Input tensor (batch, time, `*`).
|
||||
Returns:
|
||||
torch.Tensor: Encoded tensor (batch, time, `*`).
|
||||
"""
|
||||
self.extend_pe(x)
|
||||
x = x * self.xscale + self.pe[:, : x.size(1)]
|
||||
return self.dropout(x)
|
||||
|
||||
|
||||
class ScaledPositionalEncoding(PositionalEncoding):
|
||||
"""Scaled positional encoding module.
|
||||
See Sec. 3.2 https://arxiv.org/abs/1809.08895
|
||||
Args:
|
||||
d_model (int): Embedding dimension.
|
||||
dropout_rate (float): Dropout rate.
|
||||
max_len (int): Maximum input length.
|
||||
"""
|
||||
|
||||
def __init__(self, d_model, dropout_rate, max_len=5000):
|
||||
"""Initialize class."""
|
||||
super().__init__(d_model=d_model, dropout_rate=dropout_rate, max_len=max_len)
|
||||
self.alpha = torch.nn.Parameter(torch.tensor(1.0))
|
||||
|
||||
def reset_parameters(self):
|
||||
"""Reset parameters."""
|
||||
self.alpha.data = torch.tensor(1.0)
|
||||
|
||||
def forward(self, x):
|
||||
"""Add positional encoding.
|
||||
Args:
|
||||
x (torch.Tensor): Input tensor (batch, time, `*`).
|
||||
Returns:
|
||||
torch.Tensor: Encoded tensor (batch, time, `*`).
|
||||
"""
|
||||
self.extend_pe(x)
|
||||
x = x + self.alpha * self.pe[:, : x.size(1)]
|
||||
return self.dropout(x)
|
||||
|
||||
|
||||
class RelPositionalEncoding(PositionalEncoding):
|
||||
"""Relative positional encoding module.
|
||||
See : Appendix B in https://arxiv.org/abs/1901.02860
|
||||
Args:
|
||||
d_model (int): Embedding dimension.
|
||||
dropout_rate (float): Dropout rate.
|
||||
max_len (int): Maximum input length.
|
||||
"""
|
||||
|
||||
def __init__(self, d_model, dropout_rate, max_len=5000):
|
||||
"""Initialize class."""
|
||||
super().__init__(d_model, dropout_rate, max_len, reverse=True)
|
||||
|
||||
def forward(self, x):
|
||||
"""Compute positional encoding.
|
||||
Args:
|
||||
x (torch.Tensor): Input tensor (batch, time, `*`).
|
||||
Returns:
|
||||
torch.Tensor: Encoded tensor (batch, time, `*`).
|
||||
torch.Tensor: Positional embedding tensor (1, time, `*`).
|
||||
"""
|
||||
self.extend_pe(x)
|
||||
x = x * self.xscale
|
||||
pos_emb = self.pe[:, : x.size(1)]
|
||||
return self.dropout(x) + self.dropout(pos_emb)
|
||||
@@ -0,0 +1,63 @@
|
||||
import torch
|
||||
from einops import rearrange, repeat
|
||||
from torch import einsum, Tensor
|
||||
from torch.nn import Module
|
||||
|
||||
|
||||
def rotate_half(x: Tensor, interleaved=True) -> Tensor:
|
||||
if not interleaved:
|
||||
# x_half1, x_half2 = x.chunk(2, dim=-1)
|
||||
# Using torch.split instead of chunk for ONNX export compatibility.
|
||||
x1, x2 = torch.split(x, x.size(-1) // 2, dim=-1)
|
||||
return torch.cat((-x2, x1), dim=-1)
|
||||
else:
|
||||
x = rearrange(x, '... (d r) -> ... d r', r=2)
|
||||
x1, x2 = x.unbind(dim=-1)
|
||||
x = torch.stack((-x2, x1), dim=-1)
|
||||
return rearrange(x, '... d r -> ... (d r)')
|
||||
|
||||
|
||||
def apply_rotary_emb(freqs: Tensor, t: Tensor, interleaved=True) -> Tensor:
|
||||
rot_dim = freqs.shape[-1]
|
||||
t_to_rotate = t[..., :rot_dim]
|
||||
t_pass_through = t[..., rot_dim:]
|
||||
|
||||
t_rotated = (t_to_rotate * freqs.cos()) + (rotate_half(t_to_rotate, interleaved) * freqs.sin())
|
||||
|
||||
return torch.cat((t_rotated, t_pass_through), dim=-1)
|
||||
|
||||
|
||||
class RotaryEmbedding(Module):
|
||||
def __init__(
|
||||
self,
|
||||
dim,
|
||||
theta=10000,
|
||||
max_seq_len=8192,
|
||||
interleaved: bool = True
|
||||
):
|
||||
super().__init__()
|
||||
self.interleaved = interleaved
|
||||
self.cached_freqs_seq_len = max_seq_len
|
||||
inv_freq = 1. / (theta ** (torch.arange(0, dim, 2).float() / dim))
|
||||
self.register_buffer('inv_freq', inv_freq, persistent=False)
|
||||
self.register_buffer('cached_freqs', self._precompute_cache(max_seq_len), persistent=False)
|
||||
|
||||
def _precompute_cache(self, seq_len: int):
|
||||
seq = torch.arange(seq_len, device=self.inv_freq.device, dtype=self.inv_freq.dtype)
|
||||
freqs = einsum('i, j -> i j', seq, self.inv_freq)
|
||||
if self.interleaved:
|
||||
freqs = repeat(freqs, '... n -> ... (n r)', r=2)
|
||||
else:
|
||||
freqs = torch.cat((freqs, freqs), dim=-1)
|
||||
return freqs
|
||||
|
||||
def forward(self, seq_len: int) -> Tensor:
|
||||
if seq_len > self.cached_freqs_seq_len:
|
||||
raise RuntimeError("sequence exceeds RoPE max_seq_len!")
|
||||
return self.cached_freqs[0: seq_len].detach()
|
||||
|
||||
def rotate_queries_or_keys(self, t: Tensor) -> Tensor:
|
||||
device, dtype, seq_len = t.device, t.dtype, t.shape[-2]
|
||||
freqs = self.forward(seq_len=seq_len)
|
||||
|
||||
return apply_rotary_emb(freqs.to(device=device, dtype=dtype), t, self.interleaved)
|
||||
@@ -0,0 +1,24 @@
|
||||
def get_backbone_type(root_config: dict, nested_config: dict = None):
|
||||
if nested_config is None:
|
||||
nested_config = root_config
|
||||
return nested_config.get(
|
||||
'backbone_type',
|
||||
root_config.get(
|
||||
'backbone_type',
|
||||
root_config.get('diff_decoder_type', 'wavenet')
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def get_backbone_args(config: dict, backbone_type: str):
|
||||
args = config.get('backbone_args')
|
||||
if args is not None:
|
||||
return args
|
||||
elif backbone_type == 'wavenet':
|
||||
return {
|
||||
'num_layers': config.get('residual_layers'),
|
||||
'num_channels': config.get('residual_channels'),
|
||||
'dilation_cycle_length': config.get('dilation_cycle_length'),
|
||||
}
|
||||
else:
|
||||
return None
|
||||
@@ -0,0 +1,2 @@
|
||||
from .ddpm import GaussianDiffusion, PitchDiffusion, MultiVarianceDiffusion
|
||||
from .reflow import RectifiedFlow, PitchRectifiedFlow, MultiVarianceRectifiedFlow
|
||||
@@ -0,0 +1,509 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections import deque
|
||||
from functools import partial
|
||||
from typing import List, Tuple
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch import nn
|
||||
from tqdm import tqdm
|
||||
|
||||
from modules.backbones import build_backbone
|
||||
from utils.hparams import hparams
|
||||
|
||||
|
||||
def extract(a, t, x_shape):
|
||||
b, *_ = t.shape
|
||||
out = a.gather(-1, t)
|
||||
return out.reshape(b, *((1,) * (len(x_shape) - 1)))
|
||||
|
||||
|
||||
def noise_like(shape, device, repeat=False):
|
||||
repeat_noise = lambda: torch.randn((1, *shape[1:]), device=device).repeat(shape[0], *((1,) * (len(shape) - 1)))
|
||||
noise = lambda: torch.randn(shape, device=device)
|
||||
return repeat_noise() if repeat else noise()
|
||||
|
||||
|
||||
def linear_beta_schedule(timesteps, max_beta=0.01):
|
||||
"""
|
||||
linear schedule
|
||||
"""
|
||||
betas = np.linspace(1e-4, max_beta, timesteps)
|
||||
return betas
|
||||
|
||||
|
||||
def cosine_beta_schedule(timesteps, s=0.008):
|
||||
"""
|
||||
cosine schedule
|
||||
as proposed in https://openreview.net/forum?id=-NEXDKk8gZ
|
||||
"""
|
||||
steps = timesteps + 1
|
||||
x = np.linspace(0, steps, steps)
|
||||
alphas_cumprod = np.cos(((x / steps) + s) / (1 + s) * np.pi * 0.5) ** 2
|
||||
alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
|
||||
betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
|
||||
return np.clip(betas, a_min=0, a_max=0.999)
|
||||
|
||||
|
||||
beta_schedule = {
|
||||
"cosine": cosine_beta_schedule,
|
||||
"linear": linear_beta_schedule,
|
||||
}
|
||||
|
||||
|
||||
class GaussianDiffusion(nn.Module):
|
||||
def __init__(self, out_dims, num_feats=1, timesteps=1000, k_step=1000,
|
||||
backbone_type=None, backbone_args=None, betas=None,
|
||||
spec_min=None, spec_max=None):
|
||||
super().__init__()
|
||||
self.denoise_fn: nn.Module = build_backbone(out_dims, num_feats, backbone_type, backbone_args)
|
||||
self.out_dims = out_dims
|
||||
self.num_feats = num_feats
|
||||
|
||||
if betas is not None:
|
||||
betas = betas.detach().cpu().numpy() if isinstance(betas, torch.Tensor) else betas
|
||||
else:
|
||||
schedule_args = {}
|
||||
if hparams['schedule_type'] == 'linear':
|
||||
schedule_args['max_beta'] = hparams.get('max_beta', 0.01)
|
||||
betas = beta_schedule[hparams['schedule_type']](timesteps, **schedule_args)
|
||||
|
||||
alphas = 1. - betas
|
||||
alphas_cumprod = np.cumprod(alphas, axis=0)
|
||||
alphas_cumprod_prev = np.append(1., alphas_cumprod[:-1])
|
||||
|
||||
self.use_shallow_diffusion = hparams.get('use_shallow_diffusion', False)
|
||||
if self.use_shallow_diffusion:
|
||||
assert k_step <= timesteps, 'K_step should not be larger than timesteps.'
|
||||
self.timesteps = timesteps
|
||||
self.k_step = k_step if self.use_shallow_diffusion else timesteps
|
||||
self.noise_list = deque(maxlen=4)
|
||||
|
||||
to_torch = partial(torch.tensor, dtype=torch.float32)
|
||||
|
||||
self.register_buffer('betas', to_torch(betas))
|
||||
self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod))
|
||||
self.register_buffer('alphas_cumprod_prev', to_torch(alphas_cumprod_prev))
|
||||
|
||||
# calculations for diffusion q(x_t | x_{t-1}) and others
|
||||
self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod)))
|
||||
self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod)))
|
||||
self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod)))
|
||||
self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod)))
|
||||
self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod - 1)))
|
||||
|
||||
# calculations for posterior q(x_{t-1} | x_t, x_0)
|
||||
posterior_variance = betas * (1. - alphas_cumprod_prev) / (1. - alphas_cumprod)
|
||||
# above: equal to 1. / (1. / (1. - alpha_cumprod_tm1) + alpha_t / beta_t)
|
||||
self.register_buffer('posterior_variance', to_torch(posterior_variance))
|
||||
# below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain
|
||||
self.register_buffer('posterior_log_variance_clipped', to_torch(np.log(np.maximum(posterior_variance, 1e-20))))
|
||||
self.register_buffer('posterior_mean_coef1', to_torch(
|
||||
betas * np.sqrt(alphas_cumprod_prev) / (1. - alphas_cumprod)))
|
||||
self.register_buffer('posterior_mean_coef2', to_torch(
|
||||
(1. - alphas_cumprod_prev) * np.sqrt(alphas) / (1. - alphas_cumprod)))
|
||||
|
||||
# spec: [B, T, M] or [B, F, T, M]
|
||||
# spec_min and spec_max: [1, 1, M] or [1, 1, F, M] => transpose(-3, -2) => [1, 1, M] or [1, F, 1, M]
|
||||
spec_min = torch.FloatTensor(spec_min)[None, None, :out_dims].transpose(-3, -2)
|
||||
spec_max = torch.FloatTensor(spec_max)[None, None, :out_dims].transpose(-3, -2)
|
||||
self.register_buffer('spec_min', spec_min)
|
||||
self.register_buffer('spec_max', spec_max)
|
||||
|
||||
# for compatibility with ONNX continuous acceleration
|
||||
self.time_scale_factor = self.timesteps
|
||||
self.t_start = 1 - self.k_step / self.timesteps
|
||||
factors = torch.LongTensor([i for i in range(1, self.timesteps + 1) if self.timesteps % i == 0])
|
||||
self.register_buffer('timestep_factors', factors, persistent=False)
|
||||
|
||||
def q_mean_variance(self, x_start, t):
|
||||
mean = extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
|
||||
variance = extract(1. - self.alphas_cumprod, t, x_start.shape)
|
||||
log_variance = extract(self.log_one_minus_alphas_cumprod, t, x_start.shape)
|
||||
return mean, variance, log_variance
|
||||
|
||||
def predict_start_from_noise(self, x_t, t, noise):
|
||||
return (
|
||||
extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t -
|
||||
extract(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * noise
|
||||
)
|
||||
|
||||
def q_posterior(self, x_start, x_t, t):
|
||||
posterior_mean = (
|
||||
extract(self.posterior_mean_coef1, t, x_t.shape) * x_start +
|
||||
extract(self.posterior_mean_coef2, t, x_t.shape) * x_t
|
||||
)
|
||||
posterior_variance = extract(self.posterior_variance, t, x_t.shape)
|
||||
posterior_log_variance_clipped = extract(self.posterior_log_variance_clipped, t, x_t.shape)
|
||||
return posterior_mean, posterior_variance, posterior_log_variance_clipped
|
||||
|
||||
def p_mean_variance(self, x, t, cond):
|
||||
noise_pred = self.denoise_fn(x, t, cond=cond)
|
||||
x_recon = self.predict_start_from_noise(x, t=t, noise=noise_pred)
|
||||
|
||||
# This is previously inherited from original DiffSinger repository
|
||||
# and disabled due to some loudness issues when speedup = 1.
|
||||
# x_recon.clamp_(-1., 1.)
|
||||
|
||||
model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start=x_recon, x_t=x, t=t)
|
||||
return model_mean, posterior_variance, posterior_log_variance
|
||||
|
||||
@torch.no_grad()
|
||||
def p_sample(self, x, t, cond, clip_denoised=True, repeat_noise=False):
|
||||
b, *_, device = *x.shape, x.device
|
||||
model_mean, _, model_log_variance = self.p_mean_variance(x=x, t=t, cond=cond)
|
||||
noise = noise_like(x.shape, device, repeat_noise)
|
||||
# no noise when t == 0
|
||||
nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1)))
|
||||
return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise
|
||||
|
||||
@torch.no_grad()
|
||||
def p_sample_ddim(self, x, t, interval, cond):
|
||||
a_t = extract(self.alphas_cumprod, t, x.shape)
|
||||
a_prev = extract(self.alphas_cumprod, torch.max(t - interval, torch.zeros_like(t)), x.shape)
|
||||
|
||||
noise_pred = self.denoise_fn(x, t, cond=cond)
|
||||
x_prev = a_prev.sqrt() * (
|
||||
x / a_t.sqrt() + (((1 - a_prev) / a_prev).sqrt() - ((1 - a_t) / a_t).sqrt()) * noise_pred
|
||||
)
|
||||
return x_prev
|
||||
|
||||
@torch.no_grad()
|
||||
def p_sample_plms(self, x, t, interval, cond, clip_denoised=True, repeat_noise=False):
|
||||
"""
|
||||
Use the PLMS method from
|
||||
[Pseudo Numerical Methods for Diffusion Models on Manifolds](https://arxiv.org/abs/2202.09778).
|
||||
"""
|
||||
|
||||
def get_x_pred(x, noise_t, t):
|
||||
a_t = extract(self.alphas_cumprod, t, x.shape)
|
||||
a_prev = extract(self.alphas_cumprod, torch.max(t - interval, torch.zeros_like(t)), x.shape)
|
||||
a_t_sq, a_prev_sq = a_t.sqrt(), a_prev.sqrt()
|
||||
|
||||
x_delta = (a_prev - a_t) * ((1 / (a_t_sq * (a_t_sq + a_prev_sq))) * x - 1 / (
|
||||
a_t_sq * (((1 - a_prev) * a_t).sqrt() + ((1 - a_t) * a_prev).sqrt())) * noise_t)
|
||||
x_pred = x + x_delta
|
||||
|
||||
return x_pred
|
||||
|
||||
noise_list = self.noise_list
|
||||
noise_pred = self.denoise_fn(x, t, cond=cond)
|
||||
|
||||
if len(noise_list) == 0:
|
||||
x_pred = get_x_pred(x, noise_pred, t)
|
||||
noise_pred_prev = self.denoise_fn(x_pred, max(t - interval, 0), cond=cond)
|
||||
noise_pred_prime = (noise_pred + noise_pred_prev) / 2
|
||||
elif len(noise_list) == 1:
|
||||
noise_pred_prime = (3 * noise_pred - noise_list[-1]) / 2
|
||||
elif len(noise_list) == 2:
|
||||
noise_pred_prime = (23 * noise_pred - 16 * noise_list[-1] + 5 * noise_list[-2]) / 12
|
||||
else:
|
||||
noise_pred_prime = (55 * noise_pred - 59 * noise_list[-1] + 37 * noise_list[-2] - 9 * noise_list[-3]) / 24
|
||||
|
||||
x_prev = get_x_pred(x, noise_pred_prime, t)
|
||||
noise_list.append(noise_pred)
|
||||
|
||||
return x_prev
|
||||
|
||||
def q_sample(self, x_start, t, noise):
|
||||
return (
|
||||
extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start +
|
||||
extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise
|
||||
)
|
||||
|
||||
def p_losses(self, x_start, t, cond, noise=None):
|
||||
if noise is None:
|
||||
noise = torch.randn_like(x_start)
|
||||
|
||||
x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
|
||||
x_recon = self.denoise_fn(x_noisy, t, cond)
|
||||
|
||||
return x_recon, noise
|
||||
|
||||
def inference(self, cond, b=1, x_start=None, device=None):
|
||||
depth = hparams.get('K_step_infer', self.k_step)
|
||||
speedup = hparams['diff_speedup']
|
||||
if speedup > 0:
|
||||
assert depth % speedup == 0, f'Acceleration ratio must be a factor of diffusion depth {depth}.'
|
||||
|
||||
noise = torch.randn(b, self.num_feats, self.out_dims, cond.shape[2], device=device)
|
||||
if self.use_shallow_diffusion:
|
||||
t_max = min(depth, self.k_step)
|
||||
else:
|
||||
t_max = self.k_step
|
||||
|
||||
if t_max >= self.timesteps:
|
||||
x = noise
|
||||
elif t_max > 0:
|
||||
assert x_start is not None, 'Missing shallow diffusion source.'
|
||||
x = self.q_sample(
|
||||
x_start, torch.full((b,), t_max - 1, device=device, dtype=torch.long), noise
|
||||
)
|
||||
else:
|
||||
assert x_start is not None, 'Missing shallow diffusion source.'
|
||||
x = x_start
|
||||
|
||||
if speedup > 1 and t_max > 0:
|
||||
algorithm = hparams['diff_accelerator']
|
||||
if algorithm == 'dpm-solver':
|
||||
from inference.dpm_solver_pytorch import NoiseScheduleVP, model_wrapper, DPM_Solver
|
||||
# 1. Define the noise schedule.
|
||||
noise_schedule = NoiseScheduleVP(schedule='discrete', betas=self.betas[:t_max])
|
||||
|
||||
# 2. Convert your discrete-time `model` to the continuous-time
|
||||
# noise prediction model. Here is an example for a diffusion model
|
||||
# `model` with the noise prediction type ("noise") .
|
||||
def my_wrapper(fn):
|
||||
def wrapped(x, t, **kwargs):
|
||||
ret = fn(x, t, **kwargs)
|
||||
self.bar.update(1)
|
||||
return ret
|
||||
|
||||
return wrapped
|
||||
|
||||
model_fn = model_wrapper(
|
||||
my_wrapper(self.denoise_fn),
|
||||
noise_schedule,
|
||||
model_type="noise", # or "x_start" or "v" or "score"
|
||||
model_kwargs={"cond": cond}
|
||||
)
|
||||
|
||||
# 3. Define dpm-solver and sample by singlestep DPM-Solver.
|
||||
# (We recommend singlestep DPM-Solver for unconditional sampling)
|
||||
# You can adjust the `steps` to balance the computation
|
||||
# costs and the sample quality.
|
||||
dpm_solver = DPM_Solver(model_fn, noise_schedule, algorithm_type="dpmsolver++")
|
||||
|
||||
steps = t_max // hparams["diff_speedup"]
|
||||
self.bar = tqdm(desc="sample time step", total=steps, disable=not hparams['infer'], leave=False)
|
||||
x = dpm_solver.sample(
|
||||
x,
|
||||
steps=steps,
|
||||
order=2,
|
||||
skip_type="time_uniform",
|
||||
method="multistep",
|
||||
)
|
||||
self.bar.close()
|
||||
elif algorithm == 'unipc':
|
||||
from inference.uni_pc import NoiseScheduleVP, model_wrapper, UniPC
|
||||
# 1. Define the noise schedule.
|
||||
noise_schedule = NoiseScheduleVP(schedule='discrete', betas=self.betas[:t_max])
|
||||
|
||||
# 2. Convert your discrete-time `model` to the continuous-time
|
||||
# noise prediction model. Here is an example for a diffusion model
|
||||
# `model` with the noise prediction type ("noise") .
|
||||
def my_wrapper(fn):
|
||||
def wrapped(x, t, **kwargs):
|
||||
ret = fn(x, t, **kwargs)
|
||||
self.bar.update(1)
|
||||
return ret
|
||||
|
||||
return wrapped
|
||||
|
||||
model_fn = model_wrapper(
|
||||
my_wrapper(self.denoise_fn),
|
||||
noise_schedule,
|
||||
model_type="noise", # or "x_start" or "v" or "score"
|
||||
model_kwargs={"cond": cond}
|
||||
)
|
||||
|
||||
# 3. Define uni_pc and sample by multistep UniPC.
|
||||
# You can adjust the `steps` to balance the computation
|
||||
# costs and the sample quality.
|
||||
uni_pc = UniPC(model_fn, noise_schedule, variant='bh2')
|
||||
|
||||
steps = t_max // hparams["diff_speedup"]
|
||||
self.bar = tqdm(desc="sample time step", total=steps, disable=not hparams['infer'], leave=False)
|
||||
x = uni_pc.sample(
|
||||
x,
|
||||
steps=steps,
|
||||
order=2,
|
||||
skip_type="time_uniform",
|
||||
method="multistep",
|
||||
)
|
||||
self.bar.close()
|
||||
elif algorithm == 'pndm':
|
||||
self.noise_list = deque(maxlen=4)
|
||||
iteration_interval = speedup
|
||||
for i in tqdm(
|
||||
reversed(range(0, t_max, iteration_interval)), desc='sample time step',
|
||||
total=t_max // iteration_interval, disable=not hparams['infer'], leave=False
|
||||
):
|
||||
x = self.p_sample_plms(
|
||||
x, torch.full((b,), i, device=device, dtype=torch.long),
|
||||
iteration_interval, cond=cond
|
||||
)
|
||||
elif algorithm == 'ddim':
|
||||
iteration_interval = speedup
|
||||
for i in tqdm(
|
||||
reversed(range(0, t_max, iteration_interval)), desc='sample time step',
|
||||
total=t_max // iteration_interval, disable=not hparams['infer'], leave=False
|
||||
):
|
||||
x = self.p_sample_ddim(
|
||||
x, torch.full((b,), i, device=device, dtype=torch.long),
|
||||
iteration_interval, cond=cond
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unsupported acceleration algorithm for DDPM: {algorithm}.")
|
||||
else:
|
||||
for i in tqdm(reversed(range(0, t_max)), desc='sample time step', total=t_max,
|
||||
disable=not hparams['infer'], leave=False):
|
||||
x = self.p_sample(x, torch.full((b,), i, device=device, dtype=torch.long), cond)
|
||||
x = x.transpose(2, 3).squeeze(1) # [B, F, M, T] => [B, T, M] or [B, F, T, M]
|
||||
return x
|
||||
|
||||
def forward(self, condition, gt_spec=None, src_spec=None, infer=True):
|
||||
"""
|
||||
conditioning diffusion, use fastspeech2 encoder output as the condition
|
||||
"""
|
||||
cond = condition.transpose(1, 2)
|
||||
b, device = condition.shape[0], condition.device
|
||||
|
||||
if not infer:
|
||||
# gt_spec: [B, T, M] or [B, F, T, M]
|
||||
spec = self.norm_spec(gt_spec).transpose(-2, -1) # [B, M, T] or [B, F, M, T]
|
||||
if self.num_feats == 1:
|
||||
spec = spec[:, None, :, :] # [B, F=1, M, T]
|
||||
t = torch.randint(0, self.k_step, (b,), device=device).long()
|
||||
x_recon, noise = self.p_losses(spec, t, cond=cond)
|
||||
return x_recon, noise
|
||||
else:
|
||||
# src_spec: [B, T, M] or [B, F, T, M]
|
||||
if src_spec is not None:
|
||||
spec = self.norm_spec(src_spec).transpose(-2, -1)
|
||||
if self.num_feats == 1:
|
||||
spec = spec[:, None, :, :]
|
||||
else:
|
||||
spec = None
|
||||
x = self.inference(cond, b=b, x_start=spec, device=device)
|
||||
return self.denorm_spec(x)
|
||||
|
||||
def norm_spec(self, x):
|
||||
return (x - self.spec_min) / (self.spec_max - self.spec_min) * 2 - 1
|
||||
|
||||
def denorm_spec(self, x):
|
||||
return (x + 1) / 2 * (self.spec_max - self.spec_min) + self.spec_min
|
||||
|
||||
|
||||
class RepetitiveDiffusion(GaussianDiffusion):
|
||||
def __init__(self, vmin: float | int | list, vmax: float | int | list,
|
||||
repeat_bins: int, timesteps=1000, k_step=1000,
|
||||
backbone_type=None, backbone_args=None,
|
||||
betas=None):
|
||||
assert (isinstance(vmin, (float, int)) and isinstance(vmax, (float, int))) or len(vmin) == len(vmax)
|
||||
num_feats = 1 if isinstance(vmin, (float, int)) else len(vmin)
|
||||
spec_min = [vmin] if num_feats == 1 else [[v] for v in vmin]
|
||||
spec_max = [vmax] if num_feats == 1 else [[v] for v in vmax]
|
||||
self.repeat_bins = repeat_bins
|
||||
super().__init__(
|
||||
out_dims=repeat_bins, num_feats=num_feats,
|
||||
timesteps=timesteps, k_step=k_step,
|
||||
backbone_type=backbone_type, backbone_args=backbone_args,
|
||||
betas=betas, spec_min=spec_min, spec_max=spec_max
|
||||
)
|
||||
|
||||
def norm_spec(self, x):
|
||||
"""
|
||||
|
||||
:param x: [B, T] or [B, F, T]
|
||||
:return [B, T, R] or [B, F, T, R]
|
||||
"""
|
||||
if self.num_feats == 1:
|
||||
repeats = [1, 1, self.repeat_bins]
|
||||
else:
|
||||
repeats = [1, 1, 1, self.repeat_bins]
|
||||
return super().norm_spec(x.unsqueeze(-1).repeat(repeats))
|
||||
|
||||
def denorm_spec(self, x):
|
||||
"""
|
||||
|
||||
:param x: [B, T, R] or [B, F, T, R]
|
||||
:return [B, T] or [B, F, T]
|
||||
"""
|
||||
return super().denorm_spec(x).mean(dim=-1)
|
||||
|
||||
|
||||
class PitchDiffusion(RepetitiveDiffusion):
|
||||
def __init__(self, vmin: float, vmax: float,
|
||||
cmin: float, cmax: float, repeat_bins,
|
||||
timesteps=1000, k_step=1000,
|
||||
backbone_type=None, backbone_args=None,
|
||||
betas=None):
|
||||
self.vmin = vmin # norm min
|
||||
self.vmax = vmax # norm max
|
||||
self.cmin = cmin # clip min
|
||||
self.cmax = cmax # clip max
|
||||
super().__init__(
|
||||
vmin=vmin, vmax=vmax, repeat_bins=repeat_bins,
|
||||
timesteps=timesteps, k_step=k_step,
|
||||
backbone_type=backbone_type, backbone_args=backbone_args,
|
||||
betas=betas
|
||||
)
|
||||
|
||||
def norm_spec(self, x):
|
||||
return super().norm_spec(x.clamp(min=self.cmin, max=self.cmax))
|
||||
|
||||
def denorm_spec(self, x):
|
||||
return super().denorm_spec(x).clamp(min=self.cmin, max=self.cmax)
|
||||
|
||||
|
||||
class MultiVarianceDiffusion(RepetitiveDiffusion):
|
||||
def __init__(
|
||||
self, ranges: List[Tuple[float, float]],
|
||||
clamps: List[Tuple[float | None, float | None] | None],
|
||||
repeat_bins, timesteps=1000, k_step=1000,
|
||||
backbone_type=None, backbone_args=None,
|
||||
betas=None
|
||||
):
|
||||
assert len(ranges) == len(clamps)
|
||||
self.clamps = clamps
|
||||
vmin = [r[0] for r in ranges]
|
||||
vmax = [r[1] for r in ranges]
|
||||
if len(vmin) == 1:
|
||||
vmin = vmin[0]
|
||||
if len(vmax) == 1:
|
||||
vmax = vmax[0]
|
||||
super().__init__(
|
||||
vmin=vmin, vmax=vmax, repeat_bins=repeat_bins,
|
||||
timesteps=timesteps, k_step=k_step,
|
||||
backbone_type=backbone_type, backbone_args=backbone_args,
|
||||
betas=betas
|
||||
)
|
||||
|
||||
def clamp_spec(self, xs: list | tuple):
|
||||
clamped = []
|
||||
for x, c in zip(xs, self.clamps):
|
||||
if c is None:
|
||||
clamped.append(x)
|
||||
continue
|
||||
clamped.append(x.clamp(min=c[0], max=c[1]))
|
||||
return clamped
|
||||
|
||||
def norm_spec(self, xs: list | tuple):
|
||||
"""
|
||||
|
||||
:param xs: sequence of [B, T]
|
||||
:return: [B, F, T] => super().norm_spec(xs) => [B, F, T, R]
|
||||
"""
|
||||
assert len(xs) == self.num_feats
|
||||
clamped = self.clamp_spec(xs)
|
||||
xs = torch.stack(clamped, dim=1) # [B, F, T]
|
||||
if self.num_feats == 1:
|
||||
xs = xs.squeeze(1) # [B, T]
|
||||
return super().norm_spec(xs)
|
||||
|
||||
def denorm_spec(self, xs):
|
||||
"""
|
||||
|
||||
:param xs: [B, T, R] or [B, F, T, R] => super().denorm_spec(xs) => [B, T] or [B, F, T]
|
||||
:return: sequence of [B, T]
|
||||
"""
|
||||
xs = super().denorm_spec(xs)
|
||||
if self.num_feats == 1:
|
||||
xs = [xs]
|
||||
else:
|
||||
xs = xs.unbind(dim=1)
|
||||
assert len(xs) == self.num_feats
|
||||
return self.clamp_spec(xs)
|
||||
|
||||
@@ -0,0 +1,262 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import List, Tuple
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from tqdm import tqdm
|
||||
|
||||
from modules.backbones import build_backbone
|
||||
from utils.hparams import hparams
|
||||
|
||||
|
||||
class RectifiedFlow(nn.Module):
|
||||
def __init__(self, out_dims, num_feats=1, t_start=0., time_scale_factor=1000,
|
||||
backbone_type=None, backbone_args=None,
|
||||
spec_min=None, spec_max=None):
|
||||
super().__init__()
|
||||
self.velocity_fn: nn.Module = build_backbone(out_dims, num_feats, backbone_type, backbone_args)
|
||||
self.out_dims = out_dims
|
||||
self.num_feats = num_feats
|
||||
self.use_shallow_diffusion = hparams.get('use_shallow_diffusion', False)
|
||||
if self.use_shallow_diffusion:
|
||||
assert 0. <= t_start <= 1., 'T_start should be in [0, 1].'
|
||||
else:
|
||||
t_start = 0.
|
||||
self.t_start = t_start
|
||||
self.time_scale_factor = time_scale_factor
|
||||
|
||||
# spec: [B, T, M] or [B, F, T, M]
|
||||
# spec_min and spec_max: [1, 1, M] or [1, 1, F, M] => transpose(-3, -2) => [1, 1, M] or [1, F, 1, M]
|
||||
spec_min = torch.FloatTensor(spec_min)[None, None, :out_dims].transpose(-3, -2)
|
||||
spec_max = torch.FloatTensor(spec_max)[None, None, :out_dims].transpose(-3, -2)
|
||||
self.register_buffer('spec_min', spec_min, persistent=False)
|
||||
self.register_buffer('spec_max', spec_max, persistent=False)
|
||||
|
||||
def p_losses(self, x_end, t, cond):
|
||||
x_start = torch.randn_like(x_end)
|
||||
x_t = x_start + t[:, None, None, None] * (x_end - x_start)
|
||||
v_pred = self.velocity_fn(x_t, t * self.time_scale_factor, cond)
|
||||
|
||||
return v_pred, x_end - x_start
|
||||
|
||||
def forward(self, condition, gt_spec=None, src_spec=None, infer=True):
|
||||
cond = condition.transpose(1, 2)
|
||||
b, device = condition.shape[0], condition.device
|
||||
|
||||
if not infer:
|
||||
# gt_spec: [B, T, M] or [B, F, T, M]
|
||||
spec = self.norm_spec(gt_spec).transpose(-2, -1) # [B, M, T] or [B, F, M, T]
|
||||
if self.num_feats == 1:
|
||||
spec = spec[:, None, :, :] # [B, F=1, M, T]
|
||||
t = self.t_start + (1.0 - self.t_start) * torch.rand((b,), device=device)
|
||||
v_pred, v_gt = self.p_losses(spec, t, cond=cond)
|
||||
return v_pred, v_gt, t
|
||||
else:
|
||||
# src_spec: [B, T, M] or [B, F, T, M]
|
||||
if src_spec is not None:
|
||||
spec = self.norm_spec(src_spec).transpose(-2, -1)
|
||||
if self.num_feats == 1:
|
||||
spec = spec[:, None, :, :]
|
||||
else:
|
||||
spec = None
|
||||
x = self.inference(cond, b=b, x_end=spec, device=device)
|
||||
return self.denorm_spec(x)
|
||||
|
||||
@torch.no_grad()
|
||||
def sample_euler(self, x, t, dt, cond):
|
||||
x += self.velocity_fn(x, self.time_scale_factor * t, cond) * dt
|
||||
t += dt
|
||||
return x, t
|
||||
|
||||
@torch.no_grad()
|
||||
def sample_rk2(self, x, t, dt, cond):
|
||||
k_1 = self.velocity_fn(x, self.time_scale_factor * t, cond)
|
||||
k_2 = self.velocity_fn(x + 0.5 * k_1 * dt, self.time_scale_factor * (t + 0.5 * dt), cond)
|
||||
x += k_2 * dt
|
||||
t += dt
|
||||
return x, t
|
||||
|
||||
@torch.no_grad()
|
||||
def sample_rk4(self, x, t, dt, cond):
|
||||
k_1 = self.velocity_fn(x, self.time_scale_factor * t, cond)
|
||||
k_2 = self.velocity_fn(x + 0.5 * k_1 * dt, self.time_scale_factor * (t + 0.5 * dt), cond)
|
||||
k_3 = self.velocity_fn(x + 0.5 * k_2 * dt, self.time_scale_factor * (t + 0.5 * dt), cond)
|
||||
k_4 = self.velocity_fn(x + k_3 * dt, self.time_scale_factor * (t + dt), cond)
|
||||
x += (k_1 + 2 * k_2 + 2 * k_3 + k_4) * dt / 6
|
||||
t += dt
|
||||
return x, t
|
||||
|
||||
@torch.no_grad()
|
||||
def sample_rk5(self, x, t, dt, cond):
|
||||
k_1 = self.velocity_fn(x, self.time_scale_factor * t, cond)
|
||||
k_2 = self.velocity_fn(x + 0.25 * k_1 * dt, self.time_scale_factor * (t + 0.25 * dt), cond)
|
||||
k_3 = self.velocity_fn(x + 0.125 * (k_2 + k_1) * dt, self.time_scale_factor * (t + 0.25 * dt), cond)
|
||||
k_4 = self.velocity_fn(x + 0.5 * (-k_2 + 2 * k_3) * dt, self.time_scale_factor * (t + 0.5 * dt), cond)
|
||||
k_5 = self.velocity_fn(x + 0.0625 * (3 * k_1 + 9 * k_4) * dt, self.time_scale_factor * (t + 0.75 * dt), cond)
|
||||
k_6 = self.velocity_fn(x + (-3 * k_1 + 2 * k_2 + 12 * k_3 - 12 * k_4 + 8 * k_5) * dt / 7,
|
||||
self.time_scale_factor * (t + dt),
|
||||
cond)
|
||||
x += (7 * k_1 + 32 * k_3 + 12 * k_4 + 32 * k_5 + 7 * k_6) * dt / 90
|
||||
t += dt
|
||||
return x, t
|
||||
|
||||
@torch.no_grad()
|
||||
def inference(self, cond, b=1, x_end=None, device=None):
|
||||
noise = torch.randn(b, self.num_feats, self.out_dims, cond.shape[2], device=device)
|
||||
t_start = hparams.get('T_start_infer', self.t_start)
|
||||
if self.use_shallow_diffusion and t_start > 0:
|
||||
assert x_end is not None, 'Missing shallow diffusion source.'
|
||||
if t_start >= 1.:
|
||||
t_start = 1.
|
||||
x = x_end
|
||||
else:
|
||||
x = t_start * x_end + (1 - t_start) * noise
|
||||
else:
|
||||
t_start = 0.
|
||||
x = noise
|
||||
|
||||
algorithm = hparams['sampling_algorithm']
|
||||
infer_step = hparams['sampling_steps']
|
||||
|
||||
if t_start < 1:
|
||||
dt = (1.0 - t_start) / max(1, infer_step)
|
||||
algorithm_fn = {
|
||||
'euler': self.sample_euler,
|
||||
'rk2': self.sample_rk2,
|
||||
'rk4': self.sample_rk4,
|
||||
'rk5': self.sample_rk5,
|
||||
}.get(algorithm)
|
||||
if algorithm_fn is None:
|
||||
raise ValueError(f'Unsupported algorithm for Rectified Flow: {algorithm}.')
|
||||
dts = torch.tensor([dt]).to(x)
|
||||
for i in tqdm(range(infer_step), desc='sample time step', total=infer_step,
|
||||
disable=not hparams['infer'], leave=False):
|
||||
x, _ = algorithm_fn(x, t_start + i * dts, dt, cond)
|
||||
x = x.float()
|
||||
x = x.transpose(2, 3).squeeze(1) # [B, F, M, T] => [B, T, M] or [B, F, T, M]
|
||||
return x
|
||||
|
||||
def norm_spec(self, x):
|
||||
return (x - self.spec_min) / (self.spec_max - self.spec_min) * 2 - 1
|
||||
|
||||
def denorm_spec(self, x):
|
||||
return (x + 1) / 2 * (self.spec_max - self.spec_min) + self.spec_min
|
||||
|
||||
|
||||
class RepetitiveRectifiedFlow(RectifiedFlow):
|
||||
def __init__(self, vmin: float | int | list, vmax: float | int | list,
|
||||
repeat_bins: int, time_scale_factor=1000,
|
||||
backbone_type=None, backbone_args=None):
|
||||
assert (isinstance(vmin, (float, int)) and isinstance(vmax, (float, int))) or len(vmin) == len(vmax)
|
||||
num_feats = 1 if isinstance(vmin, (float, int)) else len(vmin)
|
||||
spec_min = [vmin] if num_feats == 1 else [[v] for v in vmin]
|
||||
spec_max = [vmax] if num_feats == 1 else [[v] for v in vmax]
|
||||
self.repeat_bins = repeat_bins
|
||||
super().__init__(
|
||||
out_dims=repeat_bins, num_feats=num_feats,
|
||||
time_scale_factor=time_scale_factor,
|
||||
backbone_type=backbone_type, backbone_args=backbone_args,
|
||||
spec_min=spec_min, spec_max=spec_max
|
||||
)
|
||||
|
||||
def norm_spec(self, x):
|
||||
"""
|
||||
|
||||
:param x: [B, T] or [B, F, T]
|
||||
:return [B, T, R] or [B, F, T, R]
|
||||
"""
|
||||
if self.num_feats == 1:
|
||||
repeats = [1, 1, self.repeat_bins]
|
||||
else:
|
||||
repeats = [1, 1, 1, self.repeat_bins]
|
||||
return super().norm_spec(x.unsqueeze(-1).repeat(repeats))
|
||||
|
||||
def denorm_spec(self, x):
|
||||
"""
|
||||
|
||||
:param x: [B, T, R] or [B, F, T, R]
|
||||
:return [B, T] or [B, F, T]
|
||||
"""
|
||||
return super().denorm_spec(x).mean(dim=-1)
|
||||
|
||||
|
||||
class PitchRectifiedFlow(RepetitiveRectifiedFlow):
|
||||
def __init__(self, vmin: float, vmax: float,
|
||||
cmin: float, cmax: float, repeat_bins,
|
||||
time_scale_factor=1000,
|
||||
backbone_type=None, backbone_args=None):
|
||||
self.vmin = vmin # norm min
|
||||
self.vmax = vmax # norm max
|
||||
self.cmin = cmin # clip min
|
||||
self.cmax = cmax # clip max
|
||||
super().__init__(
|
||||
vmin=vmin, vmax=vmax, repeat_bins=repeat_bins,
|
||||
time_scale_factor=time_scale_factor,
|
||||
backbone_type=backbone_type, backbone_args=backbone_args
|
||||
)
|
||||
|
||||
def norm_spec(self, x):
|
||||
return super().norm_spec(x.clamp(min=self.cmin, max=self.cmax))
|
||||
|
||||
def denorm_spec(self, x):
|
||||
return super().denorm_spec(x).clamp(min=self.cmin, max=self.cmax)
|
||||
|
||||
|
||||
class MultiVarianceRectifiedFlow(RepetitiveRectifiedFlow):
|
||||
def __init__(
|
||||
self, ranges: List[Tuple[float, float]],
|
||||
clamps: List[Tuple[float | None, float | None] | None],
|
||||
repeat_bins, time_scale_factor=1000,
|
||||
backbone_type=None, backbone_args=None
|
||||
):
|
||||
assert len(ranges) == len(clamps)
|
||||
self.clamps = clamps
|
||||
vmin = [r[0] for r in ranges]
|
||||
vmax = [r[1] for r in ranges]
|
||||
if len(vmin) == 1:
|
||||
vmin = vmin[0]
|
||||
if len(vmax) == 1:
|
||||
vmax = vmax[0]
|
||||
super().__init__(
|
||||
vmin=vmin, vmax=vmax, repeat_bins=repeat_bins,
|
||||
time_scale_factor=time_scale_factor,
|
||||
backbone_type=backbone_type, backbone_args=backbone_args
|
||||
)
|
||||
|
||||
def clamp_spec(self, xs: list | tuple):
|
||||
clamped = []
|
||||
for x, c in zip(xs, self.clamps):
|
||||
if c is None:
|
||||
clamped.append(x)
|
||||
continue
|
||||
clamped.append(x.clamp(min=c[0], max=c[1]))
|
||||
return clamped
|
||||
|
||||
def norm_spec(self, xs: list | tuple):
|
||||
"""
|
||||
|
||||
:param xs: sequence of [B, T]
|
||||
:return: [B, F, T] => super().norm_spec(xs) => [B, F, T, R]
|
||||
"""
|
||||
assert len(xs) == self.num_feats
|
||||
clamped = self.clamp_spec(xs)
|
||||
xs = torch.stack(clamped, dim=1) # [B, F, T]
|
||||
if self.num_feats == 1:
|
||||
xs = xs.squeeze(1) # [B, T]
|
||||
return super().norm_spec(xs)
|
||||
|
||||
def denorm_spec(self, xs):
|
||||
"""
|
||||
|
||||
:param xs: [B, T, R] or [B, F, T, R] => super().denorm_spec(xs) => [B, T] or [B, F, T]
|
||||
:return: sequence of [B, T]
|
||||
"""
|
||||
xs = super().denorm_spec(xs)
|
||||
if self.num_feats == 1:
|
||||
xs = [xs]
|
||||
else:
|
||||
xs = xs.unbind(dim=1)
|
||||
assert len(xs) == self.num_feats
|
||||
return self.clamp_spec(xs)
|
||||
|
||||
@@ -0,0 +1,185 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.nn import functional as F
|
||||
|
||||
from modules.commons.common_layers import (
|
||||
NormalInitEmbedding as Embedding,
|
||||
SinusoidalPosEmb,
|
||||
AdamWLinear,
|
||||
)
|
||||
from modules.fastspeech.tts_modules import FastSpeech2Encoder, mel2ph_to_dur, StretchRegulator
|
||||
from utils.hparams import hparams
|
||||
from utils.phoneme_utils import PAD_INDEX
|
||||
|
||||
|
||||
class FastSpeech2Acoustic(nn.Module):
|
||||
def __init__(self, vocab_size):
|
||||
super().__init__()
|
||||
self.txt_embed = Embedding(vocab_size, hparams['hidden_size'], PAD_INDEX)
|
||||
self.use_lang_id = hparams.get('use_lang_id', False)
|
||||
if self.use_lang_id:
|
||||
self.lang_embed = Embedding(hparams['num_lang'] + 1, hparams['hidden_size'], padding_idx=0)
|
||||
|
||||
self.use_stretch_embed = hparams.get('use_stretch_embed', False)
|
||||
if self.use_stretch_embed:
|
||||
self.sr = StretchRegulator()
|
||||
self.stretch_embed = nn.Sequential(
|
||||
SinusoidalPosEmb(hparams['hidden_size']),
|
||||
nn.Linear(hparams['hidden_size'], hparams['hidden_size'] * 4),
|
||||
nn.GELU(),
|
||||
nn.Linear(hparams['hidden_size'] * 4, hparams['hidden_size']),
|
||||
)
|
||||
self.stretch_embed_rnn = nn.GRU(hparams['hidden_size'], hparams['hidden_size'], 1, batch_first=True)
|
||||
self._stretch_embed_rnn_flattened = False
|
||||
|
||||
self.dur_embed = AdamWLinear(1, hparams['hidden_size'])
|
||||
self.use_mix_ln = hparams.get('use_mix_ln', False)
|
||||
if self.use_mix_ln:
|
||||
self.mix_ln_layer = hparams['mix_ln_layer']
|
||||
else:
|
||||
self.mix_ln_layer = []
|
||||
self.encoder = FastSpeech2Encoder(
|
||||
hidden_size=hparams['hidden_size'], num_layers=hparams['enc_layers'],
|
||||
ffn_kernel_size=hparams['enc_ffn_kernel_size'], ffn_act=hparams['ffn_act'],
|
||||
dropout=hparams['dropout'], num_heads=hparams['num_heads'],
|
||||
use_pos_embed=hparams['use_pos_embed'], rel_pos=hparams.get('rel_pos', False),
|
||||
use_rope=hparams.get('use_rope', False), rope_interleaved=hparams.get('rope_interleaved', True),
|
||||
mix_ln_layer=self.mix_ln_layer
|
||||
)
|
||||
|
||||
self.pitch_embed = AdamWLinear(1, hparams['hidden_size'])
|
||||
self.variance_embed_list = []
|
||||
self.use_energy_embed = hparams.get('use_energy_embed', False)
|
||||
self.use_breathiness_embed = hparams.get('use_breathiness_embed', False)
|
||||
self.use_voicing_embed = hparams.get('use_voicing_embed', False)
|
||||
self.use_tension_embed = hparams.get('use_tension_embed', False)
|
||||
if self.use_energy_embed:
|
||||
self.variance_embed_list.append('energy')
|
||||
if self.use_breathiness_embed:
|
||||
self.variance_embed_list.append('breathiness')
|
||||
if self.use_voicing_embed:
|
||||
self.variance_embed_list.append('voicing')
|
||||
if self.use_tension_embed:
|
||||
self.variance_embed_list.append('tension')
|
||||
|
||||
self.use_variance_embeds = len(self.variance_embed_list) > 0
|
||||
if self.use_variance_embeds:
|
||||
self.variance_embeds = nn.ModuleDict({
|
||||
v_name: AdamWLinear(1, hparams['hidden_size'])
|
||||
for v_name in self.variance_embed_list
|
||||
})
|
||||
|
||||
self.use_variance_scaling = hparams.get('use_variance_scaling', False)
|
||||
if self.use_variance_scaling:
|
||||
self.variance_scaling_factor = {
|
||||
'energy': 1. / 96, # 96 dB — max dynamic range of 16-bit audio
|
||||
'breathiness': 1. / 96,
|
||||
'voicing': 1. / 96,
|
||||
'tension': 0.1, # 1 / 10; tension logits are roughly [-10, 10]
|
||||
'key_shift': 1. / 12, # one octave — max key shift in most editors
|
||||
'speed': 1.
|
||||
}
|
||||
else:
|
||||
self.variance_scaling_factor = {
|
||||
'energy': 1.,
|
||||
'breathiness': 1.,
|
||||
'voicing': 1.,
|
||||
'tension': 1.,
|
||||
'key_shift': 1.,
|
||||
'speed': 1.
|
||||
}
|
||||
|
||||
self.use_key_shift_embed = hparams.get('use_key_shift_embed', False)
|
||||
if self.use_key_shift_embed:
|
||||
self.key_shift_embed = AdamWLinear(1, hparams['hidden_size'])
|
||||
|
||||
self.use_speed_embed = hparams.get('use_speed_embed', False)
|
||||
if self.use_speed_embed:
|
||||
self.speed_embed = AdamWLinear(1, hparams['hidden_size'])
|
||||
|
||||
self.use_spk_id = hparams['use_spk_id']
|
||||
if self.use_spk_id:
|
||||
self.spk_embed = Embedding(hparams['num_spk'], hparams['hidden_size'])
|
||||
|
||||
def forward_variance_embedding(self, condition, key_shift=None, speed=None, **variances):
|
||||
if self.use_variance_embeds:
|
||||
variance_embeds = torch.stack([
|
||||
self.variance_embeds[v_name](variances[v_name][:, :, None] * self.variance_scaling_factor[v_name])
|
||||
for v_name in self.variance_embed_list
|
||||
], dim=-1).sum(-1)
|
||||
condition += variance_embeds
|
||||
|
||||
if self.use_key_shift_embed:
|
||||
key_shift_embed = self.key_shift_embed(key_shift[:, :, None] * self.variance_scaling_factor['key_shift'])
|
||||
condition += key_shift_embed
|
||||
|
||||
if self.use_speed_embed:
|
||||
speed_embed = self.speed_embed(speed[:, :, None] * self.variance_scaling_factor['speed'])
|
||||
condition += speed_embed
|
||||
|
||||
return condition
|
||||
|
||||
def forward(
|
||||
self, txt_tokens, mel2ph, f0,
|
||||
key_shift=None, speed=None,
|
||||
spk_embed_id=None, languages=None,
|
||||
**kwargs
|
||||
):
|
||||
spk_embed = None
|
||||
if self.use_spk_id:
|
||||
spk_mix_embed = kwargs.get('spk_mix_embed')
|
||||
if spk_mix_embed is not None:
|
||||
spk_embed = spk_mix_embed
|
||||
else:
|
||||
spk_embed = self.spk_embed(spk_embed_id)[:, None, :]
|
||||
txt_embed = self.txt_embed(txt_tokens)
|
||||
dur = mel2ph_to_dur(mel2ph, txt_tokens.shape[1])
|
||||
if self.use_variance_scaling:
|
||||
dur_embed = self.dur_embed(torch.log(1 + dur[:, :, None].float()))
|
||||
else:
|
||||
dur_embed = self.dur_embed(dur[:, :, None].float())
|
||||
if self.use_lang_id:
|
||||
lang_embed = self.lang_embed(languages)
|
||||
extra_embed = dur_embed + lang_embed
|
||||
else:
|
||||
extra_embed = dur_embed
|
||||
encoder_out = self.encoder(txt_embed, extra_embed, txt_tokens == 0, spk_embed)
|
||||
|
||||
encoder_out = F.pad(encoder_out, [0, 0, 1, 0])
|
||||
mel2ph_ = mel2ph[..., None].repeat([1, 1, encoder_out.shape[-1]])
|
||||
condition = torch.gather(encoder_out, 1, mel2ph_)
|
||||
|
||||
if self.use_stretch_embed:
|
||||
stretch = torch.round(1000 * self.sr(mel2ph, dur))
|
||||
if self.training and stretch.numel() > 1000:
|
||||
# construct a phoneme stretching index lookup table with a total of 1001 indexes (0~1000)
|
||||
table = self.stretch_embed(torch.arange(0, 1001, device=stretch.device))
|
||||
stretch_embed = torch.index_select(table, 0, stretch.view(-1).long()).view_as(condition)
|
||||
else:
|
||||
stretch_embed = self.stretch_embed(stretch)
|
||||
condition += stretch_embed
|
||||
# flatten_parameters fuses the GRU weights into a contiguous buffer for cuDNN.
|
||||
# It only needs to happen once after weight init, device change, or load_state_dict.
|
||||
# We guard with a flag to avoid the redundant call on every forward.
|
||||
# Limitation: the flag lives on this module and is invisible to PyTorch. After
|
||||
# load_state_dict() or model.to(device) replaces the GRU weights, the flag stays
|
||||
# True and flatten_parameters is skipped — cuDNN will fall back to the slower path.
|
||||
# To restore the fast path, reset the flag manually: model._stretch_embed_rnn_flattened = False
|
||||
if not self._stretch_embed_rnn_flattened:
|
||||
self.stretch_embed_rnn.flatten_parameters()
|
||||
self._stretch_embed_rnn_flattened = True
|
||||
stretch_embed_rnn_out, _ = self.stretch_embed_rnn(condition)
|
||||
condition = condition + stretch_embed_rnn_out
|
||||
|
||||
if self.use_spk_id:
|
||||
condition += spk_embed
|
||||
|
||||
f0_mel = (1 + f0 / 700).log()
|
||||
pitch_embed = self.pitch_embed(f0_mel[:, :, None])
|
||||
condition += pitch_embed
|
||||
|
||||
condition = self.forward_variance_embedding(
|
||||
condition, key_shift=key_shift, speed=speed, **kwargs
|
||||
)
|
||||
|
||||
return condition
|
||||
@@ -0,0 +1,95 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
|
||||
import modules.compat as compat
|
||||
from modules.core.ddpm import MultiVarianceDiffusion
|
||||
from utils import filter_kwargs
|
||||
from utils.hparams import hparams
|
||||
|
||||
VARIANCE_CHECKLIST = ['energy', 'breathiness', 'voicing', 'tension']
|
||||
|
||||
|
||||
class ParameterAdaptorModule(torch.nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.variance_prediction_list = []
|
||||
self.predict_energy = hparams.get('predict_energy', False)
|
||||
self.predict_breathiness = hparams.get('predict_breathiness', False)
|
||||
self.predict_voicing = hparams.get('predict_voicing', False)
|
||||
self.predict_tension = hparams.get('predict_tension', False)
|
||||
if self.predict_energy:
|
||||
self.variance_prediction_list.append('energy')
|
||||
if self.predict_breathiness:
|
||||
self.variance_prediction_list.append('breathiness')
|
||||
if self.predict_voicing:
|
||||
self.variance_prediction_list.append('voicing')
|
||||
if self.predict_tension:
|
||||
self.variance_prediction_list.append('tension')
|
||||
self.predict_variances = len(self.variance_prediction_list) > 0
|
||||
|
||||
def build_adaptor(self, cls=MultiVarianceDiffusion):
|
||||
ranges = []
|
||||
clamps = []
|
||||
|
||||
if self.predict_energy:
|
||||
ranges.append((
|
||||
hparams['energy_db_min'],
|
||||
hparams['energy_db_max']
|
||||
))
|
||||
clamps.append((hparams['energy_db_min'], 0.))
|
||||
|
||||
if self.predict_breathiness:
|
||||
ranges.append((
|
||||
hparams['breathiness_db_min'],
|
||||
hparams['breathiness_db_max']
|
||||
))
|
||||
clamps.append((hparams['breathiness_db_min'], 0.))
|
||||
|
||||
if self.predict_voicing:
|
||||
ranges.append((
|
||||
hparams['voicing_db_min'],
|
||||
hparams['voicing_db_max']
|
||||
))
|
||||
clamps.append((hparams['voicing_db_min'], 0.))
|
||||
|
||||
if self.predict_tension:
|
||||
ranges.append((
|
||||
hparams['tension_logit_min'],
|
||||
hparams['tension_logit_max']
|
||||
))
|
||||
clamps.append((
|
||||
hparams['tension_logit_min'],
|
||||
hparams['tension_logit_max']
|
||||
))
|
||||
|
||||
variances_hparams = hparams['variances_prediction_args']
|
||||
total_repeat_bins = variances_hparams['total_repeat_bins']
|
||||
assert total_repeat_bins % len(self.variance_prediction_list) == 0, \
|
||||
f'Total number of repeat bins must be divisible by number of ' \
|
||||
f'variance parameters ({len(self.variance_prediction_list)}).'
|
||||
repeat_bins = total_repeat_bins // len(self.variance_prediction_list)
|
||||
backbone_type = compat.get_backbone_type(hparams, nested_config=variances_hparams)
|
||||
backbone_args = compat.get_backbone_args(variances_hparams, backbone_type=backbone_type)
|
||||
kwargs = filter_kwargs(
|
||||
{
|
||||
'ranges': ranges,
|
||||
'clamps': clamps,
|
||||
'repeat_bins': repeat_bins,
|
||||
'timesteps': hparams.get('timesteps'),
|
||||
'time_scale_factor': hparams.get('time_scale_factor'),
|
||||
'backbone_type': backbone_type,
|
||||
'backbone_args': backbone_args
|
||||
},
|
||||
cls
|
||||
)
|
||||
return cls(**kwargs)
|
||||
|
||||
def collect_variance_inputs(self, **kwargs) -> list:
|
||||
return [kwargs.get(name) for name in self.variance_prediction_list]
|
||||
|
||||
def collect_variance_outputs(self, variances: list | tuple) -> dict:
|
||||
return {
|
||||
name: pred
|
||||
for name, pred in zip(self.variance_prediction_list, variances)
|
||||
}
|
||||
@@ -0,0 +1,455 @@
|
||||
import math
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.nn import functional as F
|
||||
from modules.commons.rotary_embedding_torch import RotaryEmbedding
|
||||
from modules.commons.common_layers import SinusoidalPositionalEmbedding, EncSALayer, AdamWLinear
|
||||
from modules.commons.espnet_positional_embedding import RelPositionalEncoding
|
||||
|
||||
DEFAULT_MAX_SOURCE_POSITIONS = 2000
|
||||
DEFAULT_MAX_TARGET_POSITIONS = 2000
|
||||
|
||||
|
||||
class TransformerEncoderLayer(nn.Module):
|
||||
def __init__(self, hidden_size, dropout, kernel_size=None, act='gelu', num_heads=2, rotary_embed=None,
|
||||
layer_idx=None, mix_ln_layer=None):
|
||||
super().__init__()
|
||||
self.op = EncSALayer(
|
||||
hidden_size, num_heads, dropout=dropout,
|
||||
attention_dropout=0.0, relu_dropout=dropout,
|
||||
kernel_size=kernel_size,
|
||||
act=act, rotary_embed=rotary_embed,
|
||||
layer_idx=layer_idx, mix_ln_layer=mix_ln_layer
|
||||
)
|
||||
|
||||
def forward(self, x, **kwargs):
|
||||
return self.op(x, **kwargs)
|
||||
|
||||
|
||||
######################
|
||||
# fastspeech modules
|
||||
######################
|
||||
class LayerNorm(torch.nn.LayerNorm):
|
||||
"""Layer normalization module.
|
||||
:param int nout: output dim size
|
||||
:param int dim: dimension to be normalized
|
||||
"""
|
||||
|
||||
def __init__(self, nout, dim=-1):
|
||||
"""Construct an LayerNorm object."""
|
||||
super(LayerNorm, self).__init__(nout, eps=1e-12)
|
||||
self.dim = dim
|
||||
|
||||
def forward(self, x):
|
||||
"""Apply layer normalization.
|
||||
:param torch.Tensor x: input tensor
|
||||
:return: layer normalized tensor
|
||||
:rtype torch.Tensor
|
||||
"""
|
||||
if self.dim == -1:
|
||||
return super(LayerNorm, self).forward(x)
|
||||
return super(LayerNorm, self).forward(x.transpose(1, -1)).transpose(1, -1)
|
||||
|
||||
|
||||
class DurationPredictor(torch.nn.Module):
|
||||
"""Duration predictor module.
|
||||
This is a module of duration predictor described in `FastSpeech: Fast, Robust and Controllable Text to Speech`_.
|
||||
The duration predictor predicts a duration of each frame in log domain from the hidden embeddings of encoder.
|
||||
.. _`FastSpeech: Fast, Robust and Controllable Text to Speech`:
|
||||
https://arxiv.org/pdf/1905.09263.pdf
|
||||
Note:
|
||||
The calculation domain of outputs is different between in `forward` and in `inference`. In `forward`,
|
||||
the outputs are calculated in log domain but in `inference`, those are calculated in linear domain.
|
||||
"""
|
||||
|
||||
def __init__(self, in_dims, n_layers=2, n_chans=384, kernel_size=3,
|
||||
dropout_rate=0.1, offset=1.0, dur_loss_type='mse', arch='resnet'):
|
||||
"""Initialize duration predictor module.
|
||||
Args:
|
||||
in_dims (int): Input dimension.
|
||||
n_layers (int, optional): Number of convolutional layers.
|
||||
n_chans (int, optional): Number of channels of convolutional layers.
|
||||
kernel_size (int, optional): Kernel size of convolutional layers.
|
||||
dropout_rate (float, optional): Dropout rate.
|
||||
offset (float, optional): Offset value to avoid nan in log domain.
|
||||
"""
|
||||
super(DurationPredictor, self).__init__()
|
||||
self.offset = offset
|
||||
self.conv = torch.nn.ModuleList()
|
||||
self.kernel_size = kernel_size
|
||||
self.use_resnet = (arch == 'resnet')
|
||||
for idx in range(n_layers):
|
||||
in_chans = in_dims if idx == 0 else n_chans
|
||||
if self.use_resnet:
|
||||
self.conv.append(nn.Sequential(
|
||||
LayerNorm(in_chans, dim=1),
|
||||
nn.Conv1d(in_chans, n_chans, kernel_size, stride=1, padding=kernel_size // 2),
|
||||
nn.ReLU(),
|
||||
nn.Conv1d(n_chans, n_chans, 1),
|
||||
nn.Dropout(dropout_rate)
|
||||
))
|
||||
else:
|
||||
self.conv.append(nn.Sequential(
|
||||
nn.Identity(), # this is a placeholder for ConstantPad1d which is now merged into Conv1d
|
||||
nn.Conv1d(in_chans, n_chans, kernel_size, stride=1, padding=kernel_size // 2),
|
||||
nn.ReLU(),
|
||||
LayerNorm(n_chans, dim=1),
|
||||
nn.Dropout(dropout_rate)
|
||||
))
|
||||
if self.use_resnet and in_dims != n_chans:
|
||||
self.res_conv = nn.Conv1d(in_dims, n_chans, 1)
|
||||
else:
|
||||
self.res_conv = None
|
||||
self.loss_type = dur_loss_type
|
||||
if self.loss_type in ['mse', 'huber']:
|
||||
self.out_dims = 1
|
||||
# elif hparams['dur_loss_type'] == 'mog':
|
||||
# out_dims = 15
|
||||
# elif hparams['dur_loss_type'] == 'crf':
|
||||
# out_dims = 32
|
||||
# from torchcrf import CRF
|
||||
# self.crf = CRF(out_dims, batch_first=True)
|
||||
else:
|
||||
raise NotImplementedError()
|
||||
self.linear = AdamWLinear(n_chans, self.out_dims)
|
||||
|
||||
def out2dur(self, xs):
|
||||
if self.loss_type in ['mse', 'huber']:
|
||||
# NOTE: calculate loss in log domain
|
||||
dur = xs.squeeze(-1).exp() - self.offset # (B, Tmax)
|
||||
# elif hparams['dur_loss_type'] == 'crf':
|
||||
# dur = torch.LongTensor(self.crf.decode(xs)).cuda()
|
||||
else:
|
||||
raise NotImplementedError()
|
||||
return dur
|
||||
|
||||
def forward(self, xs, x_masks=None, infer=True):
|
||||
"""Calculate forward propagation.
|
||||
Args:
|
||||
xs (Tensor): Batch of input sequences (B, Tmax, idim).
|
||||
x_masks (BoolTensor, optional): Batch of masks indicating padded part (B, Tmax).
|
||||
infer (bool): Whether inference
|
||||
Returns:
|
||||
(train) FloatTensor, (infer) LongTensor: Batch of predicted durations in linear domain (B, Tmax).
|
||||
"""
|
||||
xs = xs.transpose(1, -1) # (B, idim, Tmax)
|
||||
masks = 1 - x_masks.float()
|
||||
masks_ = masks[:, None, :]
|
||||
for idx, f in enumerate(self.conv):
|
||||
if self.use_resnet:
|
||||
residual = self.res_conv(xs) if idx == 0 and self.res_conv is not None else xs
|
||||
xs = residual + f(xs)
|
||||
else:
|
||||
xs = f(xs)
|
||||
if x_masks is not None:
|
||||
xs = xs * masks_
|
||||
xs = self.linear(xs.transpose(1, -1)) # [B, T, C]
|
||||
xs = xs * masks[:, :, None] # (B, T, C)
|
||||
|
||||
dur_pred = self.out2dur(xs)
|
||||
if infer:
|
||||
dur_pred = dur_pred.clamp(min=0.) # avoid negative value
|
||||
return dur_pred
|
||||
|
||||
|
||||
class VariancePredictor(torch.nn.Module):
|
||||
def __init__(self, vmin, vmax, in_dims,
|
||||
n_layers=5, n_chans=512, kernel_size=5,
|
||||
dropout_rate=0.1):
|
||||
"""Initialize variance predictor module.
|
||||
Args:
|
||||
in_dims (int): Input dimension.
|
||||
n_layers (int, optional): Number of convolutional layers.
|
||||
n_chans (int, optional): Number of channels of convolutional layers.
|
||||
kernel_size (int, optional): Kernel size of convolutional layers.
|
||||
dropout_rate (float, optional): Dropout rate.
|
||||
"""
|
||||
super(VariancePredictor, self).__init__()
|
||||
|
||||
self.vmin = vmin
|
||||
self.vmax = vmax
|
||||
self.conv = torch.nn.ModuleList()
|
||||
self.kernel_size = kernel_size
|
||||
for idx in range(n_layers):
|
||||
in_chans = in_dims if idx == 0 else n_chans
|
||||
self.conv.append(torch.nn.Sequential(
|
||||
torch.nn.Conv1d(in_chans, n_chans, kernel_size, stride=1, padding=kernel_size // 2),
|
||||
torch.nn.ReLU(),
|
||||
LayerNorm(n_chans, dim=1),
|
||||
torch.nn.Dropout(dropout_rate)
|
||||
))
|
||||
self.linear = torch.nn.Linear(n_chans, 1)
|
||||
self.embed_positions = SinusoidalPositionalEmbedding(in_dims, 0, init_size=4096)
|
||||
self.pos_embed_alpha = nn.Parameter(torch.Tensor([1]))
|
||||
|
||||
def out2value(self, xs):
|
||||
return (xs + 1) / 2 * (self.vmax - self.vmin) + self.vmin
|
||||
|
||||
def forward(self, xs, infer=True):
|
||||
"""
|
||||
:param xs: [B, T, H]
|
||||
:param infer: whether inference
|
||||
:return: [B, T]
|
||||
"""
|
||||
positions = self.pos_embed_alpha * self.embed_positions(xs[..., 0])
|
||||
xs = xs + positions
|
||||
xs = xs.transpose(1, -1) # (B, idim, Tmax)
|
||||
for f in self.conv:
|
||||
xs = f(xs) # (B, C, Tmax)
|
||||
xs = self.linear(xs.transpose(1, -1)).squeeze(-1) # (B, Tmax)
|
||||
if infer:
|
||||
xs = self.out2value(xs)
|
||||
return xs
|
||||
|
||||
|
||||
class PitchPredictor(torch.nn.Module):
|
||||
def __init__(self, vmin, vmax, num_bins, deviation,
|
||||
in_dims, n_layers=5, n_chans=384, kernel_size=5,
|
||||
dropout_rate=0.1):
|
||||
"""Initialize pitch predictor module.
|
||||
Args:
|
||||
in_dims (int): Input dimension.
|
||||
n_layers (int, optional): Number of convolutional layers.
|
||||
n_chans (int, optional): Number of channels of convolutional layers.
|
||||
kernel_size (int, optional): Kernel size of convolutional layers.
|
||||
dropout_rate (float, optional): Dropout rate.
|
||||
"""
|
||||
super(PitchPredictor, self).__init__()
|
||||
self.vmin = vmin
|
||||
self.vmax = vmax
|
||||
self.interval = (vmax - vmin) / (num_bins - 1) # align with centers of bins
|
||||
self.sigma = deviation / self.interval
|
||||
self.register_buffer('x', torch.arange(num_bins).float().reshape(1, 1, -1)) # [1, 1, N]
|
||||
|
||||
self.base_pitch_embed = torch.nn.Linear(1, in_dims)
|
||||
self.conv = torch.nn.ModuleList()
|
||||
self.kernel_size = kernel_size
|
||||
for idx in range(n_layers):
|
||||
in_chans = in_dims if idx == 0 else n_chans
|
||||
self.conv.append(torch.nn.Sequential(
|
||||
torch.nn.Conv1d(in_chans, n_chans, kernel_size, stride=1, padding=kernel_size // 2),
|
||||
torch.nn.ReLU(),
|
||||
LayerNorm(n_chans, dim=1),
|
||||
torch.nn.Dropout(dropout_rate)
|
||||
))
|
||||
self.linear = torch.nn.Linear(n_chans, num_bins)
|
||||
self.embed_positions = SinusoidalPositionalEmbedding(in_dims, 0, init_size=4096)
|
||||
self.pos_embed_alpha = nn.Parameter(torch.Tensor([1]))
|
||||
|
||||
def bins_to_values(self, bins):
|
||||
return bins * self.interval + self.vmin
|
||||
|
||||
def out2pitch(self, probs):
|
||||
logits = probs.sigmoid() # [B, T, N]
|
||||
# return logits
|
||||
# logits_sum = logits.sum(dim=2) # [B, T]
|
||||
bins = torch.sum(self.x * logits, dim=2) / torch.sum(logits, dim=2) # [B, T]
|
||||
pitch = self.bins_to_values(bins)
|
||||
# uv = logits_sum / (self.sigma * math.sqrt(2 * math.pi)) < 0.3
|
||||
# pitch[uv] = torch.nan
|
||||
return pitch
|
||||
|
||||
def forward(self, xs, base):
|
||||
"""
|
||||
:param xs: [B, T, H]
|
||||
:param base: [B, T]
|
||||
:return: [B, T, N]
|
||||
"""
|
||||
xs = xs + self.base_pitch_embed(base[..., None])
|
||||
positions = self.pos_embed_alpha * self.embed_positions(xs[..., 0])
|
||||
xs = xs + positions
|
||||
xs = xs.transpose(1, -1) # (B, idim, Tmax)
|
||||
for f in self.conv:
|
||||
xs = f(xs) # (B, C, Tmax)
|
||||
xs = self.linear(xs.transpose(1, -1)) # (B, Tmax, H)
|
||||
return self.out2pitch(xs) + base, xs
|
||||
|
||||
|
||||
class RhythmRegulator(torch.nn.Module):
|
||||
def __init__(self, eps=1e-5):
|
||||
super().__init__()
|
||||
self.eps = eps
|
||||
|
||||
def forward(self, ph_dur, ph2word, word_dur):
|
||||
"""
|
||||
Example (no batch dim version):
|
||||
1. ph_dur = [4,2,3,2]
|
||||
2. word_dur = [3,4,2], ph2word = [1,2,2,3]
|
||||
3. word_dur_in = [4,5,2]
|
||||
4. alpha_w = [0.75,0.8,1], alpha_ph = [0.75,0.8,0.8,1]
|
||||
5. ph_dur_out = [3,1.6,2.4,2]
|
||||
:param ph_dur: [B, T_ph]
|
||||
:param ph2word: [B, T_ph]
|
||||
:param word_dur: [B, T_w]
|
||||
"""
|
||||
ph_dur = ph_dur.float() * (ph2word > 0)
|
||||
word_dur = word_dur.float()
|
||||
word_dur_in = ph_dur.new_zeros(ph_dur.shape[0], ph2word.max() + 1).scatter_add(
|
||||
1, ph2word, ph_dur
|
||||
)[:, 1:] # [B, T_ph] => [B, T_w]
|
||||
alpha_w = word_dur / word_dur_in.clamp(min=self.eps) # avoid dividing by zero
|
||||
alpha_ph = torch.gather(F.pad(alpha_w, [1, 0]), 1, ph2word) # [B, T_w] => [B, T_ph]
|
||||
ph_dur_out = ph_dur * alpha_ph
|
||||
return ph_dur_out.round().long()
|
||||
|
||||
|
||||
class LengthRegulator(torch.nn.Module):
|
||||
# noinspection PyMethodMayBeStatic
|
||||
def forward(self, dur, dur_padding=None, alpha=None):
|
||||
"""
|
||||
Example (no batch dim version):
|
||||
1. dur = [2,2,3]
|
||||
2. token_idx = [[1],[2],[3]], dur_cumsum = [2,4,7], dur_cumsum_prev = [0,2,4]
|
||||
3. token_mask = [[1,1,0,0,0,0,0],
|
||||
[0,0,1,1,0,0,0],
|
||||
[0,0,0,0,1,1,1]]
|
||||
4. token_idx * token_mask = [[1,1,0,0,0,0,0],
|
||||
[0,0,2,2,0,0,0],
|
||||
[0,0,0,0,3,3,3]]
|
||||
5. (token_idx * token_mask).sum(0) = [1,1,2,2,3,3,3]
|
||||
|
||||
:param dur: Batch of durations of each frame (B, T_txt)
|
||||
:param dur_padding: Batch of padding of each frame (B, T_txt)
|
||||
:param alpha: duration rescale coefficient
|
||||
:return:
|
||||
mel2ph (B, T_speech)
|
||||
"""
|
||||
assert alpha is None or alpha > 0
|
||||
if alpha is not None:
|
||||
dur = torch.round(dur.float() * alpha).long()
|
||||
if dur_padding is not None:
|
||||
dur = dur * (1 - dur_padding.long())
|
||||
token_idx = torch.arange(1, dur.shape[1] + 1)[None, :, None].to(dur.device)
|
||||
dur_cumsum = torch.cumsum(dur, 1)
|
||||
dur_cumsum_prev = F.pad(dur_cumsum, [1, -1], mode='constant', value=0)
|
||||
|
||||
pos_idx = torch.arange(dur.sum(-1).max())[None, None].to(dur.device)
|
||||
token_mask = (pos_idx >= dur_cumsum_prev[:, :, None]) & (pos_idx < dur_cumsum[:, :, None])
|
||||
mel2ph = (token_idx * token_mask.long()).sum(1)
|
||||
return mel2ph
|
||||
|
||||
|
||||
class StretchRegulator(torch.nn.Module):
|
||||
# noinspection PyMethodMayBeStatic
|
||||
def forward(self, mel2ph, dur=None):
|
||||
"""
|
||||
Example (no batch dim version):
|
||||
1. dur = [2,4,3]
|
||||
2. mel2ph = [1,1,2,2,2,2,3,3,3]
|
||||
3. mel2dur = [2,2,4,4,4,4,3,3,3]
|
||||
4. bound_mask = [0,1,0,0,0,1,0,0,1]
|
||||
5. 1 - bound_mask * mel2dur = [1,-1,1,1,1,-3,1,1,-2] => pad => [0,1,-1,1,1,1,-3,1,1]
|
||||
6. stretch_denorm = [0,1,0,1,2,3,0,1,2]
|
||||
|
||||
:param dur: Batch of durations of each frame (B, T_txt)
|
||||
:param mel2ph: Batch of mel2ph (B, T_speech)
|
||||
:return:
|
||||
stretch (B, T_speech)
|
||||
"""
|
||||
if dur is None:
|
||||
dur = mel2ph_to_dur(mel2ph, mel2ph.max())
|
||||
dur = torch.cat([torch.ones_like(dur[:, :1]), dur], dim=1) # Avoid dividing by zero
|
||||
mel2dur = torch.gather(dur, 1, mel2ph)
|
||||
bound_mask = torch.gt(mel2ph[:, 1:], mel2ph[:, :-1])
|
||||
stretch_delta = 1 - bound_mask * mel2dur[:, :-1]
|
||||
stretch_delta = F.pad(stretch_delta, [1, 0])
|
||||
stretch_denorm = torch.cumsum(stretch_delta, dim=1)
|
||||
stretch = stretch_denorm.float() / mel2dur
|
||||
return stretch * (mel2ph > 0)
|
||||
|
||||
|
||||
def mel2ph_to_dur(mel2ph, T_txt, max_dur=None):
|
||||
B, _ = mel2ph.shape
|
||||
dur = mel2ph.new_zeros(B, T_txt + 1).scatter_add(1, mel2ph, torch.ones_like(mel2ph))
|
||||
dur = dur[:, 1:]
|
||||
if max_dur is not None:
|
||||
dur = dur.clamp(max=max_dur)
|
||||
return dur
|
||||
|
||||
|
||||
class FastSpeech2Encoder(nn.Module):
|
||||
def __init__(
|
||||
self, hidden_size, num_layers,
|
||||
ffn_kernel_size=9, ffn_act='gelu',
|
||||
dropout=None, num_heads=2, use_pos_embed=True, rel_pos=True,
|
||||
use_rope=False, rope_interleaved=True, mix_ln_layer=None
|
||||
):
|
||||
super().__init__()
|
||||
self.num_layers = num_layers
|
||||
embed_dim = self.hidden_size = hidden_size
|
||||
self.dropout = dropout
|
||||
self.use_pos_embed = use_pos_embed
|
||||
if use_pos_embed and use_rope:
|
||||
if embed_dim % (num_heads * 2) != 0:
|
||||
raise ValueError(
|
||||
"RoPE requires the hidden size to be multiple of "
|
||||
f"num_heads * 2 = {num_heads * 2}, but got {embed_dim}."
|
||||
)
|
||||
rotary_embed = RotaryEmbedding(dim=embed_dim // num_heads, interleaved=rope_interleaved)
|
||||
else:
|
||||
rotary_embed = None
|
||||
self.layers = nn.ModuleList([
|
||||
TransformerEncoderLayer(
|
||||
self.hidden_size, self.dropout,
|
||||
kernel_size=ffn_kernel_size, act=ffn_act,
|
||||
num_heads=num_heads, rotary_embed=rotary_embed,
|
||||
layer_idx=i, mix_ln_layer=mix_ln_layer
|
||||
)
|
||||
for i in range(self.num_layers)
|
||||
])
|
||||
self.layer_norm = nn.LayerNorm(embed_dim)
|
||||
|
||||
self.embed_scale = math.sqrt(hidden_size)
|
||||
self.padding_idx = 0
|
||||
self.rel_pos = rel_pos
|
||||
if use_rope:
|
||||
self.embed_positions = None
|
||||
elif self.rel_pos:
|
||||
self.embed_positions = RelPositionalEncoding(hidden_size, dropout_rate=0.0)
|
||||
else:
|
||||
self.embed_positions = SinusoidalPositionalEmbedding(
|
||||
hidden_size, self.padding_idx, init_size=DEFAULT_MAX_TARGET_POSITIONS,
|
||||
)
|
||||
|
||||
def forward_embedding(self, main_embed, extra_embed=None, padding_mask=None):
|
||||
# embed tokens and positions
|
||||
x = self.embed_scale * main_embed
|
||||
if extra_embed is not None:
|
||||
x = x + extra_embed
|
||||
if self.use_pos_embed and self.embed_positions is not None:
|
||||
if self.rel_pos:
|
||||
x = self.embed_positions(x)
|
||||
else:
|
||||
positions = self.embed_positions(~padding_mask)
|
||||
x = x + positions
|
||||
x = F.dropout(x, p=self.dropout, training=self.training)
|
||||
return x
|
||||
|
||||
def forward(self, main_embed, extra_embed, padding_mask, spk_embed=None, attn_mask=None, return_hiddens=False):
|
||||
x = self.forward_embedding(main_embed, extra_embed, padding_mask=padding_mask) # [B, T, H]
|
||||
nonpadding_mask_BT = 1 - padding_mask.float()[:, :, None] # [B, T, 1]
|
||||
|
||||
# NOTICE:
|
||||
# The following codes are commented out because
|
||||
# `self.use_pos_embed` is always False in the older versions,
|
||||
# and this argument did not compat with `hparams['use_pos_embed']`,
|
||||
# which defaults to True. The new version fixed this inconsistency,
|
||||
# resulting in temporary removal of pos_embed_alpha, which has actually
|
||||
# never been used before.
|
||||
|
||||
# if self.use_pos_embed:
|
||||
# positions = self.pos_embed_alpha * self.embed_positions(x[..., 0])
|
||||
# x = x + positions
|
||||
# x = F.dropout(x, p=self.dropout, training=self.training)
|
||||
|
||||
x = x * nonpadding_mask_BT
|
||||
hiddens = []
|
||||
for layer in self.layers:
|
||||
x = layer(x, encoder_padding_mask=padding_mask, cond=spk_embed, attn_mask=attn_mask) * nonpadding_mask_BT
|
||||
if return_hiddens:
|
||||
hiddens.append(x)
|
||||
x = self.layer_norm(x) * nonpadding_mask_BT
|
||||
if return_hiddens:
|
||||
x = torch.stack(hiddens, 0) # [L, B, T, C]
|
||||
return x
|
||||
@@ -0,0 +1,158 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.nn import functional as F
|
||||
|
||||
from modules.commons.common_layers import (
|
||||
NormalInitEmbedding as Embedding,
|
||||
XavierUniformInitLinear as Linear,
|
||||
AdamWLinear,
|
||||
)
|
||||
from modules.fastspeech.tts_modules import FastSpeech2Encoder, DurationPredictor
|
||||
from utils.hparams import hparams
|
||||
from utils.phoneme_utils import PAD_INDEX
|
||||
|
||||
|
||||
class FastSpeech2Variance(nn.Module):
|
||||
def __init__(self, vocab_size):
|
||||
super().__init__()
|
||||
self.predict_dur = hparams['predict_dur']
|
||||
self.linguistic_mode = 'word' if hparams['predict_dur'] else 'phoneme'
|
||||
self.use_lang_id = hparams['use_lang_id']
|
||||
self.use_variance_scaling = hparams.get('use_variance_scaling', False)
|
||||
self.txt_embed = Embedding(vocab_size, hparams['hidden_size'], PAD_INDEX)
|
||||
if self.use_lang_id:
|
||||
self.lang_embed = Embedding(hparams['num_lang'] + 1, hparams['hidden_size'], padding_idx=0)
|
||||
|
||||
if self.predict_dur:
|
||||
self.onset_embed = Embedding(2, hparams['hidden_size'])
|
||||
self.word_dur_embed = AdamWLinear(1, hparams['hidden_size'])
|
||||
else:
|
||||
self.ph_dur_embed = AdamWLinear(1, hparams['hidden_size'])
|
||||
|
||||
self.encoder = FastSpeech2Encoder(
|
||||
hidden_size=hparams['hidden_size'], num_layers=hparams['enc_layers'],
|
||||
ffn_kernel_size=hparams['enc_ffn_kernel_size'], ffn_act=hparams['ffn_act'],
|
||||
dropout=hparams['dropout'], num_heads=hparams['num_heads'],
|
||||
use_pos_embed=hparams['use_pos_embed'], rel_pos=hparams.get('rel_pos', False),
|
||||
use_rope=hparams.get('use_rope', False), rope_interleaved=hparams.get('rope_interleaved', True)
|
||||
)
|
||||
|
||||
dur_hparams = hparams['dur_prediction_args']
|
||||
if self.predict_dur:
|
||||
self.midi_embed = Embedding(128, hparams['hidden_size'])
|
||||
self.dur_predictor = DurationPredictor(
|
||||
in_dims=hparams['hidden_size'],
|
||||
n_chans=dur_hparams['hidden_size'],
|
||||
n_layers=dur_hparams['num_layers'],
|
||||
dropout_rate=dur_hparams['dropout'],
|
||||
kernel_size=dur_hparams['kernel_size'],
|
||||
offset=dur_hparams['log_offset'],
|
||||
dur_loss_type=dur_hparams['loss_type'],
|
||||
arch=dur_hparams['arch']
|
||||
)
|
||||
|
||||
def forward(
|
||||
self, txt_tokens, midi, ph2word,
|
||||
ph_dur=None, word_dur=None,
|
||||
spk_embed=None, languages=None,
|
||||
infer=True
|
||||
):
|
||||
"""
|
||||
:param txt_tokens: (train, infer) [B, T_ph]
|
||||
:param midi: (train, infer) [B, T_ph]
|
||||
:param ph2word: (train, infer) [B, T_ph]
|
||||
:param ph_dur: (train, [infer]) [B, T_ph]
|
||||
:param word_dur: (infer) [B, T_w]
|
||||
:param spk_embed: (train) [B, T_ph, H]
|
||||
:param languages (train, infer) [B, T_ph]
|
||||
:param infer: whether inference
|
||||
:return: encoder_out, ph_dur_pred
|
||||
"""
|
||||
txt_embed = self.txt_embed(txt_tokens)
|
||||
if self.linguistic_mode == 'word':
|
||||
b = txt_tokens.shape[0]
|
||||
onset = torch.diff(ph2word, dim=1, prepend=ph2word.new_zeros(b, 1)) > 0
|
||||
onset_embed = self.onset_embed(onset.long()) # [B, T_ph, H]
|
||||
|
||||
if word_dur is None or not infer:
|
||||
word_dur = ph_dur.new_zeros(b, ph2word.max() + 1).scatter_add(
|
||||
1, ph2word, ph_dur
|
||||
)[:, 1:] # [B, T_ph] => [B, T_w]
|
||||
word_dur = torch.gather(F.pad(word_dur, [1, 0], value=0), 1, ph2word) # [B, T_w] => [B, T_ph]
|
||||
word_dur_embed = self.word_dur_embed(word_dur.float()[:, :, None])
|
||||
extra_embed = onset_embed + word_dur_embed
|
||||
elif self.use_variance_scaling:
|
||||
extra_embed = self.ph_dur_embed(torch.log(1 + ph_dur.float())[:, :, None])
|
||||
else:
|
||||
extra_embed = self.ph_dur_embed(ph_dur.float()[:, :, None])
|
||||
|
||||
if self.use_lang_id:
|
||||
lang_embed = self.lang_embed(languages)
|
||||
extra_embed += lang_embed
|
||||
encoder_out = self.encoder(txt_embed, extra_embed, txt_tokens == 0)
|
||||
|
||||
if self.predict_dur:
|
||||
midi_embed = self.midi_embed(midi) # => [B, T_ph, H]
|
||||
dur_cond = encoder_out + midi_embed
|
||||
if spk_embed is not None:
|
||||
dur_cond += spk_embed
|
||||
ph_dur_pred = self.dur_predictor(dur_cond, x_masks=txt_tokens == PAD_INDEX, infer=infer)
|
||||
|
||||
return encoder_out, ph_dur_pred
|
||||
else:
|
||||
return encoder_out, None
|
||||
|
||||
|
||||
class MelodyEncoder(nn.Module):
|
||||
def __init__(self, enc_hparams: dict):
|
||||
super().__init__()
|
||||
|
||||
def get_hparam(key):
|
||||
return enc_hparams.get(key, hparams.get(key))
|
||||
|
||||
# MIDI inputs
|
||||
hidden_size = get_hparam('hidden_size')
|
||||
self.use_variance_scaling = hparams.get('use_variance_scaling', False)
|
||||
self.note_midi_embed = AdamWLinear(1, hidden_size)
|
||||
self.note_dur_embed = AdamWLinear(1, hidden_size)
|
||||
|
||||
# ornament inputs
|
||||
self.use_glide_embed = hparams['use_glide_embed']
|
||||
self.glide_embed_scale = hparams['glide_embed_scale']
|
||||
if self.use_glide_embed:
|
||||
# 0: none, 1: up, 2: down
|
||||
self.note_glide_embed = Embedding(len(hparams['glide_types']) + 1, hidden_size, padding_idx=0)
|
||||
|
||||
self.encoder = FastSpeech2Encoder(
|
||||
hidden_size=hidden_size, num_layers=get_hparam('enc_layers'),
|
||||
ffn_kernel_size=get_hparam('enc_ffn_kernel_size'), ffn_act=get_hparam('ffn_act'),
|
||||
dropout=get_hparam('dropout'), num_heads=get_hparam('num_heads'),
|
||||
use_pos_embed=get_hparam('use_pos_embed'), rel_pos=get_hparam('rel_pos'),
|
||||
use_rope=get_hparam('use_rope'), rope_interleaved=hparams.get('rope_interleaved', True)
|
||||
)
|
||||
self.out_proj = Linear(hidden_size, hparams['hidden_size'])
|
||||
|
||||
def forward(self, note_midi, note_rest, note_dur, glide=None):
|
||||
"""
|
||||
:param note_midi: float32 [B, T_n], -1: padding
|
||||
:param note_rest: bool [B, T_n]
|
||||
:param note_dur: int64 [B, T_n]
|
||||
:param glide: int64 [B, T_n]
|
||||
:return: [B, T_n, H]
|
||||
"""
|
||||
if self.use_variance_scaling:
|
||||
midi_embed = self.note_midi_embed(note_midi[:, :, None] / 128)
|
||||
dur_embed = self.note_dur_embed(torch.log(1 + note_dur.float())[:, :, None])
|
||||
else:
|
||||
midi_embed = self.note_midi_embed(note_midi[:, :, None])
|
||||
dur_embed = self.note_dur_embed(note_dur.float()[:, :, None])
|
||||
midi_embed *= ~note_rest[:, :, None]
|
||||
ornament_embed = 0
|
||||
if self.use_glide_embed:
|
||||
ornament_embed += self.note_glide_embed(glide) * self.glide_embed_scale
|
||||
encoder_out = self.encoder(
|
||||
midi_embed, dur_embed + ornament_embed,
|
||||
padding_mask=note_midi < 0
|
||||
)
|
||||
encoder_out = self.out_proj(encoder_out)
|
||||
return encoder_out
|
||||
@@ -0,0 +1,35 @@
|
||||
import pathlib
|
||||
|
||||
import torch
|
||||
import yaml
|
||||
|
||||
from .nets import CascadedNet
|
||||
|
||||
|
||||
class DotDict(dict):
|
||||
def __getattr__(*args):
|
||||
val = dict.get(*args)
|
||||
return DotDict(val) if type(val) is dict else val
|
||||
|
||||
__setattr__ = dict.__setitem__
|
||||
__delattr__ = dict.__delitem__
|
||||
|
||||
|
||||
def load_sep_model(model_path, device='cpu'):
|
||||
model_path = pathlib.Path(model_path)
|
||||
config_file = model_path.with_name('config.yaml')
|
||||
with open(config_file, "r") as config:
|
||||
args = yaml.safe_load(config)
|
||||
args = DotDict(args)
|
||||
model = CascadedNet(
|
||||
args.n_fft,
|
||||
args.hop_length,
|
||||
args.n_out,
|
||||
args.n_out_lstm,
|
||||
True,
|
||||
is_mono=args.is_mono
|
||||
)
|
||||
model.to(device)
|
||||
model.load_state_dict(torch.load(model_path, map_location='cpu'))
|
||||
model.eval()
|
||||
return model
|
||||
@@ -0,0 +1,166 @@
|
||||
import torch
|
||||
from torch import nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
def crop_center(h1, h2):
|
||||
h1_shape = h1.size()
|
||||
h2_shape = h2.size()
|
||||
|
||||
if h1_shape[3] == h2_shape[3]:
|
||||
return h1
|
||||
elif h1_shape[3] < h2_shape[3]:
|
||||
raise ValueError('h1_shape[3] must be greater than h2_shape[3]')
|
||||
|
||||
# s_freq = (h2_shape[2] - h1_shape[2]) // 2
|
||||
# e_freq = s_freq + h1_shape[2]
|
||||
s_time = (h1_shape[3] - h2_shape[3]) // 2
|
||||
e_time = s_time + h2_shape[3]
|
||||
h1 = h1[:, :, :, s_time:e_time]
|
||||
|
||||
return h1
|
||||
|
||||
|
||||
class Conv2DBNActiv(nn.Module):
|
||||
|
||||
def __init__(self, nin, nout, ksize=3, stride=1, pad=1, dilation=1, activ=nn.ReLU):
|
||||
super(Conv2DBNActiv, self).__init__()
|
||||
self.conv = nn.Sequential(
|
||||
nn.Conv2d(
|
||||
nin, nout,
|
||||
kernel_size=ksize,
|
||||
stride=stride,
|
||||
padding=pad,
|
||||
dilation=dilation,
|
||||
bias=False
|
||||
),
|
||||
nn.BatchNorm2d(nout),
|
||||
activ()
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
return self.conv(x)
|
||||
|
||||
|
||||
class Encoder(nn.Module):
|
||||
|
||||
def __init__(self, nin, nout, ksize=3, stride=1, pad=1, activ=nn.LeakyReLU):
|
||||
super(Encoder, self).__init__()
|
||||
self.conv1 = Conv2DBNActiv(nin, nout, ksize, stride, pad, activ=activ)
|
||||
self.conv2 = Conv2DBNActiv(nout, nout, ksize, 1, pad, activ=activ)
|
||||
|
||||
def forward(self, x):
|
||||
h = self.conv1(x)
|
||||
h = self.conv2(h)
|
||||
|
||||
return h
|
||||
|
||||
|
||||
class Decoder(nn.Module):
|
||||
|
||||
def __init__(self, nin, nout, ksize=3, stride=1, pad=1, activ=nn.ReLU, dropout=False):
|
||||
super(Decoder, self).__init__()
|
||||
self.conv1 = Conv2DBNActiv(nin, nout, ksize, 1, pad, activ=activ)
|
||||
# self.conv2 = Conv2DBNActiv(nout, nout, ksize, 1, pad, activ=activ)
|
||||
self.dropout = nn.Dropout2d(0.1) if dropout else None
|
||||
|
||||
def forward(self, x, skip=None, fixed_length=True):
|
||||
if fixed_length:
|
||||
x = F.interpolate(x, scale_factor=2, mode='bilinear', align_corners=True)
|
||||
else:
|
||||
_, _, h, w = x.size()
|
||||
x = F.pad(x, (0, 1, 0, 1), mode='replicate')
|
||||
x = F.interpolate(x, size=(2*h+1,2*w+1), mode='bilinear', align_corners=True)
|
||||
x = x[:, :, :-1, :-1]
|
||||
|
||||
if skip is not None:
|
||||
skip = crop_center(skip, x)
|
||||
x = torch.cat([x, skip], dim=1)
|
||||
|
||||
h = self.conv1(x)
|
||||
# h = self.conv2(h)
|
||||
|
||||
if self.dropout is not None:
|
||||
h = self.dropout(h)
|
||||
|
||||
return h
|
||||
|
||||
|
||||
class Mean(nn.Module):
|
||||
def __init__(self, dim, keepdims=False):
|
||||
super(Mean, self).__init__()
|
||||
self.dim = dim
|
||||
self.keepdims = keepdims
|
||||
|
||||
def forward(self, x):
|
||||
return x.mean(self.dim, keepdims=self.keepdims)
|
||||
|
||||
|
||||
class ASPPModule(nn.Module):
|
||||
|
||||
def __init__(self, nin, nout, dilations=(4, 8, 12), activ=nn.ReLU, dropout=False):
|
||||
super(ASPPModule, self).__init__()
|
||||
self.conv1 = nn.Sequential(
|
||||
Mean(dim=-2, keepdims=True), # nn.AdaptiveAvgPool2d((1, None)),
|
||||
Conv2DBNActiv(nin, nout, 1, 1, 0, activ=activ)
|
||||
)
|
||||
self.conv2 = Conv2DBNActiv(
|
||||
nin, nout, 1, 1, 0, activ=activ
|
||||
)
|
||||
self.conv3 = Conv2DBNActiv(
|
||||
nin, nout, 3, 1, dilations[0], dilations[0], activ=activ
|
||||
)
|
||||
self.conv4 = Conv2DBNActiv(
|
||||
nin, nout, 3, 1, dilations[1], dilations[1], activ=activ
|
||||
)
|
||||
self.conv5 = Conv2DBNActiv(
|
||||
nin, nout, 3, 1, dilations[2], dilations[2], activ=activ
|
||||
)
|
||||
self.bottleneck = Conv2DBNActiv(
|
||||
nout * 5, nout, 1, 1, 0, activ=activ
|
||||
)
|
||||
self.dropout = nn.Dropout2d(0.1) if dropout else None
|
||||
|
||||
def forward(self, x):
|
||||
_, _, h, w = x.size()
|
||||
# feat1 = F.interpolate(self.conv1(x), size=(h, w), mode='bilinear', align_corners=True)
|
||||
feat1 = self.conv1(x).repeat(1, 1, h, 1)
|
||||
feat2 = self.conv2(x)
|
||||
feat3 = self.conv3(x)
|
||||
feat4 = self.conv4(x)
|
||||
feat5 = self.conv5(x)
|
||||
out = torch.cat((feat1, feat2, feat3, feat4, feat5), dim=1)
|
||||
out = self.bottleneck(out)
|
||||
|
||||
if self.dropout is not None:
|
||||
out = self.dropout(out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class LSTMModule(nn.Module):
|
||||
|
||||
def __init__(self, nin_conv, nin_lstm, nout_lstm):
|
||||
super(LSTMModule, self).__init__()
|
||||
self.conv = Conv2DBNActiv(nin_conv, 1, 1, 1, 0)
|
||||
self.lstm = nn.LSTM(
|
||||
input_size=nin_lstm,
|
||||
hidden_size=nout_lstm // 2,
|
||||
bidirectional=True
|
||||
)
|
||||
self.dense = nn.Sequential(
|
||||
nn.Linear(nout_lstm, nin_lstm),
|
||||
nn.BatchNorm1d(nin_lstm),
|
||||
nn.ReLU()
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
N, _, nbins, nframes = x.size()
|
||||
h = self.conv(x)[:, 0] # N, nbins, nframes
|
||||
h = h.permute(2, 0, 1) # nframes, N, nbins
|
||||
h, _ = self.lstm(h)
|
||||
h = self.dense(h.reshape(-1, h.size()[-1])) # nframes * N, nbins
|
||||
h = h.reshape(nframes, N, 1, nbins)
|
||||
h = h.permute(1, 2, 3, 0)
|
||||
|
||||
return h
|
||||
@@ -0,0 +1,202 @@
|
||||
import torch
|
||||
from torch import nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from . import layers
|
||||
|
||||
|
||||
class BaseNet(nn.Module):
|
||||
|
||||
def __init__(self, nin, nout, nin_lstm, nout_lstm, dilations=((4, 2), (8, 4), (12, 6)), fixed_length=True):
|
||||
super(BaseNet, self).__init__()
|
||||
self.enc1 = layers.Conv2DBNActiv(nin, nout, 3, 1, 1)
|
||||
self.enc2 = layers.Encoder(nout, nout * 2, 3, 2, 1)
|
||||
self.enc3 = layers.Encoder(nout * 2, nout * 4, 3, 2, 1)
|
||||
self.enc4 = layers.Encoder(nout * 4, nout * 6, 3, 2, 1)
|
||||
self.enc5 = layers.Encoder(nout * 6, nout * 8, 3, 2, 1)
|
||||
|
||||
self.aspp = layers.ASPPModule(nout * 8, nout * 8, dilations, dropout=True)
|
||||
|
||||
self.dec4 = layers.Decoder(nout * (6 + 8), nout * 6, 3, 1, 1)
|
||||
self.dec3 = layers.Decoder(nout * (4 + 6), nout * 4, 3, 1, 1)
|
||||
self.dec2 = layers.Decoder(nout * (2 + 4), nout * 2, 3, 1, 1)
|
||||
self.lstm_dec2 = layers.LSTMModule(nout * 2, nin_lstm, nout_lstm)
|
||||
self.dec1 = layers.Decoder(nout * (1 + 2) + 1, nout * 1, 3, 1, 1)
|
||||
|
||||
self.fixed_length = fixed_length
|
||||
|
||||
def __call__(self, x):
|
||||
e1 = self.enc1(x)
|
||||
e2 = self.enc2(e1)
|
||||
e3 = self.enc3(e2)
|
||||
e4 = self.enc4(e3)
|
||||
e5 = self.enc5(e4)
|
||||
|
||||
h = self.aspp(e5)
|
||||
|
||||
h = self.dec4(h, e4, fixed_length=self.fixed_length)
|
||||
h = self.dec3(h, e3, fixed_length=self.fixed_length)
|
||||
h = self.dec2(h, e2, fixed_length=self.fixed_length)
|
||||
h = torch.cat([h, self.lstm_dec2(h)], dim=1)
|
||||
h = self.dec1(h, e1, fixed_length=self.fixed_length)
|
||||
|
||||
return h
|
||||
|
||||
|
||||
class CascadedNet(nn.Module):
|
||||
|
||||
def __init__(self, n_fft, hop_length, nout=32, nout_lstm=128, is_complex=False, is_mono=False, fixed_length=True):
|
||||
super(CascadedNet, self).__init__()
|
||||
self.n_fft = n_fft
|
||||
self.hop_length = hop_length
|
||||
self.seg_length = 32 * hop_length
|
||||
self.is_complex = is_complex
|
||||
self.is_mono = is_mono
|
||||
self.register_buffer("window", torch.hann_window(n_fft), persistent=False)
|
||||
self.max_bin = n_fft // 2
|
||||
self.output_bin = n_fft // 2 + 1
|
||||
self.nin_lstm = self.max_bin // 2
|
||||
self.offset = 64
|
||||
|
||||
nin = 4 if is_complex else 2
|
||||
if is_mono:
|
||||
nin = nin // 2
|
||||
|
||||
self.stg1_low_band_net = nn.Sequential(
|
||||
BaseNet(nin, nout // 2, self.nin_lstm // 2, nout_lstm, fixed_length=fixed_length),
|
||||
layers.Conv2DBNActiv(nout // 2, nout // 4, 1, 1, 0)
|
||||
)
|
||||
self.stg1_high_band_net = BaseNet(
|
||||
nin, nout // 4, self.nin_lstm // 2, nout_lstm // 2, fixed_length=fixed_length
|
||||
)
|
||||
|
||||
self.stg2_low_band_net = nn.Sequential(
|
||||
BaseNet(nout // 4 + nin, nout, self.nin_lstm // 2, nout_lstm, fixed_length=fixed_length),
|
||||
layers.Conv2DBNActiv(nout, nout // 2, 1, 1, 0)
|
||||
)
|
||||
self.stg2_high_band_net = BaseNet(
|
||||
nout // 4 + nin, nout // 2, self.nin_lstm // 2, nout_lstm // 2, fixed_length=fixed_length
|
||||
)
|
||||
|
||||
self.stg3_full_band_net = BaseNet(
|
||||
3 * nout // 4 + nin, nout, self.nin_lstm, nout_lstm, fixed_length=fixed_length
|
||||
)
|
||||
|
||||
self.out = nn.Conv2d(nout, nin, 1, bias=False)
|
||||
self.aux_out = nn.Conv2d(3 * nout // 4, nin, 1, bias=False)
|
||||
|
||||
def forward(self, x):
|
||||
if self.is_complex:
|
||||
x = torch.cat([x.real, x.imag], dim=1)
|
||||
|
||||
x = x[:, :, :self.max_bin]
|
||||
|
||||
bandw = x.size()[2] // 2
|
||||
l1_in = x[:, :, :bandw]
|
||||
h1_in = x[:, :, bandw:]
|
||||
l1 = self.stg1_low_band_net(l1_in)
|
||||
h1 = self.stg1_high_band_net(h1_in)
|
||||
aux1 = torch.cat([l1, h1], dim=2)
|
||||
|
||||
l2_in = torch.cat([l1_in, l1], dim=1)
|
||||
h2_in = torch.cat([h1_in, h1], dim=1)
|
||||
l2 = self.stg2_low_band_net(l2_in)
|
||||
h2 = self.stg2_high_band_net(h2_in)
|
||||
aux2 = torch.cat([l2, h2], dim=2)
|
||||
|
||||
f3_in = torch.cat([x, aux1, aux2], dim=1)
|
||||
f3 = self.stg3_full_band_net(f3_in)
|
||||
|
||||
if self.is_complex:
|
||||
mask = self.out(f3)
|
||||
if self.is_mono:
|
||||
mask = torch.complex(mask[:, :1], mask[:, 1:])
|
||||
else:
|
||||
mask = torch.complex(mask[:, :2], mask[:, 2:])
|
||||
mask = self.bounded_mask(mask)
|
||||
else:
|
||||
mask = torch.sigmoid(self.out(f3))
|
||||
|
||||
mask = F.pad(
|
||||
input=mask,
|
||||
pad=(0, 0, 0, self.output_bin - mask.size()[2]),
|
||||
mode='replicate'
|
||||
)
|
||||
|
||||
return mask
|
||||
|
||||
def bounded_mask(self, mask, eps=1e-8):
|
||||
mask_mag = torch.abs(mask)
|
||||
mask = torch.tanh(mask_mag) * mask / (mask_mag + eps)
|
||||
return mask
|
||||
|
||||
def predict_mask(self, x):
|
||||
mask = self.forward(x)
|
||||
|
||||
if self.offset > 0:
|
||||
mask = mask[:, :, :, self.offset:-self.offset]
|
||||
assert mask.size()[3] > 0
|
||||
|
||||
return mask
|
||||
|
||||
def predict(self, x):
|
||||
mask = self.forward(x)
|
||||
pred = x * mask
|
||||
|
||||
if self.offset > 0:
|
||||
pred = pred[:, :, :, self.offset:-self.offset]
|
||||
assert pred.size()[3] > 0
|
||||
|
||||
return pred
|
||||
|
||||
def audio2spec(self, x, use_pad=False):
|
||||
B, C, T = x.shape
|
||||
x = x.reshape(B * C, T)
|
||||
if use_pad:
|
||||
T1 = T + self.hop_length
|
||||
T_pad = self.seg_length * ((T1 - 1) // self.seg_length + 1) - T1
|
||||
nl_pad = T_pad // 2 // self.hop_length
|
||||
Tl_pad = nl_pad * self.hop_length
|
||||
x = F.pad(x, (Tl_pad, T_pad - Tl_pad))
|
||||
spec = torch.stft(
|
||||
x,
|
||||
n_fft=self.n_fft,
|
||||
hop_length=self.hop_length,
|
||||
return_complex=True,
|
||||
window=self.window,
|
||||
pad_mode='constant'
|
||||
)
|
||||
spec = spec.reshape(B, C, spec.shape[-2], spec.shape[-1])
|
||||
return spec
|
||||
|
||||
def spec2audio(self, x):
|
||||
B, C, N, T = x.shape
|
||||
x = x.reshape(-1, N, T)
|
||||
x = torch.istft(x, self.n_fft, self.hop_length, window=self.window)
|
||||
x = x.reshape(B, C, -1)
|
||||
return x
|
||||
|
||||
def predict_from_audio(self, x):
|
||||
B, C, T = x.shape
|
||||
x = x.reshape(B * C, T)
|
||||
T1 = T + self.hop_length
|
||||
T_pad = self.seg_length * ((T1 - 1) // self.seg_length + 1) - T1
|
||||
nl_pad = T_pad // 2 // self.hop_length
|
||||
Tl_pad = nl_pad * self.hop_length
|
||||
x = F.pad(x, (Tl_pad, T_pad - Tl_pad))
|
||||
spec = torch.stft(
|
||||
x,
|
||||
n_fft=self.n_fft,
|
||||
hop_length=self.hop_length,
|
||||
return_complex=True,
|
||||
window=self.window,
|
||||
pad_mode='constant'
|
||||
)
|
||||
spec = spec.reshape(B, C, spec.shape[-2], spec.shape[-1])
|
||||
mask = self.forward(spec)
|
||||
spec_pred = spec * mask
|
||||
spec_pred = spec_pred.reshape(B * C, spec.shape[-2], spec.shape[-1])
|
||||
x_pred = torch.istft(spec_pred, self.n_fft, self.hop_length, window=self.window)
|
||||
x_pred = x_pred[:, Tl_pad: Tl_pad + T]
|
||||
x_pred = x_pred.reshape(B, C, T)
|
||||
return x_pred
|
||||
@@ -0,0 +1,3 @@
|
||||
from .diff_loss import DiffusionLoss
|
||||
from .reflow_loss import RectifiedFlowLoss
|
||||
from .dur_loss import DurationLoss
|
||||
@@ -0,0 +1,34 @@
|
||||
import torch.nn as nn
|
||||
from torch import Tensor
|
||||
|
||||
|
||||
class DiffusionLoss(nn.Module):
|
||||
def __init__(self, loss_type):
|
||||
super().__init__()
|
||||
self.loss_type = loss_type
|
||||
if self.loss_type == 'l1':
|
||||
self.loss = nn.L1Loss(reduction='none')
|
||||
elif self.loss_type == 'l2':
|
||||
self.loss = nn.MSELoss(reduction='none')
|
||||
else:
|
||||
raise NotImplementedError()
|
||||
|
||||
@staticmethod
|
||||
def _mask_non_padding(x_recon, noise, non_padding=None):
|
||||
if non_padding is not None:
|
||||
non_padding = non_padding.transpose(1, 2).unsqueeze(1)
|
||||
return x_recon * non_padding, noise * non_padding
|
||||
else:
|
||||
return x_recon, noise
|
||||
|
||||
def _forward(self, x_recon, noise):
|
||||
return self.loss(x_recon, noise)
|
||||
|
||||
def forward(self, x_recon: Tensor, noise: Tensor, non_padding: Tensor = None) -> Tensor:
|
||||
"""
|
||||
:param x_recon: [B, 1, M, T]
|
||||
:param noise: [B, 1, M, T]
|
||||
:param non_padding: [B, T, M]
|
||||
"""
|
||||
x_recon, noise = self._mask_non_padding(x_recon, noise, non_padding)
|
||||
return self._forward(x_recon, noise).mean()
|
||||
@@ -0,0 +1,56 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch import Tensor
|
||||
|
||||
|
||||
class DurationLoss(nn.Module):
|
||||
"""
|
||||
Loss module as combination of phone duration loss, word duration loss and sentence duration loss.
|
||||
"""
|
||||
|
||||
def __init__(self, offset, loss_type,
|
||||
lambda_pdur=0.6, lambda_wdur=0.3, lambda_sdur=0.1):
|
||||
super().__init__()
|
||||
self.loss_type = loss_type
|
||||
if self.loss_type == 'mse':
|
||||
self.loss = nn.MSELoss()
|
||||
elif self.loss_type == 'huber':
|
||||
self.loss = nn.HuberLoss()
|
||||
else:
|
||||
raise NotImplementedError()
|
||||
self.offset = offset
|
||||
|
||||
self.lambda_pdur = lambda_pdur
|
||||
self.lambda_wdur = lambda_wdur
|
||||
self.lambda_sdur = lambda_sdur
|
||||
|
||||
def linear2log(self, any_dur):
|
||||
return torch.log(any_dur + self.offset)
|
||||
|
||||
def forward(self, dur_pred: Tensor, dur_gt: Tensor, ph2word: Tensor) -> Tensor:
|
||||
dur_gt = dur_gt.to(dtype=dur_pred.dtype)
|
||||
|
||||
# pdur_loss
|
||||
pdur_loss = self.lambda_pdur * self.loss(self.linear2log(dur_pred), self.linear2log(dur_gt))
|
||||
|
||||
dur_pred = dur_pred.clamp(min=0.) # clip to avoid NaN loss
|
||||
|
||||
# wdur loss
|
||||
shape = dur_pred.shape[0], ph2word.max() + 1
|
||||
wdur_pred = dur_pred.new_zeros(*shape).scatter_add(
|
||||
1, ph2word, dur_pred
|
||||
)[:, 1:] # [B, T_ph] => [B, T_w]
|
||||
wdur_gt = dur_gt.new_zeros(*shape).scatter_add(
|
||||
1, ph2word, dur_gt
|
||||
)[:, 1:] # [B, T_ph] => [B, T_w]
|
||||
wdur_loss = self.lambda_wdur * self.loss(self.linear2log(wdur_pred), self.linear2log(wdur_gt))
|
||||
|
||||
# sdur loss
|
||||
sdur_pred = dur_pred.sum(dim=1)
|
||||
sdur_gt = dur_gt.sum(dim=1)
|
||||
sdur_loss = self.lambda_sdur * self.loss(self.linear2log(sdur_pred), self.linear2log(sdur_gt))
|
||||
|
||||
# combine
|
||||
dur_loss = pdur_loss + wdur_loss + sdur_loss
|
||||
|
||||
return dur_loss
|
||||
@@ -0,0 +1,50 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch import Tensor
|
||||
|
||||
|
||||
class RectifiedFlowLoss(nn.Module):
|
||||
def __init__(self, loss_type, log_norm=True):
|
||||
super().__init__()
|
||||
self.loss_type = loss_type
|
||||
self.log_norm = log_norm
|
||||
if self.loss_type == 'l1':
|
||||
self.loss = nn.L1Loss(reduction='none')
|
||||
elif self.loss_type == 'l2':
|
||||
self.loss = nn.MSELoss(reduction='none')
|
||||
else:
|
||||
raise NotImplementedError()
|
||||
|
||||
@staticmethod
|
||||
def _mask_non_padding(v_pred, v_gt, non_padding=None):
|
||||
if non_padding is not None:
|
||||
non_padding = non_padding.transpose(1, 2).unsqueeze(1)
|
||||
return v_pred * non_padding, v_gt * non_padding
|
||||
else:
|
||||
return v_pred, v_gt
|
||||
|
||||
@staticmethod
|
||||
def get_weights(t):
|
||||
eps = 1e-7
|
||||
t = t.float()
|
||||
t = torch.clip(t, 0 + eps, 1 - eps)
|
||||
weights = 0.398942 / t / (1 - t) * torch.exp(
|
||||
-0.5 * torch.log(t / (1 - t)) ** 2
|
||||
) + eps
|
||||
return weights[:, None, None, None]
|
||||
|
||||
def _forward(self, v_pred, v_gt, t=None):
|
||||
if self.log_norm:
|
||||
return self.get_weights(t) * self.loss(v_pred, v_gt)
|
||||
else:
|
||||
return self.loss(v_pred, v_gt)
|
||||
|
||||
def forward(self, v_pred: Tensor, v_gt: Tensor, t: Tensor, non_padding: Tensor = None) -> Tensor:
|
||||
"""
|
||||
:param v_pred: [B, 1, M, T]
|
||||
:param v_gt: [B, 1, M, T]
|
||||
:param t: [B,]
|
||||
:param non_padding: [B, T, M]
|
||||
"""
|
||||
v_pred, v_gt = self._mask_non_padding(v_pred, v_gt, non_padding)
|
||||
return self._forward(v_pred, v_gt, t=t).mean()
|
||||
@@ -0,0 +1,2 @@
|
||||
from .curve import RawCurveAccuracy, RawCurveR2Score
|
||||
from .duration import RhythmCorrectness, PhonemeDurationAccuracy
|
||||
@@ -0,0 +1,73 @@
|
||||
import torch
|
||||
import torchmetrics
|
||||
from torch import Tensor
|
||||
|
||||
|
||||
class RawCurveAccuracy(torchmetrics.Metric):
|
||||
def __init__(self, *, tolerance, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
self.tolerance = tolerance
|
||||
self.add_state('close', default=torch.tensor(0, dtype=torch.int), dist_reduce_fx='sum')
|
||||
self.add_state('total', default=torch.tensor(0, dtype=torch.int), dist_reduce_fx='sum')
|
||||
|
||||
def update(self, pred: Tensor, target: Tensor, mask=None) -> None:
|
||||
"""
|
||||
|
||||
:param pred: predicted curve
|
||||
:param target: reference curve
|
||||
:param mask: valid or non-padding mask
|
||||
"""
|
||||
if mask is None:
|
||||
assert pred.shape == target.shape, f'shapes of pred and target mismatch: {pred.shape}, {target.shape}'
|
||||
else:
|
||||
assert pred.shape == target.shape == mask.shape, \
|
||||
f'shapes of pred, target and mask mismatch: {pred.shape}, {target.shape}, {mask.shape}'
|
||||
close = torch.abs(pred - target) <= self.tolerance
|
||||
if mask is not None:
|
||||
close &= mask
|
||||
|
||||
self.close += close.sum()
|
||||
self.total += pred.numel() if mask is None else mask.sum()
|
||||
|
||||
def compute(self) -> Tensor:
|
||||
return self.close / self.total
|
||||
|
||||
|
||||
class RawCurveR2Score(torchmetrics.Metric):
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
self.add_state('sum_squared_error', default=torch.tensor(0.0), dist_reduce_fx='sum')
|
||||
self.add_state('sum_error', default=torch.tensor(0.0), dist_reduce_fx='sum')
|
||||
self.add_state('residual', default=torch.tensor(0.0), dist_reduce_fx='sum')
|
||||
self.add_state('total', default=torch.tensor(0), dist_reduce_fx='sum')
|
||||
|
||||
def update(self, pred: Tensor, target: Tensor, mask=None) -> None:
|
||||
"""
|
||||
|
||||
:param pred: predicted curve
|
||||
:param target: reference curve
|
||||
:param mask: valid or non-padding mask
|
||||
"""
|
||||
if mask is None:
|
||||
assert pred.shape == target.shape, f'shapes of pred and target mismatch: {pred.shape}, {target.shape}'
|
||||
else:
|
||||
assert pred.shape == target.shape == mask.shape, \
|
||||
f'shapes of pred, target and mask mismatch: {pred.shape}, {target.shape}, {mask.shape}'
|
||||
pred = pred[mask]
|
||||
target = target[mask]
|
||||
pred = pred.flatten()
|
||||
target = target.flatten()
|
||||
|
||||
sum_error = torch.sum(target)
|
||||
sum_squared_error = torch.sum(target * target)
|
||||
residual = target - pred
|
||||
rss = torch.sum(residual * residual)
|
||||
total = target.numel() if mask is None else mask.sum()
|
||||
|
||||
self.sum_squared_error += sum_squared_error
|
||||
self.sum_error += sum_error
|
||||
self.residual += rss
|
||||
self.total += total
|
||||
|
||||
def compute(self) -> Tensor:
|
||||
return 1 - self.residual / (self.sum_squared_error - self.sum_error ** 2 / self.total)
|
||||
@@ -0,0 +1,97 @@
|
||||
import torch
|
||||
import torchmetrics
|
||||
from torch import Tensor
|
||||
|
||||
from modules.fastspeech.tts_modules import RhythmRegulator
|
||||
|
||||
|
||||
def linguistic_checks(pred, target, ph2word, mask=None):
|
||||
if mask is None:
|
||||
assert pred.shape == target.shape == ph2word.shape, \
|
||||
f'shapes of pred, target and ph2word mismatch: {pred.shape}, {target.shape}, {ph2word.shape}'
|
||||
else:
|
||||
assert pred.shape == target.shape == ph2word.shape == mask.shape, \
|
||||
f'shapes of pred, target and mask mismatch: {pred.shape}, {target.shape}, {ph2word.shape}, {mask.shape}'
|
||||
assert pred.ndim == 2, f'all inputs should be 2D, but got {pred.shape}'
|
||||
assert torch.any(ph2word > 0), 'empty word sequence'
|
||||
assert torch.all(ph2word >= 0), 'unexpected negative word index'
|
||||
assert ph2word.max() <= pred.shape[1], f'word index out of range: {ph2word.max()} > {pred.shape[1]}'
|
||||
assert torch.all(pred >= 0.), f'unexpected negative ph_dur prediction'
|
||||
assert torch.all(target >= 0.), f'unexpected negative ph_dur target'
|
||||
|
||||
|
||||
class RhythmCorrectness(torchmetrics.Metric):
|
||||
def __init__(self, *, tolerance, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
assert 0. < tolerance < 1., 'tolerance should be within (0, 1)'
|
||||
self.tolerance = tolerance
|
||||
self.add_state('correct', default=torch.tensor(0, dtype=torch.int), dist_reduce_fx='sum')
|
||||
self.add_state('total', default=torch.tensor(0, dtype=torch.int), dist_reduce_fx='sum')
|
||||
|
||||
def update(self, pdur_pred: Tensor, pdur_target: Tensor, ph2word: Tensor, mask=None) -> None:
|
||||
"""
|
||||
|
||||
:param pdur_pred: predicted ph_dur
|
||||
:param pdur_target: reference ph_dur
|
||||
:param ph2word: word division sequence
|
||||
:param mask: valid or non-padding mask
|
||||
"""
|
||||
linguistic_checks(pdur_pred, pdur_target, ph2word, mask=mask)
|
||||
|
||||
shape = pdur_pred.shape[0], ph2word.max() + 1
|
||||
wdur_pred = pdur_pred.new_zeros(*shape).scatter_add(
|
||||
1, ph2word, pdur_pred
|
||||
)[:, 1:] # [B, T_ph] => [B, T_w]
|
||||
wdur_target = pdur_target.new_zeros(*shape).scatter_add(
|
||||
1, ph2word, pdur_target
|
||||
)[:, 1:] # [B, T_ph] => [B, T_w]
|
||||
if mask is None:
|
||||
wdur_mask = torch.ones_like(wdur_pred, dtype=torch.bool)
|
||||
else:
|
||||
wdur_mask = mask.new_zeros(*shape).scatter_add(
|
||||
1, ph2word, mask
|
||||
)[:, 1:].bool() # [B, T_ph] => [B, T_w]
|
||||
|
||||
correct = torch.abs(wdur_pred - wdur_target) <= wdur_target * self.tolerance
|
||||
correct &= wdur_mask
|
||||
|
||||
self.correct += correct.sum()
|
||||
self.total += wdur_mask.sum()
|
||||
|
||||
def compute(self) -> Tensor:
|
||||
return self.correct / self.total
|
||||
|
||||
|
||||
class PhonemeDurationAccuracy(torchmetrics.Metric):
|
||||
def __init__(self, *, tolerance, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
self.tolerance = tolerance
|
||||
self.rr = RhythmRegulator()
|
||||
self.add_state('accurate', default=torch.tensor(0, dtype=torch.int), dist_reduce_fx='sum')
|
||||
self.add_state('total', default=torch.tensor(0, dtype=torch.int), dist_reduce_fx='sum')
|
||||
|
||||
def update(self, pdur_pred: Tensor, pdur_target: Tensor, ph2word: Tensor, mask=None) -> None:
|
||||
"""
|
||||
|
||||
:param pdur_pred: predicted ph_dur
|
||||
:param pdur_target: reference ph_dur
|
||||
:param ph2word: word division sequence
|
||||
:param mask: valid or non-padding mask
|
||||
"""
|
||||
linguistic_checks(pdur_pred, pdur_target, ph2word, mask=mask)
|
||||
|
||||
shape = pdur_pred.shape[0], ph2word.max() + 1
|
||||
wdur_target = pdur_target.new_zeros(*shape).scatter_add(
|
||||
1, ph2word, pdur_target
|
||||
)[:, 1:] # [B, T_ph] => [B, T_w]
|
||||
pdur_align = self.rr(pdur_pred, ph2word=ph2word, word_dur=wdur_target)
|
||||
|
||||
accurate = torch.abs(pdur_align - pdur_target) <= pdur_target * self.tolerance
|
||||
if mask is not None:
|
||||
accurate &= mask
|
||||
|
||||
self.accurate += accurate.sum()
|
||||
self.total += pdur_pred.numel() if mask is None else mask.sum()
|
||||
|
||||
def compute(self) -> Tensor:
|
||||
return self.accurate / self.total
|
||||
@@ -0,0 +1,32 @@
|
||||
class AttrDict(dict):
|
||||
"""A dictionary with attribute-style access. It maps attribute access to
|
||||
the real dictionary. """
|
||||
def __init__(self, *args, **kwargs):
|
||||
dict.__init__(self, *args, **kwargs)
|
||||
|
||||
def __getstate__(self):
|
||||
return self.__dict__.items()
|
||||
|
||||
def __setstate__(self, items):
|
||||
for key, val in items:
|
||||
self.__dict__[key] = val
|
||||
|
||||
def __repr__(self):
|
||||
return "%s(%s)" % (self.__class__.__name__, dict.__repr__(self))
|
||||
|
||||
def __setitem__(self, key, value):
|
||||
return super(AttrDict, self).__setitem__(key, value)
|
||||
|
||||
def __getitem__(self, name):
|
||||
if name not in super(AttrDict, self).keys():
|
||||
return None
|
||||
return super(AttrDict, self).__getitem__(name)
|
||||
|
||||
def __delitem__(self, name):
|
||||
return super(AttrDict, self).__delitem__(name)
|
||||
|
||||
__getattr__ = __getitem__
|
||||
__setattr__ = __setitem__
|
||||
|
||||
def copy(self):
|
||||
return AttrDict(self)
|
||||
@@ -0,0 +1,303 @@
|
||||
import json
|
||||
import pathlib
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from lightning.pytorch.utilities.rank_zero import rank_zero_info
|
||||
from torch.nn import Conv1d, ConvTranspose1d
|
||||
from torch.nn.utils import weight_norm, remove_weight_norm
|
||||
|
||||
from .env import AttrDict
|
||||
from .utils import init_weights, get_padding
|
||||
|
||||
LRELU_SLOPE = 0.1
|
||||
|
||||
|
||||
def load_model(model_path: pathlib.Path):
|
||||
config_file = model_path.with_name('config.json')
|
||||
with open(config_file) as f:
|
||||
data = f.read()
|
||||
|
||||
json_config = json.loads(data)
|
||||
h = AttrDict(json_config)
|
||||
|
||||
generator = Generator(h)
|
||||
|
||||
cp_dict = torch.load(model_path, map_location='cpu')
|
||||
generator.load_state_dict(cp_dict['generator'])
|
||||
generator.eval()
|
||||
generator.remove_weight_norm()
|
||||
del cp_dict
|
||||
return generator, h
|
||||
|
||||
|
||||
class ResBlock1(torch.nn.Module):
|
||||
def __init__(self, h, channels, kernel_size=3, dilation=(1, 3, 5)):
|
||||
super(ResBlock1, self).__init__()
|
||||
self.h = h
|
||||
self.convs1 = nn.ModuleList([
|
||||
weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0],
|
||||
padding=get_padding(kernel_size, dilation[0]))),
|
||||
weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1],
|
||||
padding=get_padding(kernel_size, dilation[1]))),
|
||||
weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[2],
|
||||
padding=get_padding(kernel_size, dilation[2])))
|
||||
])
|
||||
self.convs1.apply(init_weights)
|
||||
|
||||
self.convs2 = nn.ModuleList([
|
||||
weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
|
||||
padding=get_padding(kernel_size, 1))),
|
||||
weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
|
||||
padding=get_padding(kernel_size, 1))),
|
||||
weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
|
||||
padding=get_padding(kernel_size, 1)))
|
||||
])
|
||||
self.convs2.apply(init_weights)
|
||||
|
||||
def forward(self, x):
|
||||
for c1, c2 in zip(self.convs1, self.convs2):
|
||||
xt = F.leaky_relu(x, LRELU_SLOPE)
|
||||
xt = c1(xt)
|
||||
xt = F.leaky_relu(xt, LRELU_SLOPE)
|
||||
xt = c2(xt)
|
||||
x = xt + x
|
||||
return x
|
||||
|
||||
def remove_weight_norm(self):
|
||||
for l in self.convs1:
|
||||
remove_weight_norm(l)
|
||||
for l in self.convs2:
|
||||
remove_weight_norm(l)
|
||||
|
||||
|
||||
class ResBlock2(torch.nn.Module):
|
||||
def __init__(self, h, channels, kernel_size=3, dilation=(1, 3)):
|
||||
super(ResBlock2, self).__init__()
|
||||
self.h = h
|
||||
self.convs = nn.ModuleList([
|
||||
weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0],
|
||||
padding=get_padding(kernel_size, dilation[0]))),
|
||||
weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1],
|
||||
padding=get_padding(kernel_size, dilation[1])))
|
||||
])
|
||||
self.convs.apply(init_weights)
|
||||
|
||||
def forward(self, x):
|
||||
for c in self.convs:
|
||||
xt = F.leaky_relu(x, LRELU_SLOPE)
|
||||
xt = c(xt)
|
||||
x = xt + x
|
||||
return x
|
||||
|
||||
def remove_weight_norm(self):
|
||||
for l in self.convs:
|
||||
remove_weight_norm(l)
|
||||
|
||||
|
||||
class SineGen(torch.nn.Module):
|
||||
""" Definition of sine generator
|
||||
SineGen(samp_rate, harmonic_num = 0,
|
||||
sine_amp = 0.1, noise_std = 0.003,
|
||||
voiced_threshold = 0,
|
||||
flag_for_pulse=False)
|
||||
samp_rate: sampling rate in Hz
|
||||
harmonic_num: number of harmonic overtones (default 0)
|
||||
sine_amp: amplitude of sine-waveform (default 0.1)
|
||||
noise_std: std of Gaussian noise (default 0.003)
|
||||
voiced_threshold: F0 threshold for U/V classification (default 0)
|
||||
flag_for_pulse: this SinGen is used inside PulseGen (default False)
|
||||
Note: when flag_for_pulse is True, the first time step of a voiced
|
||||
segment is always sin(np.pi) or cos(0)
|
||||
"""
|
||||
|
||||
def __init__(self, samp_rate, harmonic_num=0,
|
||||
sine_amp=0.1, noise_std=0.003,
|
||||
voiced_threshold=0):
|
||||
super(SineGen, self).__init__()
|
||||
self.sine_amp = sine_amp
|
||||
self.noise_std = noise_std
|
||||
self.harmonic_num = harmonic_num
|
||||
self.dim = self.harmonic_num + 1
|
||||
self.sampling_rate = samp_rate
|
||||
self.voiced_threshold = voiced_threshold
|
||||
|
||||
def _f02uv(self, f0):
|
||||
# generate uv signal
|
||||
uv = torch.ones_like(f0)
|
||||
uv = uv * (f0 > self.voiced_threshold)
|
||||
return uv
|
||||
|
||||
def _f02sine(self, f0, upp):
|
||||
""" f0: (batchsize, length, dim)
|
||||
where dim indicates fundamental tone and overtones
|
||||
"""
|
||||
rad = f0 / self.sampling_rate * torch.arange(1, upp + 1, device=f0.device)
|
||||
rad2 = torch.fmod(rad[..., -1:].float() + 0.5, 1.0) - 0.5
|
||||
rad_acc = rad2.cumsum(dim=1).fmod(1.0).to(f0)
|
||||
rad += F.pad(rad_acc[:, :-1, :], (0, 0, 1, 0))
|
||||
rad = rad.reshape(f0.shape[0], -1, 1)
|
||||
rad = torch.multiply(rad, torch.arange(1, self.dim + 1, device=f0.device).reshape(1, 1, -1))
|
||||
rand_ini = torch.rand(1, 1, self.dim, device=f0.device)
|
||||
rand_ini[..., 0] = 0
|
||||
rad += rand_ini
|
||||
sines = torch.sin(2 * np.pi * rad)
|
||||
return sines
|
||||
|
||||
@torch.no_grad()
|
||||
def forward(self, f0, upp):
|
||||
""" sine_tensor, uv = forward(f0)
|
||||
input F0: tensor(batchsize=1, length, dim=1)
|
||||
f0 for unvoiced steps should be 0
|
||||
output sine_tensor: tensor(batchsize=1, length, dim)
|
||||
output uv: tensor(batchsize=1, length, 1)
|
||||
"""
|
||||
f0 = f0.unsqueeze(-1)
|
||||
sine_waves = self._f02sine(f0, upp) * self.sine_amp
|
||||
uv = (f0 > self.voiced_threshold).float()
|
||||
uv = F.interpolate(uv.transpose(2, 1), scale_factor=upp, mode='nearest').transpose(2, 1)
|
||||
noise_amp = uv * self.noise_std + (1 - uv) * self.sine_amp / 3
|
||||
noise = noise_amp * torch.randn_like(sine_waves)
|
||||
sine_waves = sine_waves * uv + noise
|
||||
return sine_waves
|
||||
|
||||
|
||||
class SourceModuleHnNSF(torch.nn.Module):
|
||||
""" SourceModule for hn-nsf
|
||||
SourceModule(sampling_rate, harmonic_num=0, sine_amp=0.1,
|
||||
add_noise_std=0.003, voiced_threshod=0)
|
||||
sampling_rate: sampling_rate in Hz
|
||||
harmonic_num: number of harmonic above F0 (default: 0)
|
||||
sine_amp: amplitude of sine source signal (default: 0.1)
|
||||
add_noise_std: std of additive Gaussian noise (default: 0.003)
|
||||
note that amplitude of noise in unvoiced is decided
|
||||
by sine_amp
|
||||
voiced_threshold: threhold to set U/V given F0 (default: 0)
|
||||
Sine_source, noise_source = SourceModuleHnNSF(F0_sampled)
|
||||
F0_sampled (batchsize, length, 1)
|
||||
Sine_source (batchsize, length, 1)
|
||||
noise_source (batchsize, length 1)
|
||||
uv (batchsize, length, 1)
|
||||
"""
|
||||
|
||||
def __init__(self, sampling_rate, harmonic_num=0, sine_amp=0.1,
|
||||
add_noise_std=0.003, voiced_threshold=0):
|
||||
super(SourceModuleHnNSF, self).__init__()
|
||||
|
||||
self.sine_amp = sine_amp
|
||||
self.noise_std = add_noise_std
|
||||
|
||||
# to produce sine waveforms
|
||||
self.l_sin_gen = SineGen(sampling_rate, harmonic_num,
|
||||
sine_amp, add_noise_std, voiced_threshold)
|
||||
|
||||
# to merge source harmonics into a single excitation
|
||||
self.l_linear = torch.nn.Linear(harmonic_num + 1, 1)
|
||||
self.l_tanh = torch.nn.Tanh()
|
||||
|
||||
def forward(self, x, upp):
|
||||
sine_wavs = self.l_sin_gen(x, upp)
|
||||
sine_merge = self.l_tanh(self.l_linear(sine_wavs))
|
||||
return sine_merge
|
||||
|
||||
|
||||
class Generator(torch.nn.Module):
|
||||
def __init__(self, h):
|
||||
super(Generator, self).__init__()
|
||||
self.h = h
|
||||
self.num_kernels = len(h.resblock_kernel_sizes)
|
||||
self.num_upsamples = len(h.upsample_rates)
|
||||
self.mini_nsf = h.mini_nsf
|
||||
self.noise_sigma = h.noise_sigma
|
||||
|
||||
if h.mini_nsf:
|
||||
self.source_sr = h.sampling_rate / int(np.prod(h.upsample_rates[2: ]))
|
||||
self.upp = int(np.prod(h.upsample_rates[: 2]))
|
||||
else:
|
||||
self.source_sr = h.sampling_rate
|
||||
self.upp = int(np.prod(h.upsample_rates))
|
||||
self.m_source = SourceModuleHnNSF(
|
||||
sampling_rate=h.sampling_rate,
|
||||
harmonic_num=8
|
||||
)
|
||||
self.noise_convs = nn.ModuleList()
|
||||
|
||||
self.conv_pre = weight_norm(Conv1d(h.num_mels, h.upsample_initial_channel, 7, 1, padding=3))
|
||||
|
||||
self.ups = nn.ModuleList()
|
||||
self.resblocks = nn.ModuleList()
|
||||
resblock = ResBlock1 if h.resblock == '1' else ResBlock2
|
||||
ch = h.upsample_initial_channel
|
||||
for i, (u, k) in enumerate(zip(h.upsample_rates, h.upsample_kernel_sizes)):
|
||||
ch //= 2
|
||||
self.ups.append(weight_norm(ConvTranspose1d(2 * ch, ch, k, u, padding=(k - u) // 2)))
|
||||
for j, (k, d) in enumerate(zip(h.resblock_kernel_sizes, h.resblock_dilation_sizes)):
|
||||
self.resblocks.append(resblock(h, ch, k, d))
|
||||
if not h.mini_nsf:
|
||||
if i + 1 < len(h.upsample_rates): #
|
||||
stride_f0 = int(np.prod(h.upsample_rates[i + 1:]))
|
||||
self.noise_convs.append(Conv1d(
|
||||
1, ch, kernel_size=stride_f0 * 2, stride=stride_f0, padding=stride_f0 // 2))
|
||||
else:
|
||||
self.noise_convs.append(Conv1d(1, ch, kernel_size=1))
|
||||
elif i == 1:
|
||||
self.source_conv = Conv1d(1, ch, 1)
|
||||
self.source_conv.apply(init_weights)
|
||||
|
||||
self.conv_post = weight_norm(Conv1d(ch, 1, 7, 1, padding=3))
|
||||
|
||||
self.ups.apply(init_weights)
|
||||
self.conv_post.apply(init_weights)
|
||||
|
||||
def fastsinegen(self, f0):
|
||||
n = torch.arange(1, self.upp + 1, device=f0.device)
|
||||
s0 = f0.unsqueeze(-1) / self.source_sr
|
||||
ds0 = F.pad(s0[:, 1:, :] - s0[:, :-1, :], (0, 0, 0, 1))
|
||||
rad = s0 * n + 0.5 * ds0 * n * (n - 1) / self.upp
|
||||
rad2 = torch.fmod(rad[..., -1:].float() + 0.5, 1.0) - 0.5
|
||||
rad_acc = rad2.cumsum(dim=1).fmod(1.0).to(f0)
|
||||
rad += F.pad(rad_acc[:, :-1, :], (0, 0, 1, 0))
|
||||
rad = rad.reshape(f0.shape[0], 1, -1)
|
||||
sines = torch.sin(2 * np.pi * rad)
|
||||
return sines
|
||||
|
||||
def forward(self, x, f0):
|
||||
if self.mini_nsf:
|
||||
har_source = self.fastsinegen(f0)
|
||||
else:
|
||||
har_source = self.m_source(f0, self.upp).transpose(1, 2)
|
||||
x = self.conv_pre(x)
|
||||
if self.noise_sigma is not None and self.noise_sigma > 0:
|
||||
x += self.noise_sigma * torch.randn_like(x)
|
||||
for i in range(self.num_upsamples):
|
||||
x = F.leaky_relu(x, LRELU_SLOPE)
|
||||
x = self.ups[i](x)
|
||||
if not self.mini_nsf:
|
||||
x_source = self.noise_convs[i](har_source)
|
||||
x = x + x_source
|
||||
elif i == 1:
|
||||
x_source = self.source_conv(har_source)
|
||||
x = x + x_source
|
||||
xs = None
|
||||
for j in range(self.num_kernels):
|
||||
if xs is None:
|
||||
xs = self.resblocks[i * self.num_kernels + j](x)
|
||||
else:
|
||||
xs += self.resblocks[i * self.num_kernels + j](x)
|
||||
x = xs / self.num_kernels
|
||||
x = F.leaky_relu(x)
|
||||
x = self.conv_post(x)
|
||||
x = torch.tanh(x)
|
||||
return x
|
||||
|
||||
def remove_weight_norm(self):
|
||||
# rank_zero_info('Removing weight norm...')
|
||||
print('Removing weight norm...')
|
||||
for l in self.ups:
|
||||
remove_weight_norm(l)
|
||||
for l in self.resblocks:
|
||||
l.remove_weight_norm()
|
||||
remove_weight_norm(self.conv_pre)
|
||||
remove_weight_norm(self.conv_post)
|
||||
@@ -0,0 +1,87 @@
|
||||
import os
|
||||
|
||||
os.environ["LRU_CACHE_CAPACITY"] = "3"
|
||||
import torch
|
||||
import torch.utils.data
|
||||
import numpy as np
|
||||
from librosa.filters import mel as librosa_mel_fn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
def dynamic_range_compression(x, C=1, clip_val=1e-5):
|
||||
return np.log(np.clip(x, a_min=clip_val, a_max=None) * C)
|
||||
|
||||
|
||||
def dynamic_range_decompression(x, C=1):
|
||||
return np.exp(x) / C
|
||||
|
||||
|
||||
def dynamic_range_compression_torch(x, C=1, clip_val=1e-5):
|
||||
return torch.log(torch.clamp(x, min=clip_val) * C)
|
||||
|
||||
|
||||
def dynamic_range_decompression_torch(x, C=1):
|
||||
return torch.exp(x) / C
|
||||
|
||||
|
||||
class STFT:
|
||||
def __init__(
|
||||
self, sr=22050,
|
||||
n_mels=80, n_fft=1024, win_size=1024, hop_length=256,
|
||||
fmin=20, fmax=11025, clip_val=1e-5,
|
||||
device=None
|
||||
):
|
||||
self.target_sr = sr
|
||||
|
||||
self.n_mels = n_mels
|
||||
self.n_fft = n_fft
|
||||
self.win_size = win_size
|
||||
self.hop_length = hop_length
|
||||
self.fmin = fmin
|
||||
self.fmax = fmax
|
||||
self.clip_val = clip_val
|
||||
|
||||
if device is None:
|
||||
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||||
self.device = device
|
||||
|
||||
mel_basis = librosa_mel_fn(sr=sr, n_fft=n_fft, n_mels=n_mels, fmin=fmin, fmax=fmax)
|
||||
self.mel_basis = torch.from_numpy(mel_basis).float().to(device)
|
||||
|
||||
def get_mel(self, y, keyshift=0, speed=1, center=False):
|
||||
|
||||
factor = 2 ** (keyshift / 12)
|
||||
n_fft_new = int(np.round(self.n_fft * factor))
|
||||
win_size_new = int(np.round(self.win_size * factor))
|
||||
hop_length_new = int(np.round(self.hop_length * speed))
|
||||
|
||||
if torch.min(y) < -1.:
|
||||
print('min value is ', torch.min(y))
|
||||
if torch.max(y) > 1.:
|
||||
print('max value is ', torch.max(y))
|
||||
|
||||
window = torch.hann_window(win_size_new, device=self.device)
|
||||
|
||||
y = torch.nn.functional.pad(y.unsqueeze(1), (
|
||||
(win_size_new - hop_length_new) // 2,
|
||||
(win_size_new - hop_length_new + 1) // 2
|
||||
), mode='reflect')
|
||||
y = y.squeeze(1)
|
||||
|
||||
spec = torch.stft(
|
||||
y, n_fft_new, hop_length=hop_length_new,
|
||||
win_length=win_size_new, window=window,
|
||||
center=center, pad_mode='reflect',
|
||||
normalized=False, onesided=True, return_complex=True
|
||||
).abs()
|
||||
if keyshift != 0:
|
||||
size = self.n_fft // 2 + 1
|
||||
resize = spec.size(1)
|
||||
if resize < size:
|
||||
spec = F.pad(spec, (0, 0, 0, size - resize))
|
||||
spec = spec[:, :size, :] * self.win_size / win_size_new
|
||||
|
||||
spec = torch.matmul(self.mel_basis, spec)
|
||||
spec = dynamic_range_compression_torch(spec, clip_val=self.clip_val)
|
||||
|
||||
return spec
|
||||
@@ -0,0 +1,13 @@
|
||||
import matplotlib
|
||||
|
||||
matplotlib.use("Agg")
|
||||
|
||||
|
||||
def init_weights(m, mean=0.0, std=0.01):
|
||||
classname = m.__class__.__name__
|
||||
if classname.find("Conv") != -1:
|
||||
m.weight.data.normal_(mean, std)
|
||||
|
||||
|
||||
def get_padding(kernel_size, dilation=1):
|
||||
return int((kernel_size*dilation - dilation)/2)
|
||||
@@ -0,0 +1,132 @@
|
||||
import torch
|
||||
from torch import Tensor
|
||||
from torch.optim import Optimizer
|
||||
from torch.optim.optimizer import ParamsT
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, List, Type, Callable, Optional
|
||||
|
||||
|
||||
@dataclass
|
||||
class OptimizerSpec:
|
||||
"""Spec for creating an optimizer that is part of a `ChainedOptimizer`."""
|
||||
|
||||
class_type: Type[Optimizer]
|
||||
init_args: Dict[str, Any]
|
||||
param_filter: Optional[Callable[[Tensor], bool]]
|
||||
|
||||
|
||||
class ChainedOptimizer(Optimizer):
|
||||
"""
|
||||
A wrapper around multiple optimizers that allows for chaining them together.
|
||||
The optimizers are applied in the order they are passed in the constructor.
|
||||
Each optimizer is responsible for updating a subset of the parameters, which
|
||||
is determined by the `param_filter` function. If no optimizer is found for a
|
||||
parameter group, an exception is raised.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
params: ParamsT,
|
||||
optimizer_specs: List[OptimizerSpec],
|
||||
lr: float,
|
||||
weight_decay: float = 0.0,
|
||||
optimizer_selection_callback: Optional[Callable[[Tensor, int], None]] = None,
|
||||
**common_kwargs,
|
||||
):
|
||||
self.optimizer_specs = optimizer_specs
|
||||
self.optimizer_selection_callback = optimizer_selection_callback
|
||||
self.optimizers: List[Optimizer] = []
|
||||
defaults = dict(lr=lr, weight_decay=weight_decay)
|
||||
super().__init__(params, defaults)
|
||||
|
||||
# Split the params for each optimizer
|
||||
params_for_optimizers = [[] for _ in optimizer_specs]
|
||||
for param_group in self.param_groups:
|
||||
params = param_group["params"]
|
||||
indices = param_group["optimizer_and_param_group_indices"] = set()
|
||||
for param in params:
|
||||
assert isinstance(param, Tensor), f"Expected a Tensor, got {type(param)}"
|
||||
found_optimizer = False
|
||||
for index, spec in enumerate(optimizer_specs):
|
||||
if spec.param_filter is None or spec.param_filter(param):
|
||||
if self.optimizer_selection_callback is not None:
|
||||
self.optimizer_selection_callback(param, index)
|
||||
params_for_optimizers[index].append(param)
|
||||
indices.add((index, 0))
|
||||
found_optimizer = True
|
||||
break
|
||||
if not found_optimizer:
|
||||
raise ValueError("No valid optimizer found for the given parameter")
|
||||
|
||||
# Initialize the optimizers
|
||||
for spec, selected_params in zip(optimizer_specs, params_for_optimizers):
|
||||
optimizer_args = {
|
||||
'lr': lr,
|
||||
'weight_decay': weight_decay,
|
||||
}
|
||||
optimizer_args.update(common_kwargs)
|
||||
optimizer_args.update(spec.init_args)
|
||||
optimizer = spec.class_type(selected_params, **optimizer_args)
|
||||
self.optimizers.append(optimizer)
|
||||
|
||||
def state_dict(self) -> Dict[str, Any]:
|
||||
return {
|
||||
"optimizers": [opt.state_dict() for opt in self.optimizers],
|
||||
**super().state_dict(),
|
||||
}
|
||||
|
||||
def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
|
||||
optimizers = state_dict.pop("optimizers")
|
||||
super().load_state_dict(state_dict)
|
||||
for i in range(len(self.optimizers)):
|
||||
self.optimizers[i].load_state_dict(optimizers[i])
|
||||
|
||||
def zero_grad(self, set_to_none: bool = True) -> None:
|
||||
for opt in self.optimizers:
|
||||
opt.zero_grad(set_to_none=set_to_none)
|
||||
|
||||
def _copy_lr_to_optimizers(self) -> None:
|
||||
for param_group in self.param_groups:
|
||||
indices = param_group["optimizer_and_param_group_indices"]
|
||||
for optimizer_idx, param_group_idx in indices:
|
||||
self.optimizers[optimizer_idx].param_groups[param_group_idx]["lr"] = param_group["lr"]
|
||||
|
||||
def step(self, closure=None) -> None:
|
||||
loss = None
|
||||
if closure is not None:
|
||||
with torch.enable_grad():
|
||||
loss = closure()
|
||||
self._copy_lr_to_optimizers()
|
||||
for opt in self.optimizers:
|
||||
opt.step(closure=None)
|
||||
return loss
|
||||
|
||||
def add_param_group(self, param_group: Dict[str, Any]) -> None:
|
||||
super().add_param_group(param_group)
|
||||
|
||||
# If optimizer has not been initialized, skip adding the param groups
|
||||
if not self.optimizers:
|
||||
return
|
||||
|
||||
# Split the params for each optimizer
|
||||
params_for_optimizers = [[] for _ in self.optimizer_specs]
|
||||
params = param_group["params"]
|
||||
indices = param_group["optimizer_and_param_group_indices"] = set()
|
||||
for param in params:
|
||||
assert isinstance(param, Tensor), f"Expected a Tensor, got {type(param)}"
|
||||
found_optimizer = False
|
||||
for index, spec in enumerate(self.optimizer_specs):
|
||||
if spec.param_filter is None or spec.param_filter(param):
|
||||
if self.optimizer_selection_callback is not None:
|
||||
self.optimizer_selection_callback(param, index)
|
||||
params_for_optimizers[index].append(param)
|
||||
indices.add((index, len(self.optimizers[index].param_groups)))
|
||||
found_optimizer = True
|
||||
break
|
||||
if not found_optimizer:
|
||||
raise ValueError("No valid optimizer found for the given parameter group")
|
||||
|
||||
# Add the selected param group to the optimizers
|
||||
for optimizer, selected_params in zip(self.optimizers, params_for_optimizers):
|
||||
if selected_params:
|
||||
optimizer.add_param_group({"params": selected_params})
|
||||
@@ -0,0 +1,200 @@
|
||||
import collections
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from torch import Tensor
|
||||
from torch.nn import Parameter
|
||||
from typing import List
|
||||
from .chained_optimizer import ChainedOptimizer, OptimizerSpec
|
||||
|
||||
from modules.commons.common_layers import AdamWLinear, AdamWConv1d
|
||||
|
||||
|
||||
def zeropower_via_newtonschulz5(G: Tensor, steps: int) -> Tensor:
|
||||
"""
|
||||
Newton-Schulz iteration to compute the zeroth power / orthogonalization of G. We opt to use a
|
||||
quintic iteration whose coefficients are selected to maximize the slope at zero. For the purpose
|
||||
of minimizing steps, it turns out to be empirically effective to keep increasing the slope at
|
||||
zero even beyond the point where the iteration no longer converges all the way to one everywhere
|
||||
on the interval. This iteration therefore does not produce UV^T but rather something like US'V^T
|
||||
where S' is diagonal with S_{ii}' ~ Uniform(0.5, 1.5), which turns out not to hurt model
|
||||
performance at all relative to UV^T, where USV^T = G is the SVD.
|
||||
"""
|
||||
assert G.ndim == 3 # batched Muon implementation by @scottjmaddox, and put into practice in the record by @YouJiacheng
|
||||
a, b, c = (3.4445, -4.7750, 2.0315)
|
||||
|
||||
X = G.to(torch.float32)
|
||||
|
||||
# Ensure spectral norm is at most 1
|
||||
X = F.normalize(X, p=2.0, dim=(-2, -1), eps=1e-7)
|
||||
|
||||
X = X.to(torch.float16)
|
||||
|
||||
# Perform the NS iterations
|
||||
if X.size(-2) < X.size(-1):
|
||||
for _ in range(steps):
|
||||
A = torch.bmm(X, X.mT)
|
||||
A = torch.baddbmm(A, A, A, beta=b, alpha=c)
|
||||
X = torch.baddbmm(X, A, X, beta=a, alpha=1)
|
||||
else:
|
||||
for _ in range(steps):
|
||||
A = torch.bmm(X.mT, X)
|
||||
A = torch.baddbmm(A, A, A, beta=b, alpha=c)
|
||||
X = torch.baddbmm(X, X, A, beta=a, alpha=1)
|
||||
|
||||
return X
|
||||
|
||||
|
||||
def gram_newton_schulz(G: Tensor, steps: int) -> Tensor:
|
||||
"""
|
||||
Refer to:
|
||||
Gram Newton-Schulz: A Fast, Hardware-Aware Newton-Schulz Algorithm for Muon
|
||||
Authors: Jack Zhang, Noah Amsel, Berlin Chen, Tri Dao
|
||||
Blogpost: https://dao-ailab.github.io/blog/2026/gram-newton-schulz/
|
||||
|
||||
Gram Newton-Schulz iteration to compute the orthogonalization of G.
|
||||
Mathematically identical to standard Newton-Schulz but computes iterating
|
||||
on the smaller NxN Gram matrix to save up to 50% FLOPs.
|
||||
"""
|
||||
assert G.ndim == 3
|
||||
reset_iterations = [2]
|
||||
original_shape = G.shape
|
||||
dtype = G.dtype
|
||||
|
||||
X = G.to(torch.float32)
|
||||
X = F.normalize(X, p=2.0, dim=(-2, -1), eps=1e-7)
|
||||
should_transpose = X.size(-2) > X.size(-1)
|
||||
if should_transpose:
|
||||
X = X.mT
|
||||
X = X.to(torch.float16)
|
||||
|
||||
a, b, c = (3.4445, -4.7750, 2.0315)
|
||||
|
||||
if X.size(-2) != X.size(-1):
|
||||
R = torch.bmm(X, X.mT)
|
||||
Q = None
|
||||
for i in range(steps):
|
||||
if i in reset_iterations and i != 0:
|
||||
X = torch.bmm(Q, X)
|
||||
R = torch.bmm(X, X.mT)
|
||||
Q = None
|
||||
Z = torch.baddbmm(R, R, R, beta=b, alpha=c)
|
||||
if i != 0 and i not in reset_iterations:
|
||||
Q = torch.baddbmm(Q, Q, Z, beta=a, alpha=1.0)
|
||||
else:
|
||||
Q = Z.clone()
|
||||
Q.diagonal(dim1=-2, dim2=-1).add_(a)
|
||||
if i < steps - 1 and (i + 1) not in reset_iterations:
|
||||
RZ = torch.baddbmm(R, R, Z, beta=a, alpha=1.0)
|
||||
R = torch.baddbmm(RZ, Z, RZ, beta=a, alpha=1.0)
|
||||
X = torch.bmm(Q, X) if not should_transpose else torch.bmm(X.mT, Q)
|
||||
else:
|
||||
for _ in range(steps):
|
||||
A = torch.bmm(X, X.mT)
|
||||
B = torch.baddbmm(A, A, A, beta=b, alpha=c)
|
||||
X = torch.baddbmm(X, B, X, beta=a, alpha=1.0)
|
||||
|
||||
return X.to(dtype).view(original_shape)
|
||||
|
||||
|
||||
class Muon(torch.optim.Optimizer):
|
||||
"""
|
||||
Muon - MomentUm Orthogonalized by Newton-schulz
|
||||
|
||||
https://kellerjordan.github.io/posts/muon/
|
||||
|
||||
Muon internally runs standard SGD-momentum, and then performs an orthogonalization post-
|
||||
processing step, in which each 2D parameter's update is replaced with the nearest orthogonal
|
||||
matrix. To efficiently orthogonalize each update, we use a Newton-Schulz iteration, which has
|
||||
the advantage that it can be stably run in float16 on the GPU.
|
||||
|
||||
Some warnings:
|
||||
- This optimizer should not be used for the embedding layer, the final fully connected layer,
|
||||
or any {0,1}-D parameters; those should all be optimized by a standard method (e.g., AdamW).
|
||||
- To use it with 4D convolutional filters, it works well to just flatten their last 3 dimensions.
|
||||
|
||||
Arguments:
|
||||
lr: The learning rate used by the internal SGD.
|
||||
momentum: The momentum used by the internal SGD.
|
||||
nesterov: Whether to use Nesterov-style momentum in the internal SGD. (recommended)
|
||||
ns_steps: The number of Newton-Schulz iteration steps to use.
|
||||
"""
|
||||
|
||||
def __init__(self, params, lr=5e-4, weight_decay=0.1, momentum=0.95, nesterov=True, ns_steps=5):
|
||||
defaults = dict(lr=lr, weight_decay=weight_decay, momentum=momentum, nesterov=nesterov, ns_steps=ns_steps)
|
||||
super().__init__(params, defaults)
|
||||
|
||||
@torch.no_grad()
|
||||
def step(self, closure=None):
|
||||
for group in self.param_groups:
|
||||
shape_groups = {}
|
||||
for p in filter(lambda p: p.grad is not None, group["params"]):
|
||||
g = p.grad
|
||||
state = self.state[p]
|
||||
if "momentum_buffer" not in state:
|
||||
state["momentum_buffer"] = torch.zeros_like(g)
|
||||
key = (p.shape, p.device, p.dtype)
|
||||
if key not in shape_groups:
|
||||
shape_groups[key] = {"params": [], "grads": [], "buffers": []}
|
||||
shape_groups[key]["params"].append(p)
|
||||
shape_groups[key]["grads"].append(g)
|
||||
shape_groups[key]["buffers"].append(state["momentum_buffer"])
|
||||
for key in shape_groups:
|
||||
group_data = shape_groups[key]
|
||||
p, g, buf, m = group_data["params"], group_data["grads"], group_data["buffers"], group["momentum"]
|
||||
torch._foreach_lerp_(buf, g, 1-m)
|
||||
if group["nesterov"]:
|
||||
torch._foreach_lerp_(g, buf, m)
|
||||
g = torch.stack(g)
|
||||
else:
|
||||
g = torch.stack(buf)
|
||||
original_shape = g.shape
|
||||
if g.ndim >= 4: # for the case of conv filters
|
||||
g = g.view(g.size(0), g.size(1), -1)
|
||||
g = gram_newton_schulz(g, steps=group["ns_steps"])
|
||||
|
||||
if group["weight_decay"] > 0:
|
||||
torch._foreach_mul_(p, 1 - group["lr"] * group["weight_decay"])
|
||||
torch._foreach_add_(p, g.view(original_shape).unbind(0), alpha=-group["lr"] * max(g[0].size()) ** 0.5)
|
||||
|
||||
|
||||
def get_params_for_muon(model) -> List[Parameter]:
|
||||
"""
|
||||
Filter parameters of a module into two groups: those that can be optimized by Muon,
|
||||
and those that should be optimized by a standard optimizer.
|
||||
Args:
|
||||
module: The module to filter parameters for.
|
||||
Returns:
|
||||
A list of parameters that should be optimized with muon.
|
||||
"""
|
||||
excluded_module_classes = (nn.Embedding, AdamWLinear, AdamWConv1d)
|
||||
muon_params = []
|
||||
# BFS through all submodules and exclude parameters from certain module types
|
||||
queue = collections.deque([model])
|
||||
while queue:
|
||||
module = queue.popleft()
|
||||
if isinstance(module, excluded_module_classes):
|
||||
continue
|
||||
for param in module.parameters(recurse=False):
|
||||
if not param.requires_grad:
|
||||
continue
|
||||
if param.ndim >= 2:
|
||||
muon_params.append(param)
|
||||
queue.extend(list(module.children()))
|
||||
return muon_params
|
||||
|
||||
|
||||
class Muon_AdamW(ChainedOptimizer):
|
||||
def __init__(self, model, lr=0.0005, weight_decay=0.0, muon_args=None, adamw_args=None, verbose=False):
|
||||
muon_args = {} if muon_args is None else muon_args
|
||||
adamw_args = {} if adamw_args is None else adamw_args
|
||||
muon_params_id_set = set(id(p) for p in get_params_for_muon(model))
|
||||
spec_muon = OptimizerSpec(Muon, muon_args, lambda param: id(param) in muon_params_id_set)
|
||||
spec_adamw = OptimizerSpec(torch.optim.AdamW, adamw_args, None)
|
||||
specs = [spec_muon, spec_adamw]
|
||||
callback = None
|
||||
if verbose:
|
||||
callback = lambda p, spec_idx: print(
|
||||
f"Adding param {p.shape} to optimizer{spec_idx} {str(specs[spec_idx].class_type)}"
|
||||
)
|
||||
super().__init__(model.parameters(), specs, lr=lr, weight_decay=weight_decay, optimizer_selection_callback=callback)
|
||||
@@ -0,0 +1,18 @@
|
||||
from utils import hparams
|
||||
|
||||
from .pm import ParselmouthPE
|
||||
from .pw import HarvestPE
|
||||
from .rmvpe import RMVPE
|
||||
|
||||
|
||||
def initialize_pe():
|
||||
pe = hparams['pe']
|
||||
pe_ckpt = hparams['pe_ckpt']
|
||||
if pe == 'parselmouth':
|
||||
return ParselmouthPE()
|
||||
elif pe == 'rmvpe':
|
||||
return RMVPE(pe_ckpt)
|
||||
elif pe == 'harvest':
|
||||
return HarvestPE()
|
||||
else:
|
||||
raise ValueError(f" [x] Unknown f0 extractor: {pe}")
|
||||
@@ -0,0 +1,15 @@
|
||||
from basics.base_pe import BasePE
|
||||
from utils.binarizer_utils import get_pitch_parselmouth
|
||||
|
||||
|
||||
class ParselmouthPE(BasePE):
|
||||
def get_pitch(
|
||||
self,waveform, samplerate, length,
|
||||
*, hop_size, f0_min=65, f0_max=1100,
|
||||
speed=1, interp_uv=False
|
||||
):
|
||||
return get_pitch_parselmouth(
|
||||
waveform, samplerate=samplerate, length=length,
|
||||
hop_size=hop_size, f0_min=f0_min, f0_max=f0_max,
|
||||
speed=speed, interp_uv=interp_uv
|
||||
)
|
||||
@@ -0,0 +1,29 @@
|
||||
from basics.base_pe import BasePE
|
||||
import numpy as np
|
||||
import pyworld as pw
|
||||
from utils.pitch_utils import interp_f0
|
||||
|
||||
|
||||
class HarvestPE(BasePE):
|
||||
def get_pitch(
|
||||
self, waveform, samplerate, length,
|
||||
*, hop_size, f0_min=65, f0_max=1100,
|
||||
speed=1, interp_uv=False
|
||||
):
|
||||
hop_size = int(np.round(hop_size * speed))
|
||||
time_step = 1000 * hop_size / samplerate
|
||||
|
||||
f0, _ = pw.harvest(
|
||||
waveform.astype(np.float64), samplerate,
|
||||
f0_floor=f0_min, f0_ceil=f0_max, frame_period=time_step
|
||||
)
|
||||
f0 = f0.astype(np.float32)
|
||||
|
||||
if f0.size < length:
|
||||
f0 = np.pad(f0, (0, length - f0.size))
|
||||
f0 = f0[:length]
|
||||
uv = f0 == 0
|
||||
|
||||
if interp_uv:
|
||||
f0, uv = interp_f0(f0, uv)
|
||||
return f0, uv
|
||||
@@ -0,0 +1,5 @@
|
||||
from .constants import *
|
||||
from .model import E2E0
|
||||
from .utils import to_local_average_f0, to_viterbi_f0
|
||||
from .inference import RMVPE
|
||||
from .spec import MelSpectrogram
|
||||
@@ -0,0 +1,9 @@
|
||||
SAMPLE_RATE = 16000
|
||||
|
||||
N_CLASS = 360
|
||||
|
||||
N_MELS = 128
|
||||
MEL_FMIN = 30
|
||||
MEL_FMAX = 8000
|
||||
WINDOW_LENGTH = 1024
|
||||
CONST = 1997.3794084376191
|
||||
@@ -0,0 +1,173 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from .constants import N_MELS
|
||||
|
||||
|
||||
class ConvBlockRes(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, momentum=0.01):
|
||||
super(ConvBlockRes, self).__init__()
|
||||
self.conv = nn.Sequential(
|
||||
nn.Conv2d(in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
kernel_size=(3, 3),
|
||||
stride=(1, 1),
|
||||
padding=(1, 1),
|
||||
bias=False),
|
||||
nn.BatchNorm2d(out_channels, momentum=momentum),
|
||||
nn.ReLU(),
|
||||
|
||||
nn.Conv2d(in_channels=out_channels,
|
||||
out_channels=out_channels,
|
||||
kernel_size=(3, 3),
|
||||
stride=(1, 1),
|
||||
padding=(1, 1),
|
||||
bias=False),
|
||||
nn.BatchNorm2d(out_channels, momentum=momentum),
|
||||
nn.ReLU(),
|
||||
)
|
||||
if in_channels != out_channels:
|
||||
self.shortcut = nn.Conv2d(in_channels, out_channels, (1, 1))
|
||||
self.is_shortcut = True
|
||||
else:
|
||||
self.is_shortcut = False
|
||||
|
||||
def forward(self, x):
|
||||
if self.is_shortcut:
|
||||
return self.conv(x) + self.shortcut(x)
|
||||
else:
|
||||
return self.conv(x) + x
|
||||
|
||||
|
||||
class ResEncoderBlock(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, kernel_size, n_blocks=1, momentum=0.01):
|
||||
super(ResEncoderBlock, self).__init__()
|
||||
self.n_blocks = n_blocks
|
||||
self.conv = nn.ModuleList()
|
||||
self.conv.append(ConvBlockRes(in_channels, out_channels, momentum))
|
||||
for i in range(n_blocks - 1):
|
||||
self.conv.append(ConvBlockRes(out_channels, out_channels, momentum))
|
||||
self.kernel_size = kernel_size
|
||||
if self.kernel_size is not None:
|
||||
self.pool = nn.AvgPool2d(kernel_size=kernel_size)
|
||||
|
||||
def forward(self, x):
|
||||
for i in range(self.n_blocks):
|
||||
x = self.conv[i](x)
|
||||
if self.kernel_size is not None:
|
||||
return x, self.pool(x)
|
||||
else:
|
||||
return x
|
||||
|
||||
|
||||
class ResDecoderBlock(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, stride, n_blocks=1, momentum=0.01):
|
||||
super(ResDecoderBlock, self).__init__()
|
||||
out_padding = (0, 1) if stride == (1, 2) else (1, 1)
|
||||
self.n_blocks = n_blocks
|
||||
self.conv1 = nn.Sequential(
|
||||
nn.ConvTranspose2d(in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
kernel_size=(3, 3),
|
||||
stride=stride,
|
||||
padding=(1, 1),
|
||||
output_padding=out_padding,
|
||||
bias=False),
|
||||
nn.BatchNorm2d(out_channels, momentum=momentum),
|
||||
nn.ReLU(),
|
||||
)
|
||||
self.conv2 = nn.ModuleList()
|
||||
self.conv2.append(ConvBlockRes(out_channels * 2, out_channels, momentum))
|
||||
for i in range(n_blocks-1):
|
||||
self.conv2.append(ConvBlockRes(out_channels, out_channels, momentum))
|
||||
|
||||
def forward(self, x, concat_tensor):
|
||||
x = self.conv1(x)
|
||||
x = torch.cat((x, concat_tensor), dim=1)
|
||||
for i in range(self.n_blocks):
|
||||
x = self.conv2[i](x)
|
||||
return x
|
||||
|
||||
|
||||
class Encoder(nn.Module):
|
||||
def __init__(self, in_channels, in_size, n_encoders, kernel_size, n_blocks, out_channels=16, momentum=0.01):
|
||||
super(Encoder, self).__init__()
|
||||
self.n_encoders = n_encoders
|
||||
self.bn = nn.BatchNorm2d(in_channels, momentum=momentum)
|
||||
self.layers = nn.ModuleList()
|
||||
self.latent_channels = []
|
||||
for i in range(self.n_encoders):
|
||||
self.layers.append(ResEncoderBlock(in_channels, out_channels, kernel_size, n_blocks, momentum=momentum))
|
||||
self.latent_channels.append([out_channels, in_size])
|
||||
in_channels = out_channels
|
||||
out_channels *= 2
|
||||
in_size //= 2
|
||||
self.out_size = in_size
|
||||
self.out_channel = out_channels
|
||||
|
||||
def forward(self, x):
|
||||
concat_tensors = []
|
||||
x = self.bn(x)
|
||||
for i in range(self.n_encoders):
|
||||
_, x = self.layers[i](x)
|
||||
concat_tensors.append(_)
|
||||
return x, concat_tensors
|
||||
|
||||
|
||||
class Intermediate(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, n_inters, n_blocks, momentum=0.01):
|
||||
super(Intermediate, self).__init__()
|
||||
self.n_inters = n_inters
|
||||
self.layers = nn.ModuleList()
|
||||
self.layers.append(ResEncoderBlock(in_channels, out_channels, None, n_blocks, momentum))
|
||||
for i in range(self.n_inters-1):
|
||||
self.layers.append(ResEncoderBlock(out_channels, out_channels, None, n_blocks, momentum))
|
||||
|
||||
def forward(self, x):
|
||||
for i in range(self.n_inters):
|
||||
x = self.layers[i](x)
|
||||
return x
|
||||
|
||||
|
||||
class Decoder(nn.Module):
|
||||
def __init__(self, in_channels, n_decoders, stride, n_blocks, momentum=0.01):
|
||||
super(Decoder, self).__init__()
|
||||
self.layers = nn.ModuleList()
|
||||
self.n_decoders = n_decoders
|
||||
for i in range(self.n_decoders):
|
||||
out_channels = in_channels // 2
|
||||
self.layers.append(ResDecoderBlock(in_channels, out_channels, stride, n_blocks, momentum))
|
||||
in_channels = out_channels
|
||||
|
||||
def forward(self, x, concat_tensors):
|
||||
for i in range(self.n_decoders):
|
||||
x = self.layers[i](x, concat_tensors[-1-i])
|
||||
return x
|
||||
|
||||
|
||||
class TimbreFilter(nn.Module):
|
||||
def __init__(self, latent_rep_channels):
|
||||
super(TimbreFilter, self).__init__()
|
||||
self.layers = nn.ModuleList()
|
||||
for latent_rep in latent_rep_channels:
|
||||
self.layers.append(ConvBlockRes(latent_rep[0], latent_rep[0]))
|
||||
|
||||
def forward(self, x_tensors):
|
||||
out_tensors = []
|
||||
for i, layer in enumerate(self.layers):
|
||||
out_tensors.append(layer(x_tensors[i]))
|
||||
return out_tensors
|
||||
|
||||
|
||||
class DeepUnet0(nn.Module):
|
||||
def __init__(self, kernel_size, n_blocks, en_de_layers=5, inter_layers=4, in_channels=1, en_out_channels=16):
|
||||
super(DeepUnet0, self).__init__()
|
||||
self.encoder = Encoder(in_channels, N_MELS, en_de_layers, kernel_size, n_blocks, en_out_channels)
|
||||
self.intermediate = Intermediate(self.encoder.out_channel // 2, self.encoder.out_channel, inter_layers, n_blocks)
|
||||
self.tf = TimbreFilter(self.encoder.latent_channels)
|
||||
self.decoder = Decoder(self.encoder.out_channel, en_de_layers, kernel_size, n_blocks)
|
||||
|
||||
def forward(self, x):
|
||||
x, concat_tensors = self.encoder(x)
|
||||
x = self.intermediate(x)
|
||||
x = self.decoder(x, concat_tensors)
|
||||
return x
|
||||
@@ -0,0 +1,78 @@
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torchaudio.transforms import Resample
|
||||
|
||||
from basics.base_pe import BasePE
|
||||
from utils.infer_utils import resample_align_curve
|
||||
from utils.pitch_utils import interp_f0
|
||||
from .constants import *
|
||||
from .model import E2E0
|
||||
from .spec import MelSpectrogram
|
||||
from .utils import to_local_average_f0, to_viterbi_f0
|
||||
|
||||
|
||||
class RMVPE(BasePE):
|
||||
def __init__(self, model_path, hop_length=160):
|
||||
self.resample_kernel = {}
|
||||
self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
||||
self.model = E2E0(4, 1, (2, 2)).eval().to(self.device)
|
||||
ckpt = torch.load(model_path, map_location=self.device)
|
||||
self.model.load_state_dict(ckpt['model'], strict=False)
|
||||
self.hop_length = hop_length
|
||||
self.seg_length = 32 * hop_length
|
||||
self.mel_extractor = MelSpectrogram(
|
||||
N_MELS, SAMPLE_RATE, WINDOW_LENGTH, hop_length, None, MEL_FMIN, MEL_FMAX
|
||||
).to(self.device)
|
||||
|
||||
@torch.no_grad()
|
||||
def mel2hidden(self, mel):
|
||||
n_frames = mel.shape[-1]
|
||||
mel = F.pad(mel, (0, 32 * ((n_frames - 1) // 32 + 1) - n_frames), mode='reflect')
|
||||
hidden = self.model(mel)
|
||||
return hidden[:, :n_frames]
|
||||
|
||||
def decode(self, hidden, thred=0.03, use_viterbi=False):
|
||||
if use_viterbi:
|
||||
f0 = to_viterbi_f0(hidden, thred=thred)
|
||||
else:
|
||||
f0 = to_local_average_f0(hidden, thred=thred)
|
||||
return f0
|
||||
|
||||
def infer_from_audio(self, audio, sample_rate=16000, thred=0.03, use_viterbi=False):
|
||||
audio = torch.from_numpy(audio).float().unsqueeze(0).to(self.device)
|
||||
if sample_rate == 16000:
|
||||
audio_res = audio
|
||||
else:
|
||||
key_str = str(sample_rate)
|
||||
if key_str not in self.resample_kernel:
|
||||
self.resample_kernel[key_str] = Resample(sample_rate, 16000, lowpass_filter_width=128)
|
||||
self.resample_kernel[key_str] = self.resample_kernel[key_str].to(self.device)
|
||||
audio_res = self.resample_kernel[key_str](audio)
|
||||
B, T = audio_res.shape
|
||||
n_frames = T // self.hop_length + 1
|
||||
T1 = T + self.hop_length
|
||||
T_pad = self.seg_length * ((T1 - 1) // self.seg_length + 1) - T1
|
||||
audio_res = F.pad(audio_res, (0, T_pad))
|
||||
mel = self.mel_extractor(audio_res, center=True)
|
||||
with torch.no_grad():
|
||||
hidden = self.model(mel)
|
||||
f0 = self.decode(hidden[:, :n_frames], thred=thred, use_viterbi=use_viterbi)
|
||||
return f0
|
||||
|
||||
def get_pitch(
|
||||
self, waveform, samplerate, length,
|
||||
*, hop_size, f0_min=65, f0_max=1100,
|
||||
speed=1, interp_uv=False
|
||||
):
|
||||
f0 = self.infer_from_audio(waveform, sample_rate=samplerate)
|
||||
uv = f0 == 0
|
||||
f0, uv = interp_f0(f0, uv)
|
||||
|
||||
hop_size = int(np.round(hop_size * speed))
|
||||
time_step = hop_size / samplerate
|
||||
f0_res = resample_align_curve(f0, 0.01, time_step, length)
|
||||
uv_res = resample_align_curve(uv.astype(np.float32), 0.01, time_step, length) > 0.5
|
||||
if not interp_uv:
|
||||
f0_res[uv_res] = 0
|
||||
return f0_res, uv_res
|
||||
@@ -0,0 +1,32 @@
|
||||
from torch import nn
|
||||
|
||||
from .constants import *
|
||||
from .deepunet import DeepUnet0
|
||||
from .seq import BiGRU
|
||||
|
||||
|
||||
class E2E0(nn.Module):
|
||||
def __init__(self, n_blocks, n_gru, kernel_size, en_de_layers=5, inter_layers=4, in_channels=1,
|
||||
en_out_channels=16):
|
||||
super(E2E0, self).__init__()
|
||||
self.unet = DeepUnet0(kernel_size, n_blocks, en_de_layers, inter_layers, in_channels, en_out_channels)
|
||||
self.cnn = nn.Conv2d(en_out_channels, 3, (3, 3), padding=(1, 1))
|
||||
if n_gru:
|
||||
self.fc = nn.Sequential(
|
||||
BiGRU(3 * N_MELS, 256, n_gru),
|
||||
nn.Linear(512, N_CLASS),
|
||||
nn.Dropout(0.25),
|
||||
nn.Sigmoid()
|
||||
)
|
||||
else:
|
||||
self.fc = nn.Sequential(
|
||||
nn.Linear(3 * N_MELS, N_CLASS),
|
||||
nn.Dropout(0.25),
|
||||
nn.Sigmoid()
|
||||
)
|
||||
|
||||
def forward(self, mel):
|
||||
mel = mel.transpose(-1, -2).unsqueeze(1)
|
||||
x = self.cnn(self.unet(mel)).transpose(1, 2).flatten(-2)
|
||||
x = self.fc(x)
|
||||
return x
|
||||
@@ -0,0 +1,10 @@
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class BiGRU(nn.Module):
|
||||
def __init__(self, input_features, hidden_features, num_layers):
|
||||
super(BiGRU, self).__init__()
|
||||
self.gru = nn.GRU(input_features, hidden_features, num_layers=num_layers, batch_first=True, bidirectional=True)
|
||||
|
||||
def forward(self, x):
|
||||
return self.gru(x)[0]
|
||||
@@ -0,0 +1,68 @@
|
||||
import torch
|
||||
import numpy as np
|
||||
import torch.nn.functional as F
|
||||
from librosa.filters import mel
|
||||
|
||||
|
||||
class MelSpectrogram(torch.nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
n_mel_channels,
|
||||
sampling_rate,
|
||||
win_length,
|
||||
hop_length,
|
||||
n_fft=None,
|
||||
mel_fmin=0,
|
||||
mel_fmax=None,
|
||||
clamp=1e-5
|
||||
):
|
||||
super().__init__()
|
||||
n_fft = win_length if n_fft is None else n_fft
|
||||
self.hann_window = {}
|
||||
mel_basis = mel(
|
||||
sr=sampling_rate,
|
||||
n_fft=n_fft,
|
||||
n_mels=n_mel_channels,
|
||||
fmin=mel_fmin,
|
||||
fmax=mel_fmax,
|
||||
htk=True)
|
||||
mel_basis = torch.from_numpy(mel_basis).float()
|
||||
self.register_buffer("mel_basis", mel_basis)
|
||||
self.n_fft = win_length if n_fft is None else n_fft
|
||||
self.hop_length = hop_length
|
||||
self.win_length = win_length
|
||||
self.sampling_rate = sampling_rate
|
||||
self.n_mel_channels = n_mel_channels
|
||||
self.clamp = clamp
|
||||
|
||||
def forward(self, audio, keyshift=0, speed=1, center=True):
|
||||
factor = 2 ** (keyshift / 12)
|
||||
n_fft_new = int(np.round(self.n_fft * factor))
|
||||
win_length_new = int(np.round(self.win_length * factor))
|
||||
hop_length_new = int(np.round(self.hop_length * speed))
|
||||
|
||||
keyshift_key = str(keyshift) + '_' + str(audio.device)
|
||||
if keyshift_key not in self.hann_window:
|
||||
self.hann_window[keyshift_key] = torch.hann_window(win_length_new).to(audio.device)
|
||||
|
||||
fft = torch.stft(
|
||||
audio,
|
||||
n_fft=n_fft_new,
|
||||
hop_length=hop_length_new,
|
||||
win_length=win_length_new,
|
||||
window=self.hann_window[keyshift_key],
|
||||
center=center,
|
||||
return_complex=True
|
||||
)
|
||||
magnitude = fft.abs()
|
||||
|
||||
if keyshift != 0:
|
||||
size = self.n_fft // 2 + 1
|
||||
resize = magnitude.size(1)
|
||||
if resize < size:
|
||||
magnitude = F.pad(magnitude, (0, 0, 0, size - resize))
|
||||
magnitude = magnitude[:, :size, :] * self.win_length / win_length_new
|
||||
|
||||
mel_output = torch.matmul(self.mel_basis, magnitude)
|
||||
log_mel_spec = torch.log(torch.clamp(mel_output, min=self.clamp))
|
||||
return log_mel_spec
|
||||
@@ -0,0 +1,43 @@
|
||||
import librosa
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from .constants import *
|
||||
|
||||
|
||||
def to_local_average_f0(hidden, center=None, thred=0.03):
|
||||
idx = torch.arange(N_CLASS, device=hidden.device)[None, None, :] # [B=1, T=1, N]
|
||||
idx_cents = idx * 20 + CONST # [B=1, N]
|
||||
if center is None:
|
||||
center = torch.argmax(hidden, dim=2, keepdim=True) # [B, T, 1]
|
||||
start = torch.clip(center - 4, min=0) # [B, T, 1]
|
||||
end = torch.clip(center + 5, max=N_CLASS) # [B, T, 1]
|
||||
idx_mask = (idx >= start) & (idx < end) # [B, T, N]
|
||||
weights = hidden * idx_mask # [B, T, N]
|
||||
product_sum = torch.sum(weights * idx_cents, dim=2) # [B, T]
|
||||
weight_sum = torch.sum(weights, dim=2) # [B, T]
|
||||
cents = product_sum / (weight_sum + (weight_sum == 0)) # avoid dividing by zero, [B, T]
|
||||
f0 = 10 * 2 ** (cents / 1200)
|
||||
uv = hidden.max(dim=2)[0] < thred # [B, T]
|
||||
f0 = f0 * ~uv
|
||||
return f0.squeeze(0).cpu().numpy()
|
||||
|
||||
|
||||
def to_viterbi_f0(hidden, thred=0.03):
|
||||
# Create viterbi transition matrix
|
||||
if not hasattr(to_viterbi_f0, 'transition'):
|
||||
xx, yy = np.meshgrid(range(N_CLASS), range(N_CLASS))
|
||||
transition = np.maximum(30 - abs(xx - yy), 0)
|
||||
transition = transition / transition.sum(axis=1, keepdims=True)
|
||||
to_viterbi_f0.transition = transition
|
||||
|
||||
# Convert to probability
|
||||
prob = hidden.squeeze(0).cpu().numpy()
|
||||
prob = prob.T
|
||||
prob = prob / prob.sum(axis=0)
|
||||
|
||||
# Perform viterbi decoding
|
||||
path = librosa.sequence.viterbi(prob, to_viterbi_f0.transition).astype(np.int64)
|
||||
center = torch.from_numpy(path).unsqueeze(0).unsqueeze(-1).to(hidden.device)
|
||||
|
||||
return to_local_average_f0(hidden, center=center, thred=thred)
|
||||
@@ -0,0 +1,366 @@
|
||||
from typing import Dict
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from torch import Tensor
|
||||
|
||||
import modules.compat as compat
|
||||
from basics.base_module import CategorizedModule
|
||||
from modules.aux_decoder import AuxDecoderAdaptor
|
||||
from modules.commons.common_layers import (
|
||||
NormalInitEmbedding as Embedding,
|
||||
SinusoidalPosEmb, AdamWLinear,
|
||||
)
|
||||
from modules.core import (
|
||||
GaussianDiffusion, PitchDiffusion, MultiVarianceDiffusion,
|
||||
RectifiedFlow, PitchRectifiedFlow, MultiVarianceRectifiedFlow
|
||||
)
|
||||
from modules.fastspeech.acoustic_encoder import FastSpeech2Acoustic
|
||||
from modules.fastspeech.param_adaptor import ParameterAdaptorModule
|
||||
from modules.fastspeech.tts_modules import RhythmRegulator, LengthRegulator, StretchRegulator
|
||||
from modules.fastspeech.variance_encoder import FastSpeech2Variance, MelodyEncoder
|
||||
from utils.hparams import hparams
|
||||
|
||||
|
||||
class ShallowDiffusionOutput:
|
||||
def __init__(self, *, aux_out=None, diff_out=None):
|
||||
self.aux_out = aux_out
|
||||
self.diff_out = diff_out
|
||||
|
||||
|
||||
class DiffSingerAcoustic(CategorizedModule, ParameterAdaptorModule):
|
||||
@property
|
||||
def category(self):
|
||||
return 'acoustic'
|
||||
|
||||
def __init__(self, vocab_size, out_dims):
|
||||
CategorizedModule.__init__(self)
|
||||
ParameterAdaptorModule.__init__(self)
|
||||
self.fs2 = FastSpeech2Acoustic(
|
||||
vocab_size=vocab_size
|
||||
)
|
||||
|
||||
self.use_shallow_diffusion = hparams.get('use_shallow_diffusion', False)
|
||||
self.shallow_args = hparams.get('shallow_diffusion_args', {})
|
||||
if self.use_shallow_diffusion:
|
||||
self.train_aux_decoder = self.shallow_args['train_aux_decoder']
|
||||
self.train_diffusion = self.shallow_args['train_diffusion']
|
||||
self.aux_decoder_grad = self.shallow_args['aux_decoder_grad']
|
||||
self.aux_decoder = AuxDecoderAdaptor(
|
||||
in_dims=hparams['hidden_size'], out_dims=out_dims, num_feats=1,
|
||||
spec_min=hparams['spec_min'], spec_max=hparams['spec_max'],
|
||||
aux_decoder_arch=self.shallow_args['aux_decoder_arch'],
|
||||
aux_decoder_args=self.shallow_args['aux_decoder_args']
|
||||
)
|
||||
self.diffusion_type = hparams.get('diffusion_type', 'ddpm')
|
||||
self.backbone_type = compat.get_backbone_type(hparams)
|
||||
self.backbone_args = compat.get_backbone_args(hparams, self.backbone_type)
|
||||
if self.diffusion_type == 'ddpm':
|
||||
self.diffusion = GaussianDiffusion(
|
||||
out_dims=out_dims,
|
||||
num_feats=1,
|
||||
timesteps=hparams['timesteps'],
|
||||
k_step=hparams['K_step'],
|
||||
backbone_type=self.backbone_type,
|
||||
backbone_args=self.backbone_args,
|
||||
spec_min=hparams['spec_min'],
|
||||
spec_max=hparams['spec_max']
|
||||
)
|
||||
elif self.diffusion_type == 'reflow':
|
||||
self.diffusion = RectifiedFlow(
|
||||
out_dims=out_dims,
|
||||
num_feats=1,
|
||||
t_start=hparams['T_start'],
|
||||
time_scale_factor=hparams['time_scale_factor'],
|
||||
backbone_type=self.backbone_type,
|
||||
backbone_args=self.backbone_args,
|
||||
spec_min=hparams['spec_min'],
|
||||
spec_max=hparams['spec_max']
|
||||
)
|
||||
else:
|
||||
raise NotImplementedError(self.diffusion_type)
|
||||
|
||||
def forward(
|
||||
self, txt_tokens, mel2ph, f0, key_shift=None, speed=None,
|
||||
spk_embed_id=None, languages=None, gt_mel=None, infer=True, **kwargs
|
||||
) -> ShallowDiffusionOutput:
|
||||
condition = self.fs2(
|
||||
txt_tokens, mel2ph, f0, key_shift=key_shift, speed=speed,
|
||||
spk_embed_id=spk_embed_id, languages=languages,
|
||||
**kwargs
|
||||
)
|
||||
if infer:
|
||||
if self.use_shallow_diffusion:
|
||||
aux_mel_pred = self.aux_decoder(condition, infer=True)
|
||||
aux_mel_pred *= ((mel2ph > 0).float()[:, :, None])
|
||||
if gt_mel is not None and self.shallow_args['val_gt_start']:
|
||||
src_mel = gt_mel
|
||||
else:
|
||||
src_mel = aux_mel_pred
|
||||
else:
|
||||
aux_mel_pred = src_mel = None
|
||||
mel_pred = self.diffusion(condition, src_spec=src_mel, infer=True)
|
||||
mel_pred *= ((mel2ph > 0).float()[:, :, None])
|
||||
return ShallowDiffusionOutput(aux_out=aux_mel_pred, diff_out=mel_pred)
|
||||
else:
|
||||
if self.use_shallow_diffusion:
|
||||
if self.train_aux_decoder:
|
||||
aux_cond = condition * self.aux_decoder_grad + condition.detach() * (1 - self.aux_decoder_grad)
|
||||
aux_out = self.aux_decoder(aux_cond, infer=False)
|
||||
else:
|
||||
aux_out = None
|
||||
if self.train_diffusion:
|
||||
diff_out = self.diffusion(condition, gt_spec=gt_mel, infer=False)
|
||||
else:
|
||||
diff_out = None
|
||||
return ShallowDiffusionOutput(aux_out=aux_out, diff_out=diff_out)
|
||||
|
||||
else:
|
||||
aux_out = None
|
||||
diff_out = self.diffusion(condition, gt_spec=gt_mel, infer=False)
|
||||
return ShallowDiffusionOutput(aux_out=aux_out, diff_out=diff_out)
|
||||
|
||||
|
||||
class DiffSingerVariance(CategorizedModule, ParameterAdaptorModule):
|
||||
@property
|
||||
def category(self):
|
||||
return 'variance'
|
||||
|
||||
def __init__(self, vocab_size):
|
||||
CategorizedModule.__init__(self)
|
||||
ParameterAdaptorModule.__init__(self)
|
||||
self.predict_dur = hparams['predict_dur']
|
||||
self.predict_pitch = hparams['predict_pitch']
|
||||
|
||||
self.use_stretch_embed = hparams.get('use_stretch_embed', None)
|
||||
assert self.use_stretch_embed is not None, "You may be loading an old version of the model checkpoint, which is incompatible with the new version due to some bug fixes. It is recommended to roll back to the old version (commit id: 6df0ee977c3728f14cb79c2db8b19df30b23a0bf)"
|
||||
if self.use_stretch_embed and (self.predict_pitch or self.predict_variances):
|
||||
self.sr = StretchRegulator()
|
||||
self.stretch_embed = nn.Sequential(
|
||||
SinusoidalPosEmb(hparams['hidden_size']),
|
||||
nn.Linear(hparams['hidden_size'], hparams['hidden_size'] * 4),
|
||||
nn.GELU(),
|
||||
nn.Linear(hparams['hidden_size'] * 4, hparams['hidden_size']),
|
||||
)
|
||||
self.stretch_embed_rnn = nn.GRU(hparams['hidden_size'], hparams['hidden_size'], 1, batch_first=True)
|
||||
|
||||
self.use_spk_id = hparams['use_spk_id']
|
||||
if self.use_spk_id:
|
||||
self.spk_embed = Embedding(hparams['num_spk'], hparams['hidden_size'])
|
||||
|
||||
self.fs2 = FastSpeech2Variance(
|
||||
vocab_size=vocab_size
|
||||
)
|
||||
self.rr = RhythmRegulator()
|
||||
self.lr = LengthRegulator()
|
||||
self.diffusion_type = hparams.get('diffusion_type', 'ddpm')
|
||||
if self.predict_pitch:
|
||||
self.use_melody_encoder = hparams.get('use_melody_encoder', False)
|
||||
if self.use_melody_encoder:
|
||||
self.melody_encoder = MelodyEncoder(enc_hparams=hparams['melody_encoder_args'])
|
||||
self.delta_pitch_embed = AdamWLinear(1, hparams['hidden_size'])
|
||||
else:
|
||||
self.base_pitch_embed = AdamWLinear(1, hparams['hidden_size'])
|
||||
|
||||
self.pitch_retake_embed = Embedding(2, hparams['hidden_size'])
|
||||
pitch_hparams = hparams['pitch_prediction_args']
|
||||
self.pitch_backbone_type = compat.get_backbone_type(hparams, nested_config=pitch_hparams)
|
||||
self.pitch_backbone_args = compat.get_backbone_args(pitch_hparams, backbone_type=self.pitch_backbone_type)
|
||||
if self.diffusion_type == 'ddpm':
|
||||
self.pitch_predictor = PitchDiffusion(
|
||||
vmin=pitch_hparams['pitd_norm_min'],
|
||||
vmax=pitch_hparams['pitd_norm_max'],
|
||||
cmin=pitch_hparams['pitd_clip_min'],
|
||||
cmax=pitch_hparams['pitd_clip_max'],
|
||||
repeat_bins=pitch_hparams['repeat_bins'],
|
||||
timesteps=hparams['timesteps'],
|
||||
k_step=hparams['K_step'],
|
||||
backbone_type=self.pitch_backbone_type,
|
||||
backbone_args=self.pitch_backbone_args
|
||||
)
|
||||
elif self.diffusion_type == 'reflow':
|
||||
self.pitch_predictor = PitchRectifiedFlow(
|
||||
vmin=pitch_hparams['pitd_norm_min'],
|
||||
vmax=pitch_hparams['pitd_norm_max'],
|
||||
cmin=pitch_hparams['pitd_clip_min'],
|
||||
cmax=pitch_hparams['pitd_clip_max'],
|
||||
repeat_bins=pitch_hparams['repeat_bins'],
|
||||
time_scale_factor=hparams['time_scale_factor'],
|
||||
backbone_type=self.pitch_backbone_type,
|
||||
backbone_args=self.pitch_backbone_args
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Invalid diffusion type: {self.diffusion_type}")
|
||||
|
||||
if self.predict_variances:
|
||||
self.pitch_embed = AdamWLinear(1, hparams['hidden_size'])
|
||||
self.variance_embeds = nn.ModuleDict({
|
||||
v_name: AdamWLinear(1, hparams['hidden_size'])
|
||||
for v_name in self.variance_prediction_list
|
||||
})
|
||||
|
||||
if self.diffusion_type == 'ddpm':
|
||||
self.variance_predictor = self.build_adaptor(cls=MultiVarianceDiffusion)
|
||||
elif self.diffusion_type == 'reflow':
|
||||
self.variance_predictor = self.build_adaptor(cls=MultiVarianceRectifiedFlow)
|
||||
else:
|
||||
raise NotImplementedError(self.diffusion_type)
|
||||
|
||||
self.use_variance_scaling = hparams.get('use_variance_scaling', False)
|
||||
self.custom_variance_scaling_factor = {
|
||||
'energy': 1. / 96,
|
||||
'breathiness': 1. / 96,
|
||||
'voicing': 1. / 96,
|
||||
'tension': 0.1,
|
||||
'key_shift': 1. / 12,
|
||||
'speed': 1.
|
||||
}
|
||||
self.default_variance_scaling_factor = {
|
||||
'energy': 1.,
|
||||
'breathiness': 1.,
|
||||
'voicing': 1.,
|
||||
'tension': 1.,
|
||||
'key_shift': 1.,
|
||||
'speed': 1.
|
||||
}
|
||||
if self.use_variance_scaling:
|
||||
self.variance_retake_scaling = self.custom_variance_scaling_factor
|
||||
else:
|
||||
self.variance_retake_scaling = self.default_variance_scaling_factor
|
||||
|
||||
def forward(
|
||||
self, txt_tokens, midi, ph2word, ph_dur=None, word_dur=None, mel2ph=None,
|
||||
note_midi=None, note_rest=None, note_dur=None, note_glide=None, mel2note=None,
|
||||
base_pitch=None, pitch=None, pitch_expr=None, pitch_retake=None,
|
||||
variance_retake: Dict[str, Tensor] = None,
|
||||
spk_id=None, languages=None,
|
||||
infer=True, **kwargs
|
||||
):
|
||||
if self.use_spk_id:
|
||||
ph_spk_mix_embed = kwargs.get('ph_spk_mix_embed')
|
||||
spk_mix_embed = kwargs.get('spk_mix_embed')
|
||||
if ph_spk_mix_embed is not None and spk_mix_embed is not None:
|
||||
ph_spk_embed = ph_spk_mix_embed
|
||||
spk_embed = spk_mix_embed
|
||||
else:
|
||||
ph_spk_embed = spk_embed = self.spk_embed(spk_id)[:, None, :] # [B,] => [B, T=1, H]
|
||||
else:
|
||||
ph_spk_embed = spk_embed = None
|
||||
|
||||
encoder_out, dur_pred_out = self.fs2(
|
||||
txt_tokens, midi=midi, ph2word=ph2word,
|
||||
ph_dur=ph_dur, word_dur=word_dur,
|
||||
spk_embed=ph_spk_embed, languages=languages,
|
||||
infer=infer
|
||||
)
|
||||
|
||||
if not self.predict_pitch and not self.predict_variances:
|
||||
return dur_pred_out, None, ({} if infer else None)
|
||||
|
||||
if mel2ph is None and word_dur is not None: # inference from file
|
||||
dur_pred_align = self.rr(dur_pred_out, ph2word, word_dur)
|
||||
mel2ph = self.lr(dur_pred_align)
|
||||
mel2ph = F.pad(mel2ph, [0, base_pitch.shape[1] - mel2ph.shape[1]])
|
||||
|
||||
encoder_out = F.pad(encoder_out, [0, 0, 1, 0])
|
||||
mel2ph_ = mel2ph[..., None].repeat([1, 1, hparams['hidden_size']])
|
||||
condition = torch.gather(encoder_out, 1, mel2ph_)
|
||||
|
||||
if self.use_stretch_embed:
|
||||
stretch = torch.round(1000 * self.sr(mel2ph, ph_dur))
|
||||
if self.training and stretch.numel() > 1000:
|
||||
# construct a phoneme stretching index lookup table with a total of 1001 indexes (0~1000)
|
||||
table = self.stretch_embed(torch.arange(0, 1001, device=stretch.device))
|
||||
stretch_embed = torch.index_select(table, 0, stretch.view(-1).long()).view_as(condition)
|
||||
else:
|
||||
stretch_embed = self.stretch_embed(stretch)
|
||||
condition += stretch_embed
|
||||
self.stretch_embed_rnn.flatten_parameters()
|
||||
stretch_embed_rnn_out, _ = self.stretch_embed_rnn(condition)
|
||||
condition = condition + stretch_embed_rnn_out
|
||||
|
||||
if self.use_spk_id:
|
||||
condition += spk_embed
|
||||
|
||||
if self.predict_pitch:
|
||||
if self.use_melody_encoder:
|
||||
melody_encoder_out = self.melody_encoder(
|
||||
note_midi, note_rest, note_dur,
|
||||
glide=note_glide
|
||||
)
|
||||
melody_encoder_out = F.pad(melody_encoder_out, [0, 0, 1, 0])
|
||||
mel2note_ = mel2note[..., None].repeat([1, 1, hparams['hidden_size']])
|
||||
melody_condition = torch.gather(melody_encoder_out, 1, mel2note_)
|
||||
pitch_cond = condition + melody_condition
|
||||
else:
|
||||
pitch_cond = condition.clone() # preserve the original tensor to avoid further inplace operations
|
||||
|
||||
retake_unset = pitch_retake is None
|
||||
if retake_unset:
|
||||
pitch_retake = torch.ones_like(mel2ph, dtype=torch.bool)
|
||||
|
||||
if pitch_expr is None:
|
||||
pitch_retake_embed = self.pitch_retake_embed(pitch_retake.long())
|
||||
else:
|
||||
retake_true_embed = self.pitch_retake_embed(
|
||||
torch.ones(1, 1, dtype=torch.long, device=txt_tokens.device)
|
||||
) # [B=1, T=1] => [B=1, T=1, H]
|
||||
retake_false_embed = self.pitch_retake_embed(
|
||||
torch.zeros(1, 1, dtype=torch.long, device=txt_tokens.device)
|
||||
) # [B=1, T=1] => [B=1, T=1, H]
|
||||
pitch_expr = (pitch_expr * pitch_retake)[:, :, None] # [B, T, 1]
|
||||
pitch_retake_embed = pitch_expr * retake_true_embed + (1. - pitch_expr) * retake_false_embed
|
||||
|
||||
pitch_cond += pitch_retake_embed
|
||||
if self.use_melody_encoder:
|
||||
if retake_unset: # generate from scratch
|
||||
delta_pitch_in = torch.zeros_like(base_pitch)
|
||||
else:
|
||||
delta_pitch_in = (pitch - base_pitch) * ~pitch_retake
|
||||
if self.use_variance_scaling:
|
||||
pitch_cond += self.delta_pitch_embed(delta_pitch_in[:, :, None] / 12)
|
||||
else:
|
||||
pitch_cond += self.delta_pitch_embed(delta_pitch_in[:, :, None])
|
||||
else:
|
||||
if not retake_unset: # retake
|
||||
base_pitch = base_pitch * pitch_retake + pitch * ~pitch_retake
|
||||
if self.use_variance_scaling:
|
||||
pitch_cond += self.base_pitch_embed(base_pitch[:, :, None] / 128)
|
||||
else:
|
||||
pitch_cond += self.base_pitch_embed(base_pitch[:, :, None])
|
||||
|
||||
if infer:
|
||||
pitch_pred_out = self.pitch_predictor(pitch_cond, infer=True)
|
||||
else:
|
||||
pitch_pred_out = self.pitch_predictor(pitch_cond, pitch - base_pitch, infer=False)
|
||||
else:
|
||||
pitch_pred_out = None
|
||||
|
||||
if not self.predict_variances:
|
||||
return dur_pred_out, pitch_pred_out, ({} if infer else None)
|
||||
|
||||
if pitch is None:
|
||||
pitch = base_pitch + pitch_pred_out
|
||||
if self.use_variance_scaling:
|
||||
var_cond = condition + self.pitch_embed(pitch[:, :, None] / 12)
|
||||
else:
|
||||
var_cond = condition + self.pitch_embed(pitch[:, :, None])
|
||||
|
||||
variance_inputs = self.collect_variance_inputs(**kwargs)
|
||||
|
||||
if variance_retake is not None:
|
||||
variance_embeds = [
|
||||
self.variance_embeds[v_name](v_input[:, :, None] * self.variance_retake_scaling[v_name]) * ~variance_retake[v_name][:, :, None]
|
||||
for v_name, v_input in zip(self.variance_prediction_list, variance_inputs)
|
||||
]
|
||||
var_cond += torch.stack(variance_embeds, dim=-1).sum(-1)
|
||||
|
||||
variance_outputs = self.variance_predictor(var_cond, variance_inputs, infer=infer)
|
||||
|
||||
if infer:
|
||||
variances_pred_out = self.collect_variance_outputs(variance_outputs)
|
||||
else:
|
||||
variances_pred_out = variance_outputs
|
||||
|
||||
return dur_pred_out, pitch_pred_out, variances_pred_out
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user