| |
| |
| |
| |
|
|
| import shutil |
| import warnings |
| import argparse |
| import torch |
| import os |
| import yaml |
|
|
| warnings.simplefilter("ignore") |
|
|
| from .modules.commons import * |
| import time |
|
|
| import torchaudio |
| import librosa |
| from collections import OrderedDict |
|
|
|
|
| class FAcodecInference(object): |
| def __init__(self, args=None, cfg=None): |
| self.args = args |
| self.cfg = cfg |
| self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| self.model = self._build_model() |
| self._load_checkpoint() |
|
|
| def _build_model(self): |
| model = build_model(self.cfg.model_params) |
| _ = [model[key].to(self.device) for key in model] |
| return model |
|
|
| def _load_checkpoint(self): |
| sd = torch.load(self.args.checkpoint_path, map_location="cpu") |
| sd = sd["net"] if "net" in sd else sd |
| new_params = dict() |
| for key, state_dict in sd.items(): |
| new_state_dict = OrderedDict() |
| for k, v in state_dict.items(): |
| if k.startswith("module."): |
| k = k[7:] |
| new_state_dict[k] = v |
| new_params[key] = new_state_dict |
| for key in new_params: |
| if key in self.model: |
| self.model[key].load_state_dict(new_params[key]) |
| _ = [self.model[key].eval() for key in self.model] |
|
|
| @torch.no_grad() |
| def inference(self, source, output_dir): |
| source_audio = librosa.load(source, sr=self.cfg.preprocess_params.sr)[0] |
| source_audio = torch.tensor(source_audio).unsqueeze(0).float().to(self.device) |
|
|
| z = self.model.encoder(source_audio[None, ...].to(self.device).float()) |
| ( |
| z, |
| quantized, |
| commitment_loss, |
| codebook_loss, |
| timbre, |
| codes, |
| ) = self.model.quantizer( |
| z, |
| source_audio[None, ...].to(self.device).float(), |
| n_c=self.cfg.model_params.n_c_codebooks, |
| return_codes=True, |
| ) |
|
|
| full_pred_wave = self.model.decoder(z) |
|
|
| os.makedirs(output_dir, exist_ok=True) |
| source_name = source.split("/")[-1].split(".")[0] |
| torchaudio.save( |
| f"{output_dir}/reconstructed_{source_name}.wav", |
| full_pred_wave[0].cpu(), |
| self.cfg.preprocess_params.sr, |
| ) |
|
|
| print( |
| "Reconstructed audio saved as: ", |
| f"{output_dir}/reconstructed_{source_name}.wav", |
| ) |
|
|
| return quantized, codes |
|
|
| @torch.no_grad() |
| def voice_conversion(self, source, reference, output_dir): |
| source_audio = librosa.load(source, sr=self.cfg.preprocess_params.sr)[0] |
| source_audio = torch.tensor(source_audio).unsqueeze(0).float().to(self.device) |
|
|
| reference_audio = librosa.load(reference, sr=self.cfg.preprocess_params.sr)[0] |
| reference_audio = ( |
| torch.tensor(reference_audio).unsqueeze(0).float().to(self.device) |
| ) |
|
|
| z = self.model.encoder(source_audio[None, ...].to(self.device).float()) |
| z, quantized, commitment_loss, codebook_loss, timbre = self.model.quantizer( |
| z, |
| source_audio[None, ...].to(self.device).float(), |
| n_c=self.cfg.model_params.n_c_codebooks, |
| ) |
|
|
| z_ref = self.model.encoder(reference_audio[None, ...].to(self.device).float()) |
| ( |
| z_ref, |
| quantized_ref, |
| commitment_loss_ref, |
| codebook_loss_ref, |
| timbre_ref, |
| ) = self.model.quantizer( |
| z_ref, |
| reference_audio[None, ...].to(self.device).float(), |
| n_c=self.cfg.model_params.n_c_codebooks, |
| ) |
|
|
| z_conv = self.model.quantizer.voice_conversion( |
| quantized[0] + quantized[1], |
| reference_audio[None, ...].to(self.device).float(), |
| ) |
| full_pred_wave = self.model.decoder(z_conv) |
|
|
| os.makedirs(output_dir, exist_ok=True) |
| source_name = source.split("/")[-1].split(".")[0] |
| reference_name = reference.split("/")[-1].split(".")[0] |
| torchaudio.save( |
| f"{output_dir}/converted_{source_name}_to_{reference_name}.wav", |
| full_pred_wave[0].cpu(), |
| self.cfg.preprocess_params.sr, |
| ) |
|
|
| print( |
| "Voice conversion results saved as: ", |
| f"{output_dir}/converted_{source_name}_to_{reference_name}.wav", |
| ) |
|
|
