lenML/ChatTTS-Forge

1import sys2from io import BytesIO3 4import numpy as np5import pyrubberband as pyrb6import soundfile as sf7from pydub import AudioSegment, effects8 9INT16_MAX = np.iinfo(np.int16).max10 11 12def audio_to_int16(audio_data: np.ndarray) -> np.ndarray:13    if (14        audio_data.dtype == np.float3215        or audio_data.dtype == np.float6416        or audio_data.dtype == np.float12817        or audio_data.dtype == np.float1618    ):19        audio_data = (audio_data * INT16_MAX).astype(np.int16)20    return audio_data21 22 23def pydub_to_np(audio: AudioSegment) -> tuple[int, np.ndarray]:24    """25    Converts pydub audio segment into np.float32 of shape [duration_in_seconds*sample_rate, channels],26    where each value is in range [-1.0, 1.0].27    Returns tuple (audio_np_array, sample_rate).28    """29    nd_array = np.array(audio.get_array_of_samples(), dtype=np.float32)30    if audio.channels != 1:31        nd_array = nd_array.reshape((-1, audio.channels))32    nd_array = nd_array / (1 << (8 * audio.sample_width - 1))33 34    return (35        audio.frame_rate,36        nd_array,37    )38 39 40def audiosegment_to_librosawav(audiosegment: AudioSegment) -> np.ndarray:41    """42    Converts pydub audio segment into np.float32 of shape [duration_in_seconds*sample_rate, channels],43    where each value is in range [-1.0, 1.0].44    """45    channel_sounds = audiosegment.split_to_mono()46    samples = [s.get_array_of_samples() for s in channel_sounds]47 48    fp_arr = np.array(samples).T.astype(np.float32)49    fp_arr /= np.iinfo(samples[0].typecode).max50    fp_arr = fp_arr.reshape(-1)51 52    return fp_arr53 54 55def ndarray_to_segment(56    ndarray: np.ndarray, frame_rate: int, sample_width: int = None, channels: int = None57) -> AudioSegment:58    buffer = BytesIO()59    sf.write(buffer, ndarray, frame_rate, format="wav", subtype="PCM_16")60    buffer.seek(0)61    sound: AudioSegment = AudioSegment.from_wav(buffer)62 63    if sample_width is None:64        sample_width = sound.sample_width65    if channels is None:66        channels = sound.channels67 68    return (69        sound.set_frame_rate(frame_rate)70        .set_sample_width(sample_width)71        .set_channels(channels)72    )73 74 75def apply_prosody_to_audio_segment(76    audio_segment: AudioSegment,77    rate: float = 1,78    volume: float = 0,79    pitch: int = 0,80    sr: int = 24000,81) -> AudioSegment:82    audio_data = audiosegment_to_librosawav(audio_segment)83 84    audio_data = apply_prosody_to_audio_data(audio_data, rate, volume, pitch, sr)85 86    audio_segment = ndarray_to_segment(87        audio_data, sr, audio_segment.sample_width, audio_segment.channels88    )89 90    return audio_segment91 92 93def apply_prosody_to_audio_data(94    audio_data: np.ndarray,95    rate: float = 1,96    volume: float = 0,97    pitch: float = 0,98    sr: int = 24000,99) -> np.ndarray:100    if rate != 1:101        audio_data = pyrb.time_stretch(audio_data, sr=sr, rate=rate)102 103    if volume != 0:104        audio_data = audio_data * volume105 106    if pitch != 0:107        audio_data = pyrb.pitch_shift(audio_data, sr=sr, n_steps=pitch)108 109    return audio_data110 111 112def apply_normalize(113    audio_data: np.ndarray,114    headroom: float = 1,115    sr: int = 24000,116):117    segment = ndarray_to_segment(audio_data, sr)118    segment = effects.normalize(seg=segment, headroom=headroom)119 120    return pydub_to_np(segment)121 122 123if __name__ == "__main__":124    input_file = sys.argv[1]125 126    time_stretch_factors = [0.5, 0.75, 1.5, 1.0]127    pitch_shift_factors = [-12, -5, 0, 5, 12]128 129    input_sound = AudioSegment.from_mp3(input_file)130 131    for time_factor in time_stretch_factors:132        output_wav = f"{input_file}_time_{time_factor}.wav"133        output_sound = apply_prosody_to_audio_segment(input_sound, rate=time_factor)134        output_sound.export(output_wav, format="wav")135 136    for pitch_factor in pitch_shift_factors:137        output_wav = f"{input_file}_pitch_{pitch_factor}.wav"138        output_sound = apply_prosody_to_audio_segment(input_sound, pitch=pitch_factor)139        output_sound.export(output_wav, format="wav")140