File size: 4,202 Bytes
8f40dc9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 | import sys
import os
current_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.append(current_dir)
from transformers import PreTrainedModel, PretrainedConfig
import torch
import numpy as np
from f5_tts.infer.utils_infer import (
infer_process,
load_model,
load_vocoder,
preprocess_ref_audio_text,
)
from f5_tts.model import DiT
import soundfile as sf
import io
from pydub import AudioSegment, silence
from huggingface_hub import hf_hub_download
import os
class SPRING_F5Config(PretrainedConfig):
model_type = "SPRING_F5"
def __init__(self, ckpt_path: str = "checkpoints/model_170000.pt", vocab_path: str = "checkpoints/vocab.txt",
speed: float = 1.0, remove_sil: bool = True, **kwargs):
super().__init__(**kwargs)
self.ckpt_path = ckpt_path
self.vocab_path = vocab_path
self.speed = speed
self.remove_sil = remove_sil
class SPRING_F5Model(PreTrainedModel):
config_class = SPRING_F5Config
def __init__(self, config):
super().__init__(config)
self._device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Load vocoder
self.vocoder = load_vocoder(vocoder_name="vocos", is_local=False, device=self._device)
ckpt_file = hf_hub_download( repo_id=config.name_or_path, filename=config.ckpt_path)
vocab_path = hf_hub_download(repo_id=config.name_or_path, filename=config.vocab_path )
self.ema_model = load_model(
DiT,
dict(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4),
ckpt_file,
mel_spec_type="vocos",
vocab_file=vocab_path,
device=self._device
)
@classmethod
def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
config = kwargs.pop("config", None)
if config is None:
config = SPRING_F5Config.from_pretrained(pretrained_model_name_or_path, **kwargs)
config.name_or_path = pretrained_model_name_or_path
return cls(config)
def forward(self, text: str, ref_audio_path: str, ref_text: str, lang: str):
"""
Generate speech given a reference audio & text input.
Args:
text (str): The text to be synthesized.
ref_audio_path (str): Path to the reference audio file.
ref_text (str): The reference text.
Returns:
np.array: Generated waveform.
"""
if not os.path.exists(ref_audio_path):
raise FileNotFoundError(f"Reference audio file {ref_audio_path} not found.")
# Load reference audio & text
ref_audio, ref_text = preprocess_ref_audio_text(ref_audio_path, ref_text)
self.ema_model.to(self._device)
self.vocoder.to(self._device)
# Perform inference
audio, final_sample_rate, _ = infer_process(
ref_audio,
ref_text,
text,
self.ema_model,
self.vocoder,
mel_spec_type="vocos",
speed=self.config.speed,
device=self._device,
lang=lang # Language ID is used for number-to-Indic word conversion.
)
# Convert to pydub format and remove silence if needed
buffer = io.BytesIO()
sf.write(buffer, audio, samplerate=24000, format="WAV")
buffer.seek(0)
audio_segment = AudioSegment.from_file(buffer, format="wav")
if self.config.remove_sil:
non_silent_segs = silence.split_on_silence(
audio_segment,
min_silence_len=1000,
silence_thresh=-50,
keep_silence=500,
seek_step=10,
)
non_silent_wave = sum(non_silent_segs, AudioSegment.silent(duration=0))
audio_segment = non_silent_wave
# Normalize loudness
target_dBFS = -20.0
change_in_dBFS = target_dBFS - audio_segment.dBFS
audio_segment = audio_segment.apply_gain(change_in_dBFS)
return np.array(audio_segment.get_array_of_samples())
|