Spaces:
Running
on
A10G
Running
on
A10G
| # Copyright (c) 2023 Amphion. | |
| # | |
| # This source code is licensed under the MIT license found in the | |
| # LICENSE file in the root directory of this source tree. | |
| import argparse | |
| import os | |
| import glob | |
| from tqdm import tqdm | |
| import json | |
| import torch | |
| import time | |
| from models.svc.diffusion.diffusion_inference import DiffusionInference | |
| from models.svc.comosvc.comosvc_inference import ComoSVCInference | |
| from models.svc.transformer.transformer_inference import TransformerInference | |
| from utils.util import load_config | |
| from utils.audio_slicer import split_audio, merge_segments_encodec | |
| from processors import acoustic_extractor, content_extractor | |
| def build_inference(args, cfg, infer_type="from_dataset"): | |
| supported_inference = { | |
| "DiffWaveNetSVC": DiffusionInference, | |
| "DiffComoSVC": ComoSVCInference, | |
| "TransformerSVC": TransformerInference, | |
| } | |
| inference_class = supported_inference[cfg.model_type] | |
| return inference_class(args, cfg, infer_type) | |
| def prepare_for_audio_file(args, cfg, num_workers=1): | |
| preprocess_path = cfg.preprocess.processed_dir | |
| audio_name = cfg.inference.source_audio_name | |
| temp_audio_dir = os.path.join(preprocess_path, audio_name) | |
| ### eval file | |
| t = time.time() | |
| eval_file = prepare_source_eval_file(cfg, temp_audio_dir, audio_name) | |
| args.source = eval_file | |
| with open(eval_file, "r") as f: | |
| metadata = json.load(f) | |
| print("Prepare for meta eval data: {:.1f}s".format(time.time() - t)) | |
| ### acoustic features | |
| t = time.time() | |
| acoustic_extractor.extract_utt_acoustic_features_serial( | |
| metadata, temp_audio_dir, cfg | |
| ) | |
| acoustic_extractor.cal_mel_min_max( | |
| dataset=audio_name, output_path=preprocess_path, cfg=cfg, metadata=metadata | |
| ) | |
| acoustic_extractor.cal_pitch_statistics_svc( | |
| dataset=audio_name, output_path=preprocess_path, cfg=cfg, metadata=metadata | |
| ) | |
| print("Prepare for acoustic features: {:.1f}s".format(time.time() - t)) | |
| ### content features | |
| t = time.time() | |
| content_extractor.extract_utt_content_features_dataloader( | |
| cfg, metadata, num_workers | |
| ) | |
| print("Prepare for content features: {:.1f}s".format(time.time() - t)) | |
| return args, cfg, temp_audio_dir | |
| def merge_for_audio_segments(audio_files, args, cfg): | |
| audio_name = cfg.inference.source_audio_name | |
| target_singer_name = args.target_singer | |
| merge_segments_encodec( | |
| wav_files=audio_files, | |
| fs=cfg.preprocess.sample_rate, | |
| output_path=os.path.join( | |
| args.output_dir, "{}_{}.wav".format(audio_name, target_singer_name) | |
| ), | |
| overlap_duration=cfg.inference.segments_overlap_duration, | |
| ) | |
| for tmp_file in audio_files: | |
| os.remove(tmp_file) | |
| def prepare_source_eval_file(cfg, temp_audio_dir, audio_name): | |
| """ | |
| Prepare the eval file (json) for an audio | |
| """ | |
| audio_chunks_results = split_audio( | |
| wav_file=cfg.inference.source_audio_path, | |
| target_sr=cfg.preprocess.sample_rate, | |
| output_dir=os.path.join(temp_audio_dir, "wavs"), | |
| max_duration_of_segment=cfg.inference.segments_max_duration, | |
| overlap_duration=cfg.inference.segments_overlap_duration, | |
| ) | |
| metadata = [] | |
| for i, res in enumerate(audio_chunks_results): | |
| res["index"] = i | |
| res["Dataset"] = audio_name | |
| res["Singer"] = audio_name | |
| res["Uid"] = "{}_{}".format(audio_name, res["Uid"]) | |
| metadata.append(res) | |
| eval_file = os.path.join(temp_audio_dir, "eval.json") | |
| with open(eval_file, "w") as f: | |
| json.dump(metadata, f, indent=4, ensure_ascii=False, sort_keys=True) | |
| return eval_file | |
| def cuda_relevant(deterministic=False): | |
| torch.cuda.empty_cache() | |
| # TF32 on Ampere and above | |
| torch.backends.cuda.matmul.allow_tf32 = True | |
| torch.backends.cudnn.enabled = True | |
| torch.backends.cudnn.allow_tf32 = True | |
| # Deterministic | |
| torch.backends.cudnn.deterministic = deterministic | |
| torch.backends.cudnn.benchmark = not deterministic | |
| torch.use_deterministic_algorithms(deterministic) | |
| def infer(args, cfg, infer_type): | |
| # Build inference | |
| t = time.time() | |
| trainer = build_inference(args, cfg, infer_type) | |
| print("Model Init: {:.1f}s".format(time.time() - t)) | |
| # Run inference | |
| t = time.time() | |
| output_audio_files = trainer.inference() | |
| print("Model inference: {:.1f}s".format(time.time() - t)) | |
| return output_audio_files | |
| def build_parser(): | |
| r"""Build argument parser for inference.py. | |
| Anything else should be put in an extra config YAML file. | |
| """ | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument( | |
| "--config", | |
| type=str, | |
| required=True, | |
| help="JSON/YAML file for configurations.", | |
| ) | |
| parser.add_argument( | |
| "--acoustics_dir", | |
| type=str, | |
| help="Acoustics model checkpoint directory. If a directory is given, " | |
| "search for the latest checkpoint dir in the directory. If a specific " | |
| "checkpoint dir is given, directly load the checkpoint.", | |
| ) | |
| parser.add_argument( | |
| "--vocoder_dir", | |
| type=str, | |
| required=True, | |
| help="Vocoder checkpoint directory. Searching behavior is the same as " | |
| "the acoustics one.", | |
| ) | |
| parser.add_argument( | |
| "--target_singer", | |
| type=str, | |
| required=True, | |
| help="convert to a specific singer (e.g. --target_singers singer_id).", | |
| ) | |
| parser.add_argument( | |
| "--trans_key", | |
| default=0, | |
| help="0: no pitch shift; autoshift: pitch shift; int: key shift.", | |
| ) | |
| parser.add_argument( | |
| "--source", | |
| type=str, | |
| default="source_audio", | |
| help="Source audio file or directory. If a JSON file is given, " | |
| "inference from dataset is applied. If a directory is given, " | |
| "inference from all wav/flac/mp3 audio files in the directory is applied. " | |
| "Default: inference from all wav/flac/mp3 audio files in ./source_audio", | |
| ) | |
| parser.add_argument( | |
| "--output_dir", | |
| type=str, | |
| default="conversion_results", | |
| help="Output directory. Default: ./conversion_results", | |
| ) | |
| parser.add_argument( | |
| "--log_level", | |
| type=str, | |
| default="warning", | |
| help="Logging level. Default: warning", | |
| ) | |
| parser.add_argument( | |
| "--keep_cache", | |
| action="store_true", | |
| default=True, | |
| help="Keep cache files. Only applicable to inference from files.", | |
| ) | |
| parser.add_argument( | |
| "--diffusion_inference_steps", | |
| type=int, | |
| default=1000, | |
| help="Number of inference steps. Only applicable to diffusion inference.", | |
| ) | |
| return parser | |
| def main(args_list): | |
| ### Parse arguments and config | |
| args = build_parser().parse_args(args_list) | |
| cfg = load_config(args.config) | |
| # CUDA settings | |
| cuda_relevant() | |
| if os.path.isdir(args.source): | |
| ### Infer from file | |
| # Get all the source audio files (.wav, .flac, .mp3) | |
| source_audio_dir = args.source | |
| audio_list = [] | |
| for suffix in ["wav", "flac", "mp3"]: | |
| audio_list += glob.glob( | |
| os.path.join(source_audio_dir, "**/*.{}".format(suffix)), recursive=True | |
| ) | |
| print("There are {} source audios: ".format(len(audio_list))) | |
| # Infer for every file as dataset | |
| output_root_path = args.output_dir | |
| for audio_path in tqdm(audio_list): | |
| audio_name = audio_path.split("/")[-1].split(".")[0] | |
| args.output_dir = os.path.join(output_root_path, audio_name) | |
| print("\n{}\nConversion for {}...\n".format("*" * 10, audio_name)) | |
| cfg.inference.source_audio_path = audio_path | |
| cfg.inference.source_audio_name = audio_name | |
| cfg.inference.segments_max_duration = 10.0 | |
| cfg.inference.segments_overlap_duration = 1.0 | |
| # Prepare metadata and features | |
| args, cfg, cache_dir = prepare_for_audio_file(args, cfg) | |
| # Infer from file | |
| output_audio_files = infer(args, cfg, infer_type="from_file") | |
| # Merge the split segments | |
| merge_for_audio_segments(output_audio_files, args, cfg) | |
| # Keep or remove caches | |
| if not args.keep_cache: | |
| os.removedirs(cache_dir) | |
| else: | |
| ### Infer from dataset | |
| infer(args, cfg, infer_type="from_dataset") | |
| if __name__ == "__main__": | |
| main() | |