kromskii2/alltalk_tts
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add support for the E2 model to F5-TTS

Browse Source
This commit is contained in:
erew123 authored and GitHub committed 2024-11-05 10:19:45 +00:00
1 parent 1a5945235d
commit 5a81e4818b
3 files changed
+107 -38
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

system/tts_engines/f5tts/model_engine.py
+88 -37
View File
@@ -27,6 +27,7 @@ deepspeed_available = False
# In this section you will import any imports that your specific TTS Engine will use. You will provide any
# start-up errors for those bits, as if you were starting up a normal Python script. Note the logging.disable
# a few lines up from here, you may want to # that out while debugging!
from vocos import Vocos
import soundfile as sf
import tempfile
from pydub import AudioSegment, silence
@@ -38,22 +39,9 @@ import sys
def install_and_restart():
try:
print("#################################################################")
print("Installing required packages for F5-TTS... This may take a while.")
print("################################################################")
# Install vocos first
print("Installing vocos...")
subprocess.check_call([
sys.executable,
"-m",
"pip",
"install",
"vocos"
])
# Then install F5-TTS
print("Installing F5-TTS...")
print("##########################################")
print("F5-TTS not found. Attempting to install...")
print("##########################################")
subprocess.check_call([
sys.executable,
"-m",
@@ -61,10 +49,9 @@ def install_and_restart():
"install",
"git+https://github.com/SWivid/F5-TTS.git"
])
print("##############################################################")
print("All packages installed successfully! Restarting application...")
print("##############################################################")
print("########################################################")
print("F5-TTS installed successfully! Restarting application...")
print("########################################################")
# Get the current script's path
script_path = sys.argv[0]
@@ -74,17 +61,16 @@ def install_and_restart():
except subprocess.CalledProcessError as e:
print("########################################################")
print(f"Failed to install required packages: {str(e)}")
print(f"Failed to install F5-TTS: {str(e)}")
print("########################################################")
raise ImportError("Could not install required packages")
raise ImportError("Could not install required package F5-TTS")
try:
from f5_tts.model import CFM, DiT
from f5_tts.model import CFM, DiT, UNetT
from f5_tts.model.utils import (
get_tokenizer,
convert_char_to_pinyin,
)
from vocos import Vocos
except ImportError:
install_and_restart()
@@ -190,14 +176,21 @@ class tts_class:
self.sway_sampling_coef = -1.0
self.speed = 1.0
# F5-TTS model configuration
self.model_cfg = dict(
self.f5_model_cfg = dict(
dim=1024,
depth=22,
heads=16,
ff_mult=2,
text_dim=512,
conv_layers=4
)
)
# Add E2-TTS model configuration
self.e2_model_cfg = dict(
dim=1024,
depth=24,
heads=16,
ff_mult=4
)
################################################################
# DONT CHANGE # Print out Python, CUDA, DeepSpeed versions ####
################################################################
@@ -387,7 +380,7 @@ class tts_class:
return self.available_models
# Look for model directories that match the pattern f5tts_v*
for model_dir in models_dir.glob("f5tts_v*"):
for model_dir in models_dir.glob("*tts_v*"):
if model_dir.is_dir():
# First try to find model_*.safetensors files
model_files = list(model_dir.glob("model_*.safetensors"))
@@ -507,7 +500,10 @@ class tts_class:
self.vocoder.load_state_dict(vocoder_state)
self.vocoder = self.vocoder.eval()
# Initialize F5-TTS model with same mel_spec_kwargs as original
# Determine if this is an E2-TTS model from the folder name
is_e2_model = "e2tts" in model_folder.lower()
# Initialize model with appropriate architecture
vocab_char_map, vocab_size = get_tokenizer(str(vocab_path), "custom")
mel_spec_kwargs = dict(
@@ -519,14 +515,34 @@ class tts_class:
mel_spec_type=self.mel_spec_type,
)
self.model = CFM(
transformer=DiT(**self.model_cfg, text_num_embeds=vocab_size, mel_dim=self.n_mel_channels),
mel_spec_kwargs=mel_spec_kwargs,
odeint_kwargs=dict(
method=self.ode_method,
),
vocab_char_map=vocab_char_map
).float() # Ensure float32
if is_e2_model:
print(f"[{self.branding}ENG] Initializing E2-TTS model") if self.debug_tts else None
self.model = CFM(
transformer=UNetT(
**self.e2_model_cfg,
text_num_embeds=vocab_size,
mel_dim=self.n_mel_channels
),
mel_spec_kwargs=mel_spec_kwargs,
odeint_kwargs=dict(
method=self.ode_method,
),
vocab_char_map=vocab_char_map
).float() # Ensure float32
else:
print(f"[{self.branding}ENG] Initializing F5-TTS model") if self.debug_tts else None
self.model = CFM(
transformer=DiT(
**self.f5_model_cfg,
text_num_embeds=vocab_size,
mel_dim=self.n_mel_channels
),
mel_spec_kwargs=mel_spec_kwargs,
odeint_kwargs=dict(
method=self.ode_method,
),
vocab_char_map=vocab_char_map
).float() # Ensure float32
# Load model weights
if str(model_path).endswith('.safetensors'):
@@ -572,7 +588,8 @@ class tts_class:
self.vocoder = self.vocoder.to(self.device)
self.is_tts_model_loaded = True
print(f"[{self.branding}ENG] F5-TTS model loaded successfully") if self.debug_tts else None
model_type = "E2-TTS" if is_e2_model else "F5-TTS"
print(f"[{self.branding}ENG] {model_type} model loaded successfully") if self.debug_tts else None
except Exception as e:
print(f"[{self.branding}ENG] \033[91mError loading model: {str(e)}\033[0m")
@@ -895,6 +912,38 @@ class tts_class:
return chunks
#############################################################################################################################
#############################################################################################################################
# DONT CHANGE ME # De-Capitalise words (Not used currently) #################################################################
#############################################################################################################################
#############################################################################################################################
def process_text_for_tts(self, text):
"""
Process text by converting fully capitalized words to lowercase,
while preserving other words' capitalization.
Args:
self: The class instance
text (str): Input text string
Returns:
str: Processed text with fully capitalized words converted to lowercase
"""
# Split text into words
words = text.split()
# Process each word
processed_words = []
for word in words:
# Check if word is fully uppercase and longer than 1 character
if word.isupper() and len(word) > 1:
processed_words.append(word.lower())
else:
processed_words.append(word)
# Join words back together
return ' '.join(processed_words)
##########################################################################################################################################
##########################################################################################################################################
# CHANGE ME # Model changing. Unload out old model and load in a new one #################################################################
@@ -981,6 +1030,8 @@ class tts_class:
str(ref_audio_path), ref_text
)
text = self.process_text_for_tts(text)
# Generate the audio using infer_process
final_wave, final_sample_rate, _ = await self.infer_process(
ref_audio,