6Simple9/ChatTTS-OpenVoice

1import os2import glob3import torch4from glob import glob5import numpy as np6from pydub import AudioSegment7from faster_whisper import WhisperModel8from whisper_timestamped.transcribe import get_audio_tensor, get_vad_segments9 10model_size = "medium"11# Run on GPU with FP1612model = None13def split_audio_whisper(audio_path, target_dir='processed'):14    global model15    if model is None:16        model = WhisperModel(model_size, device="cuda", compute_type="float16")17    audio = AudioSegment.from_file(audio_path)18    max_len = len(audio)19 20    audio_name = os.path.basename(audio_path).rsplit('.', 1)[0]21    target_folder = os.path.join(target_dir, audio_name)22    23    segments, info = model.transcribe(audio_path, beam_size=5, word_timestamps=True)24    segments = list(segments)    25 26    # create directory27    os.makedirs(target_folder, exist_ok=True)28    wavs_folder = os.path.join(target_folder, 'wavs')29    os.makedirs(wavs_folder, exist_ok=True)30 31    # segments32    s_ind = 033    start_time = None34    35    for k, w in enumerate(segments):36        # process with the time37        if k == 0:38            start_time = max(0, w.start)39 40        end_time = w.end41 42        # calculate confidence43        if len(w.words) > 0:44            confidence = sum([s.probability for s in w.words]) / len(w.words)45        else:46            confidence = 0.47        # clean text48        text = w.text.replace('...', '')49 50        # left 0.08s for each audios51        audio_seg = audio[int( start_time * 1000) : min(max_len, int(end_time * 1000) + 80)]52 53        # segment file name54        fname = f"{audio_name}_seg{s_ind}.wav"55 56        # filter out the segment shorter than 1.5s and longer than 20s57        save = audio_seg.duration_seconds > 1.5 and \58                audio_seg.duration_seconds < 20. and \59                len(text) >= 2 and len(text) < 200 60 61        if save:62            output_file = os.path.join(wavs_folder, fname)63            audio_seg.export(output_file, format='wav')64 65        if k < len(segments) - 1:66            start_time = max(0, segments[k+1].start - 0.08)67 68        s_ind = s_ind + 169    return wavs_folder70 71 72def split_audio_vad(audio_path, target_dir, split_seconds=10.0):73    SAMPLE_RATE = 1600074    audio_vad = get_audio_tensor(audio_path)75    segments = get_vad_segments(76        audio_vad,77        output_sample=True,78        min_speech_duration=0.1,79        min_silence_duration=1,80        method="silero",81    )82    segments = [(seg["start"], seg["end"]) for seg in segments]83    segments = [(float(s) / SAMPLE_RATE, float(e) / SAMPLE_RATE) for s,e in segments]84    print(segments)85    audio_active = AudioSegment.silent(duration=0)86    audio = AudioSegment.from_file(audio_path)87 88    for start_time, end_time in segments:89        audio_active += audio[int( start_time * 1000) : int(end_time * 1000)]90    91    audio_dur = audio_active.duration_seconds92    print(f'after vad: dur = {audio_dur}')93    audio_name = os.path.basename(audio_path).rsplit('.', 1)[0]94    target_folder = os.path.join(target_dir, audio_name)95    wavs_folder = os.path.join(target_folder, 'wavs')96    os.makedirs(wavs_folder, exist_ok=True)97    start_time = 0.98    count = 099    num_splits = int(np.round(audio_dur / split_seconds))100    assert num_splits > 0, 'input audio is too short'101    interval = audio_dur / num_splits102 103    for i in range(num_splits):104        end_time = min(start_time + interval, audio_dur)105        if i == num_splits - 1:106            end_time = audio_dur107        output_file = f"{wavs_folder}/{audio_name}_seg{count}.wav"108        audio_seg = audio_active[int(start_time * 1000): int(end_time * 1000)]109        audio_seg.export(output_file, format='wav')110        start_time = end_time111        count += 1112    return wavs_folder113 114 115    116 117 118def get_se(audio_path, vc_model, target_dir='processed', vad=True):119    device = vc_model.device120 121    audio_name = os.path.basename(audio_path).rsplit('.', 1)[0]122    se_path = os.path.join(target_dir, audio_name, 'se.pth')123 124    if os.path.isfile(se_path):125        se = torch.load(se_path).to(device)126        return se, audio_name127    if os.path.isdir(audio_path):128        wavs_folder = audio_path129    elif vad:130        wavs_folder = split_audio_vad(audio_path, target_dir)131    else:132        wavs_folder = split_audio_whisper(audio_path, target_dir)133    134    audio_segs = glob(f'{wavs_folder}/*.wav')135    if len(audio_segs) == 0:136        raise NotImplementedError('No audio segments found!')137    138    return vc_model.extract_se(audio_segs, se_save_path=se_path), audio_name139 140