Spaces:
Paused
Paused
| # 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 torch | |
| import numpy as np | |
| import librosa | |
| from safetensors.torch import load_model | |
| import os | |
| from utils.util import load_config | |
| from models.vc.Noro.noro_trainer import NoroTrainer | |
| from models.vc.Noro.noro_model import Noro_VCmodel | |
| from processors.content_extractor import HubertExtractor | |
| from utils.mel import mel_spectrogram_torch | |
| from utils.f0 import get_f0_features_using_dio, interpolate | |
| from torch.nn.utils.rnn import pad_sequence | |
| def build_trainer(args, cfg): | |
| supported_trainer = { | |
| "VC": NoroTrainer, | |
| } | |
| trainer_class = supported_trainer[cfg.model_type] | |
| trainer = trainer_class(args, cfg) | |
| return trainer | |
| def main(): | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument( | |
| "--config", | |
| default="config.json", | |
| help="JSON file for configurations.", | |
| required=True, | |
| ) | |
| parser.add_argument( | |
| "--checkpoint_path", | |
| type=str, | |
| help="Checkpoint for resume training or fine-tuning.", | |
| required=True, | |
| ) | |
| parser.add_argument( | |
| "--output_dir", | |
| help="Output path", | |
| required=True, | |
| ) | |
| parser.add_argument( | |
| "--ref_path", | |
| type=str, | |
| help="Reference voice path", | |
| ) | |
| parser.add_argument( | |
| "--source_path", | |
| type=str, | |
| help="Source voice path", | |
| ) | |
| parser.add_argument("--cuda_id", type=int, default=0, help="CUDA id for training.") | |
| parser.add_argument("--local_rank", default=-1, type=int) | |
| args = parser.parse_args() | |
| cfg = load_config(args.config) | |
| if not os.path.exists(args.output_dir): | |
| os.makedirs(args.output_dir) | |
| cuda_id = args.cuda_id | |
| args.local_rank = torch.device(f"cuda:{cuda_id}") | |
| print("Local rank:", args.local_rank) | |
| args.content_extractor = "mhubert" | |
| with torch.cuda.device(args.local_rank): | |
| torch.cuda.empty_cache() | |
| ckpt_path = args.checkpoint_path | |
| w2v = HubertExtractor(cfg) | |
| w2v = w2v.to(device=args.local_rank) | |
| w2v.eval() | |
| model = Noro_VCmodel(cfg=cfg.model) | |
| print("Loading model") | |
| load_model(model, ckpt_path) | |
| print("Model loaded") | |
| model.cuda(args.local_rank) | |
| model.eval() | |
| wav_path = args.source_path | |
| ref_wav_path = args.ref_path | |
| wav, _ = librosa.load(wav_path, sr=16000) | |
| wav = np.pad(wav, (0, 1600 - len(wav) % 1600)) | |
| audio = torch.from_numpy(wav).to(args.local_rank) | |
| audio = audio[None, :] | |
| ref_wav, _ = librosa.load(ref_wav_path, sr=16000) | |
| ref_wav = np.pad(ref_wav, (0, 200 - len(ref_wav) % 200)) | |
| ref_audio = torch.from_numpy(ref_wav).to(args.local_rank) | |
| ref_audio = ref_audio[None, :] | |
| with torch.no_grad(): | |
| ref_mel = mel_spectrogram_torch(ref_audio, cfg) | |
| ref_mel = ref_mel.transpose(1, 2).to(device=args.local_rank) | |
| ref_mask = ( | |
| torch.ones(ref_mel.shape[0], ref_mel.shape[1]).to(args.local_rank).bool() | |
| ) | |
| _, content_feature = w2v.extract_content_features(audio) | |
| content_feature = content_feature.to(device=args.local_rank) | |
| wav = audio.cpu().numpy() | |
| wav = wav[0, :] | |
| f0s = [] | |
| pitch_raw = get_f0_features_using_dio(wav, cfg.preprocess) | |
| pitch_raw, _ = interpolate(pitch_raw) | |
| frame_num = len(wav) // cfg.preprocess.hop_size | |
| pitch_raw = torch.from_numpy(pitch_raw[:frame_num]).float() | |
| f0s.append(pitch_raw) | |
| pitch = pad_sequence(f0s, batch_first=True, padding_value=0).float() | |
| pitch = (pitch - pitch.mean(dim=1, keepdim=True)) / ( | |
| pitch.std(dim=1, keepdim=True) + 1e-6 | |
| ) | |
| pitch = pitch.to(device=args.local_rank) | |
| x0 = model.inference( | |
| content_feature=content_feature, | |
| pitch=pitch, | |
| x_ref=ref_mel, | |
| x_ref_mask=ref_mask, | |
| inference_steps=200, | |
| sigma=1.2, | |
| ) # 150-300 0.95-1.5 | |
| recon_path = f"{args.output_dir}/recon_mel.npy" | |
| np.save(recon_path, x0.transpose(1, 2).detach().cpu().numpy()) | |
| print(f"Mel spectrogram saved to: {recon_path}") | |
| if __name__ == "__main__": | |
| main() | |