File size: 4,233 Bytes
8f40dc9 898dd2a 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 130 | 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,
use_ema=False,
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())
|