6Simple9/ChatTTS-OpenVoice

1import torch2import numpy as np3import re4import soundfile5from . import utils6from . import commons7import os8import librosa9from .text import text_to_sequence10from .mel_processing import spectrogram_torch11from .models import SynthesizerTrn12 13 14class OpenVoiceBaseClass(object):15    def __init__(self, 16                config_path, 17                #device='cuda:0'):18                device="cpu"):19        #if 'cuda' in device:20        #    assert torch.cuda.is_available()21 22        hps = utils.get_hparams_from_file(config_path)23 24        model = SynthesizerTrn(25            len(getattr(hps, 'symbols', [])),26            hps.data.filter_length // 2 + 1,27            n_speakers=hps.data.n_speakers,28            **hps.model,29        ).to(device)30 31        model.eval()32        self.model = model33        self.hps = hps34        self.device = device35 36    def load_ckpt(self, ckpt_path):37        checkpoint_dict = torch.load(ckpt_path, map_location=torch.device('cpu'))38        a, b = self.model.load_state_dict(checkpoint_dict['model'], strict=False)39        print("Loaded checkpoint '{}'".format(ckpt_path))40        print('missing/unexpected keys:', a, b)41 42 43class BaseSpeakerTTS(OpenVoiceBaseClass):44    language_marks = {45        "english": "EN",46        "chinese": "ZH",47    }48 49    @staticmethod50    def get_text(text, hps, is_symbol):51        text_norm = text_to_sequence(text, hps.symbols, [] if is_symbol else hps.data.text_cleaners)52        if hps.data.add_blank:53            text_norm = commons.intersperse(text_norm, 0)54        text_norm = torch.LongTensor(text_norm)55        return text_norm56 57    @staticmethod58    def audio_numpy_concat(segment_data_list, sr, speed=1.):59        audio_segments = []60        for segment_data in segment_data_list:61            audio_segments += segment_data.reshape(-1).tolist()62            audio_segments += [0] * int((sr * 0.05)/speed)63        audio_segments = np.array(audio_segments).astype(np.float32)64        return audio_segments65 66    @staticmethod67    def split_sentences_into_pieces(text, language_str):68        texts = utils.split_sentence(text, language_str=language_str)69        print(" > Text splitted to sentences.")70        print('\n'.join(texts))71        print(" > ===========================")72        return texts73 74    def tts(self, text, output_path, speaker, language='English', speed=1.0):75        mark = self.language_marks.get(language.lower(), None)76        assert mark is not None, f"language {language} is not supported"77 78        texts = self.split_sentences_into_pieces(text, mark)79 80        audio_list = []81        for t in texts:82            t = re.sub(r'([a-z])([A-Z])', r'\1 \2', t)83            t = f'[{mark}]{t}[{mark}]'84            stn_tst = self.get_text(t, self.hps, False)85            device = self.device86            speaker_id = self.hps.speakers[speaker]87            with torch.no_grad():88                x_tst = stn_tst.unsqueeze(0).to(device)89                x_tst_lengths = torch.LongTensor([stn_tst.size(0)]).to(device)90                sid = torch.LongTensor([speaker_id]).to(device)91                audio = self.model.infer(x_tst, x_tst_lengths, sid=sid, noise_scale=0.667, noise_scale_w=0.6,92                                    length_scale=1.0 / speed)[0][0, 0].data.cpu().float().numpy()93            audio_list.append(audio)94        audio = self.audio_numpy_concat(audio_list, sr=self.hps.data.sampling_rate, speed=speed)95 96        if output_path is None:97            return audio98        else:99            soundfile.write(output_path, audio, self.hps.data.sampling_rate)100 101 102class ToneColorConverter(OpenVoiceBaseClass):103    def __init__(self, *args, **kwargs):104        super().__init__(*args, **kwargs)105 106        if kwargs.get('enable_watermark', True):107            import wavmark108            self.watermark_model = wavmark.load_model().to(self.device)109        else:110            self.watermark_model = None111 112 113 114    def extract_se(self, ref_wav_list, se_save_path=None):115        if isinstance(ref_wav_list, str):116            ref_wav_list = [ref_wav_list]117        118        device = self.device119        hps = self.hps120        gs = []121        122        for fname in ref_wav_list:123            audio_ref, sr = librosa.load(fname, sr=hps.data.sampling_rate)124            y = torch.FloatTensor(audio_ref)125            y = y.to(device)126            y = y.unsqueeze(0)127            y = spectrogram_torch(y, hps.data.filter_length,128                                        hps.data.sampling_rate, hps.data.hop_length, hps.data.win_length,129                                        center=False).to(device)130            with torch.no_grad():131                g = self.model.ref_enc(y.transpose(1, 2)).unsqueeze(-1)132                gs.append(g.detach())133        gs = torch.stack(gs).mean(0)134 135        if se_save_path is not None:136            os.makedirs(os.path.dirname(se_save_path), exist_ok=True)137            torch.save(gs.cpu(), se_save_path)138 139        return gs140 141    def convert(self, audio_src_path, src_se, tgt_se, output_path=None, tau=0.3, message="@Hilley-MyShell"):142        hps = self.hps143        # load audio144        audio, sample_rate = librosa.load(audio_src_path, sr=hps.data.sampling_rate)145        audio = torch.tensor(audio).float()146        147        with torch.no_grad():148            y = torch.FloatTensor(audio).to(self.device)149            y = y.unsqueeze(0)150            spec = spectrogram_torch(y, hps.data.filter_length,151                                    hps.data.sampling_rate, hps.data.hop_length, hps.data.win_length,152                                    center=False).to(self.device)153            spec_lengths = torch.LongTensor([spec.size(-1)]).to(self.device)154            audio = self.model.voice_conversion(spec, spec_lengths, sid_src=src_se, sid_tgt=tgt_se, tau=tau)[0][155                        0, 0].data.cpu().float().numpy()156            audio = self.add_watermark(audio, message)157            if output_path is None:158                return audio159            else:160                soundfile.write(output_path, audio, hps.data.sampling_rate)161    162    def add_watermark(self, audio, message):163        if self.watermark_model is None:164            return audio165        device = self.device166        bits = utils.string_to_bits(message).reshape(-1)167        n_repeat = len(bits) // 32168 169        K = 16000170        coeff = 2171        for n in range(n_repeat):172            trunck = audio[(coeff * n) * K: (coeff * n + 1) * K]173            if len(trunck) != K:174                print('Audio too short, fail to add watermark')175                break176            message_npy = bits[n * 32: (n + 1) * 32]177            178            with torch.no_grad():179                signal = torch.FloatTensor(trunck).to(device)[None]180                message_tensor = torch.FloatTensor(message_npy).to(device)[None]181                signal_wmd_tensor = self.watermark_model.encode(signal, message_tensor)182                signal_wmd_npy = signal_wmd_tensor.detach().cpu().squeeze()183            audio[(coeff * n) * K: (coeff * n + 1) * K] = signal_wmd_npy184        return audio185 186    def detect_watermark(self, audio, n_repeat):187        bits = []188        K = 16000189        coeff = 2190        for n in range(n_repeat):191            trunck = audio[(coeff * n) * K: (coeff * n + 1) * K]192            if len(trunck) != K:193                print('Audio too short, fail to detect watermark')194                return 'Fail'195            with torch.no_grad():196                signal = torch.FloatTensor(trunck).to(self.device).unsqueeze(0)197                message_decoded_npy = (self.watermark_model.decode(signal) >= 0.5).int().detach().cpu().numpy().squeeze()198            bits.append(message_decoded_npy)199        bits = np.stack(bits).reshape(-1, 8)200        message = utils.bits_to_string(bits)201        return message202    203