# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. # # 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. """ This script is to generate mel spectrograms from a Fastpitch model checkpoint. Please see general usage below. It runs on GPUs by default, but you can add `--num-workers 5 --cpu` as an option to run on CPUs. $ python scripts/dataset_processing/tts/generate_mels.py \ --fastpitch-model-ckpt ./models/fastpitch/multi_spk/FastPitch--val_loss\=1.4473-epoch\=209.ckpt \ --input-json-manifests /home/xueyang/HUI-Audio-Corpus-German-clean/test_manifest_text_normed_phonemes.json --output-json-manifest-root /home/xueyang/experiments/multi_spk_tts_de """ import argparse import json from pathlib import Path import numpy as np import soundfile as sf import torch from joblib import Parallel, delayed from tqdm import tqdm from nemo.collections.tts.models import FastPitchModel from nemo.collections.tts.parts.utils.tts_dataset_utils import ( BetaBinomialInterpolator, beta_binomial_prior_distribution, ) from nemo.utils import logging def get_args(): parser = argparse.ArgumentParser( formatter_class=argparse.ArgumentDefaultsHelpFormatter, description="Generate mel spectrograms with pretrained FastPitch model, and create manifests for finetuning Hifigan.", ) parser.add_argument( "--fastpitch-model-ckpt", required=True, type=Path, help="Specify a full path of a fastpitch model checkpoint with the suffix of either .ckpt or .nemo.", ) parser.add_argument( "--input-json-manifests", nargs="+", required=True, type=Path, help="Specify a full path of a JSON manifest. You could add multiple manifests.", ) parser.add_argument( "--output-json-manifest-root", required=True, type=Path, help="Specify a full path of output root that would contain new manifests.", ) parser.add_argument( "--num-workers", default=-1, type=int, help="Specify the max number of concurrently Python workers processes. " "If -1 all CPUs are used. If 1 no parallel computing is used.", ) parser.add_argument("--cpu", action='store_true', default=False, help="Generate mel spectrograms using CPUs.") args = parser.parse_args() return args def __load_wav(audio_file): with sf.SoundFile(audio_file, 'r') as f: samples = f.read(dtype='float32') return samples.transpose() def __generate_mels(entry, spec_model, device, use_beta_binomial_interpolator, mel_root): # Generate a spectrograms (we need to use ground truth alignment for correct matching between audio and mels) audio = __load_wav(entry["audio_filepath"]) audio = torch.from_numpy(audio).unsqueeze(0).to(device) audio_len = torch.tensor(audio.shape[1], dtype=torch.long, device=device).unsqueeze(0) if spec_model.fastpitch.speaker_emb is not None and "speaker" in entry: speaker = torch.tensor([entry['speaker']]).to(device) else: speaker = None with torch.no_grad(): if "normalized_text" in entry: text = spec_model.parse(entry["normalized_text"], normalize=False) else: text = spec_model.parse(entry['text']) text_len = torch.tensor(text.shape[-1], dtype=torch.long, device=device).unsqueeze(0) spect, spect_len = spec_model.preprocessor(input_signal=audio, length=audio_len) # Generate attention prior and spectrogram inputs for HiFi-GAN if use_beta_binomial_interpolator: beta_binomial_interpolator = BetaBinomialInterpolator() attn_prior = ( torch.from_numpy(beta_binomial_interpolator(spect_len.item(), text_len.item())) .unsqueeze(0) .to(text.device) ) else: attn_prior = ( torch.from_numpy(beta_binomial_prior_distribution(text_len.item(), spect_len.item())) .unsqueeze(0) .to(text.device) ) spectrogram = spec_model.forward( text=text, input_lens=text_len, spec=spect, mel_lens=spect_len, attn_prior=attn_prior, speaker=speaker, )[0] save_path = mel_root / f"{Path(entry['audio_filepath']).stem}.npy" np.save(save_path, spectrogram[0].to('cpu').numpy()) entry["mel_filepath"] = str(save_path) return entry def main(): args = get_args() ckpt_path = args.fastpitch_model_ckpt input_manifest_filepaths = args.input_json_manifests output_json_manifest_root = args.output_json_manifest_root mel_root = output_json_manifest_root / "mels" mel_root.mkdir(exist_ok=True, parents=True) # load pretrained FastPitch model checkpoint suffix = ckpt_path.suffix if suffix == ".nemo": spec_model = FastPitchModel.restore_from(ckpt_path).eval() elif suffix == ".ckpt": spec_model = FastPitchModel.load_from_checkpoint(ckpt_path).eval() else: raise ValueError(f"Unsupported suffix: {suffix}") if not args.cpu: spec_model.cuda() device = spec_model.device use_beta_binomial_interpolator = spec_model.cfg.train_ds.dataset.get("use_beta_binomial_interpolator", False) for manifest in input_manifest_filepaths: logging.info(f"Processing {manifest}.") entries = [] with open(manifest, "r") as fjson: for line in fjson: entries.append(json.loads(line.strip())) if device == "cpu": new_entries = Parallel(n_jobs=args.num_workers)( delayed(__generate_mels)(entry, spec_model, device, use_beta_binomial_interpolator, mel_root) for entry in entries ) else: new_entries = [] for entry in tqdm(entries): new_entry = __generate_mels(entry, spec_model, device, use_beta_binomial_interpolator, mel_root) new_entries.append(new_entry) mel_manifest_path = output_json_manifest_root / f"{manifest.stem}_mel{manifest.suffix}" with open(mel_manifest_path, "w") as fmel: for entry in new_entries: fmel.write(json.dumps(entry) + "\n") logging.info(f"Processing {manifest} is complete --> {mel_manifest_path}") if __name__ == "__main__": main()