| import torch
|
| import numpy as np
|
| import re
|
| import soundfile
|
| from . import utils
|
| from . import commons
|
| import os
|
| import librosa
|
|
|
| from .mel_processing import spectrogram_torch
|
| from .models import SynthesizerTrn
|
|
|
|
|
| class OpenVoiceBaseClass(object):
|
| def __init__(self,
|
| config_path,
|
| device='cuda:0'):
|
| if 'cuda' in device:
|
| assert torch.cuda.is_available()
|
|
|
| hps = utils.get_hparams_from_file(config_path)
|
|
|
| model = SynthesizerTrn(
|
| len(getattr(hps, 'symbols', [])),
|
| hps.data.filter_length // 2 + 1,
|
| n_speakers=hps.data.n_speakers,
|
| **hps.model,
|
| ).to(device)
|
|
|
| model.eval()
|
| self.model = model
|
| self.hps = hps
|
| self.device = device
|
|
|
| def load_ckpt(self, ckpt_path):
|
| checkpoint_dict = torch.load(ckpt_path, map_location=torch.device(self.device))
|
| a, b = self.model.load_state_dict(checkpoint_dict['model'], strict=False)
|
| print("Loaded checkpoint '{}'".format(ckpt_path))
|
| print('missing/unexpected keys:', a, b)
|
|
|
|
|
| class BaseSpeakerTTS(OpenVoiceBaseClass):
|
| language_marks = {
|
| "english": "EN",
|
| "chinese": "ZH",
|
| }
|
|
|
| @staticmethod
|
| def get_text(text, hps, is_symbol):
|
| text_norm = text_to_sequence(text, hps.symbols, [] if is_symbol else hps.data.text_cleaners)
|
| if hps.data.add_blank:
|
| text_norm = commons.intersperse(text_norm, 0)
|
| text_norm = torch.LongTensor(text_norm)
|
| return text_norm
|
|
|
| @staticmethod
|
| def audio_numpy_concat(segment_data_list, sr, speed=1.):
|
| audio_segments = []
|
| for segment_data in segment_data_list:
|
| audio_segments += segment_data.reshape(-1).tolist()
|
| audio_segments += [0] * int((sr * 0.05)/speed)
|
| audio_segments = np.array(audio_segments).astype(np.float32)
|
| return audio_segments
|
|
|
| @staticmethod
|
| def split_segments_into_pieces(text, language_str):
|
| texts = utils.split_segment(text, language_str=language_str)
|
| print(" > Text split into segments.")
|
| print('\n'.join(texts))
|
| print(" > ===========================")
|
| return texts
|
|
|
| def tts(self, text, output_path, speaker, language='English', speed=1.0):
|
| mark = self.language_marks.get(language.lower(), None)
|
| assert mark is not None, f"language {language} is not supported"
|
|
|
| texts = self.split_segments_into_pieces(text, mark)
|
|
|
| audio_list = []
|
| for t in texts:
|
| t = re.sub(r'([a-z])([A-Z])', r'\1 \2', t)
|
| t = f'[{mark}]{t}[{mark}]'
|
| stn_tst = self.get_text(t, self.hps, False)
|
| device = self.device
|
| speaker_id = self.hps.speakers[speaker]
|
| with torch.no_grad():
|
| x_tst = stn_tst.unsqueeze(0).to(device)
|
| x_tst_lengths = torch.LongTensor([stn_tst.size(0)]).to(device)
|
| sid = torch.LongTensor([speaker_id]).to(device)
|
| audio = self.model.infer(x_tst, x_tst_lengths, sid=sid, noise_scale=0.667, noise_scale_w=0.6,
|
| length_scale=1.0 / speed)[0][0, 0].data.cpu().float().numpy()
|
| audio_list.append(audio)
|
| audio = self.audio_numpy_concat(audio_list, sr=self.hps.data.sampling_rate, speed=speed)
|
|
|
| if output_path is None:
|
| return audio
|
| else:
|
| soundfile.write(output_path, audio, self.hps.data.sampling_rate)
|
|
|
|
|
| class ToneColorConverter(OpenVoiceBaseClass):
|
| def __init__(self, *args, **kwargs):
|
| super().__init__(*args, **kwargs)
|
|
|
|
|
|
|
|
|
|
|
|
|
| self.version = getattr(self.hps, '_version_', "v1")
|
|
|
|
|
|
|
| def extract_se(self, waves, wave_lengths):
|
|
|
| device = self.device
|
| hps = self.hps
|
| gs = []
|
|
|
| for wav_tensor, wav_len in zip(waves, wave_lengths):
|
| y = wav_tensor[:wav_len]
|
| y = y[None, :]
|
| y = spectrogram_torch(y, hps.data.filter_length,
|
| hps.data.sampling_rate, hps.data.hop_length, hps.data.win_length,
|
| center=False).to(device)
|
| with torch.no_grad():
|
| g = self.model.ref_enc(y.transpose(1, 2)).unsqueeze(-1)
|
| gs.append(g.detach())
|
| gs = torch.stack(gs)
|
| gs = gs.squeeze(1).squeeze(-1)
|
| return gs
|
|
|
| def convert(self, src_waves, src_wave_lengths, src_se, tgt_se, tau=0.3, message="default"):
|
| hps = self.hps
|
|
|
| with torch.no_grad():
|
| y = src_waves
|
| spec = spectrogram_torch(y, hps.data.filter_length,
|
| hps.data.sampling_rate, hps.data.hop_length, hps.data.win_length,
|
| center=False).to(self.device)
|
| spec_lengths = src_wave_lengths // hps.data.hop_length
|
| spec_lengths = spec_lengths.clamp(min=1, max=spec.size(2))
|
| audio = self.model.voice_conversion(spec, spec_lengths, sid_src=src_se.unsqueeze(-1), sid_tgt=tgt_se.unsqueeze(-1), tau=tau)[0]
|
| return audio
|
|
|
| def add_watermark(self, audio, message):
|
|
|
| return audio
|
| device = self.device
|
| bits = utils.string_to_bits(message).reshape(-1)
|
| n_repeat = len(bits) // 32
|
|
|
| K = 16000
|
| coeff = 2
|
| for n in range(n_repeat):
|
| trunck = audio[(coeff * n) * K: (coeff * n + 1) * K]
|
| if len(trunck) != K:
|
| print('Audio too short, fail to add watermark')
|
| break
|
| message_npy = bits[n * 32: (n + 1) * 32]
|
|
|
| with torch.no_grad():
|
| signal = torch.FloatTensor(trunck).to(device)[None]
|
| message_tensor = torch.FloatTensor(message_npy).to(device)[None]
|
| signal_wmd_tensor = self.watermark_model.encode(signal, message_tensor)
|
| signal_wmd_npy = signal_wmd_tensor.detach().cpu().squeeze()
|
| audio[(coeff * n) * K: (coeff * n + 1) * K] = signal_wmd_npy
|
| return audio
|
|
|
| def detect_watermark(self, audio, n_repeat):
|
| bits = []
|
| K = 16000
|
| coeff = 2
|
| for n in range(n_repeat):
|
| trunck = audio[(coeff * n) * K: (coeff * n + 1) * K]
|
| if len(trunck) != K:
|
| print('Audio too short, fail to detect watermark')
|
| return 'Fail'
|
| with torch.no_grad():
|
| signal = torch.FloatTensor(trunck).to(self.device).unsqueeze(0)
|
| message_decoded_npy = (self.watermark_model.decode(signal) >= 0.5).int().detach().cpu().numpy().squeeze()
|
| bits.append(message_decoded_npy)
|
| bits = np.stack(bits).reshape(-1, 8)
|
| message = utils.bits_to_string(bits)
|
| return message
|
|
|
|
|