diodiogod
829b0aed81
- Bundle F5-TTS source code to resolve numpy dependency conflicts - Fix all import paths to use bundled engines.f5_tts package - Update ChatterBox model detection to properly scan language subdirectories - Fix missing return statement in F5-TTS engine instantiation - Add missing sys import in Higgs Audio downloader - Update requirements.txt with bundled F5-TTS approach
269 lines
8.9 KiB
Python
269 lines
8.9 KiB
Python
import argparse
|
|
import gc
|
|
import logging
|
|
import queue
|
|
import socket
|
|
import struct
|
|
import threading
|
|
import traceback
|
|
import wave
|
|
from importlib.resources import files
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torchaudio
|
|
from huggingface_hub import hf_hub_download
|
|
from hydra.utils import get_class
|
|
from omegaconf import OmegaConf
|
|
|
|
from .infer.utils_infer import (
|
|
chunk_text,
|
|
infer_batch_process,
|
|
load_model,
|
|
load_vocoder,
|
|
preprocess_ref_audio_text,
|
|
)
|
|
|
|
|
|
logging.basicConfig(level=logging.INFO)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class AudioFileWriterThread(threading.Thread):
|
|
"""Threaded file writer to avoid blocking the TTS streaming process."""
|
|
|
|
def __init__(self, output_file, sampling_rate):
|
|
super().__init__()
|
|
self.output_file = output_file
|
|
self.sampling_rate = sampling_rate
|
|
self.queue = queue.Queue()
|
|
self.stop_event = threading.Event()
|
|
self.audio_data = []
|
|
|
|
def run(self):
|
|
"""Process queued audio data and write it to a file."""
|
|
logger.info("AudioFileWriterThread started.")
|
|
with wave.open(self.output_file, "wb") as wf:
|
|
wf.setnchannels(1)
|
|
wf.setsampwidth(2)
|
|
wf.setframerate(self.sampling_rate)
|
|
|
|
while not self.stop_event.is_set() or not self.queue.empty():
|
|
try:
|
|
chunk = self.queue.get(timeout=0.1)
|
|
if chunk is not None:
|
|
chunk = np.int16(chunk * 32767)
|
|
self.audio_data.append(chunk)
|
|
wf.writeframes(chunk.tobytes())
|
|
except queue.Empty:
|
|
continue
|
|
|
|
def add_chunk(self, chunk):
|
|
"""Add a new chunk to the queue."""
|
|
self.queue.put(chunk)
|
|
|
|
def stop(self):
|
|
"""Stop writing and ensure all queued data is written."""
|
|
self.stop_event.set()
|
|
self.join()
|
|
logger.info("Audio writing completed.")
|
|
|
|
|
|
class TTSStreamingProcessor:
|
|
def __init__(self, model, ckpt_file, vocab_file, ref_audio, ref_text, device=None, dtype=torch.float32):
|
|
self.device = device or (
|
|
"cuda"
|
|
if torch.cuda.is_available()
|
|
else "xpu"
|
|
if torch.xpu.is_available()
|
|
else "mps"
|
|
if torch.backends.mps.is_available()
|
|
else "cpu"
|
|
)
|
|
model_cfg = OmegaConf.load(str(files(__package__ or "f5_tts").joinpath(f"configs/{model}.yaml")))
|
|
self.model_cls = get_class(f"engines.f5_tts.model.{model_cfg.model.backbone}")
|
|
self.model_arc = model_cfg.model.arch
|
|
self.mel_spec_type = model_cfg.model.mel_spec.mel_spec_type
|
|
self.sampling_rate = model_cfg.model.mel_spec.target_sample_rate
|
|
|
|
self.model = self.load_ema_model(ckpt_file, vocab_file, dtype)
|
|
self.vocoder = self.load_vocoder_model()
|
|
|
|
self.update_reference(ref_audio, ref_text)
|
|
self._warm_up()
|
|
self.file_writer_thread = None
|
|
self.first_package = True
|
|
|
|
def load_ema_model(self, ckpt_file, vocab_file, dtype):
|
|
return load_model(
|
|
self.model_cls,
|
|
self.model_arc,
|
|
ckpt_path=ckpt_file,
|
|
mel_spec_type=self.mel_spec_type,
|
|
vocab_file=vocab_file,
|
|
ode_method="euler",
|
|
use_ema=True,
|
|
device=self.device,
|
|
).to(self.device, dtype=dtype)
|
|
|
|
def load_vocoder_model(self):
|
|
return load_vocoder(vocoder_name=self.mel_spec_type, is_local=False, local_path=None, device=self.device)
|
|
|
|
def update_reference(self, ref_audio, ref_text):
|
|
self.ref_audio, self.ref_text = preprocess_ref_audio_text(ref_audio, ref_text)
|
|
self.audio, self.sr = torchaudio.load(self.ref_audio)
|
|
|
|
ref_audio_duration = self.audio.shape[-1] / self.sr
|
|
ref_text_byte_len = len(self.ref_text.encode("utf-8"))
|
|
self.max_chars = int(ref_text_byte_len / (ref_audio_duration) * (25 - ref_audio_duration))
|
|
self.few_chars = int(ref_text_byte_len / (ref_audio_duration) * (25 - ref_audio_duration) / 2)
|
|
self.min_chars = int(ref_text_byte_len / (ref_audio_duration) * (25 - ref_audio_duration) / 4)
|
|
|
|
def _warm_up(self):
|
|
logger.info("Warming up the model...")
|
|
gen_text = "Warm-up text for the model."
|
|
for _ in infer_batch_process(
|
|
(self.audio, self.sr),
|
|
self.ref_text,
|
|
[gen_text],
|
|
self.model,
|
|
self.vocoder,
|
|
progress=None,
|
|
device=self.device,
|
|
streaming=True,
|
|
):
|
|
pass
|
|
logger.info("Warm-up completed.")
|
|
|
|
def generate_stream(self, text, conn):
|
|
text_batches = chunk_text(text, max_chars=self.max_chars)
|
|
if self.first_package:
|
|
text_batches = chunk_text(text_batches[0], max_chars=self.few_chars) + text_batches[1:]
|
|
text_batches = chunk_text(text_batches[0], max_chars=self.min_chars) + text_batches[1:]
|
|
self.first_package = False
|
|
|
|
audio_stream = infer_batch_process(
|
|
(self.audio, self.sr),
|
|
self.ref_text,
|
|
text_batches,
|
|
self.model,
|
|
self.vocoder,
|
|
progress=None,
|
|
device=self.device,
|
|
streaming=True,
|
|
chunk_size=2048,
|
|
)
|
|
|
|
# Reset the file writer thread
|
|
if self.file_writer_thread is not None:
|
|
self.file_writer_thread.stop()
|
|
self.file_writer_thread = AudioFileWriterThread("output.wav", self.sampling_rate)
|
|
self.file_writer_thread.start()
|
|
|
|
for audio_chunk, _ in audio_stream:
|
|
if len(audio_chunk) > 0:
|
|
logger.info(f"Generated audio chunk of size: {len(audio_chunk)}")
|
|
|
|
# Send audio chunk via socket
|
|
conn.sendall(struct.pack(f"{len(audio_chunk)}f", *audio_chunk))
|
|
|
|
# Write to file asynchronously
|
|
self.file_writer_thread.add_chunk(audio_chunk)
|
|
|
|
logger.info("Finished sending audio stream.")
|
|
conn.sendall(b"END") # Send end signal
|
|
|
|
# Ensure all audio data is written before exiting
|
|
self.file_writer_thread.stop()
|
|
|
|
|
|
def handle_client(conn, processor):
|
|
try:
|
|
with conn:
|
|
conn.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
|
|
while True:
|
|
data = conn.recv(1024)
|
|
if not data:
|
|
processor.first_package = True
|
|
break
|
|
data_str = data.decode("utf-8").strip()
|
|
logger.info(f"Received text: {data_str}")
|
|
|
|
try:
|
|
processor.generate_stream(data_str, conn)
|
|
except Exception as inner_e:
|
|
logger.error(f"Error during processing: {inner_e}")
|
|
traceback.print_exc()
|
|
break
|
|
except Exception as e:
|
|
logger.error(f"Error handling client: {e}")
|
|
traceback.print_exc()
|
|
|
|
|
|
def start_server(host, port, processor):
|
|
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
|
s.bind((host, port))
|
|
s.listen()
|
|
logger.info(f"Server started on {host}:{port}")
|
|
while True:
|
|
conn, addr = s.accept()
|
|
logger.info(f"Connected by {addr}")
|
|
handle_client(conn, processor)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
|
|
parser.add_argument("--host", default="0.0.0.0")
|
|
parser.add_argument("--port", default=9998)
|
|
|
|
parser.add_argument(
|
|
"--model",
|
|
default="F5TTS_v1_Base",
|
|
help="The model name, e.g. F5TTS_v1_Base",
|
|
)
|
|
parser.add_argument(
|
|
"--ckpt_file",
|
|
default=str(hf_hub_download(repo_id="SWivid/F5-TTS", filename="F5TTS_v1_Base/model_1250000.safetensors")),
|
|
help="Path to the model checkpoint file",
|
|
)
|
|
parser.add_argument(
|
|
"--vocab_file",
|
|
default="",
|
|
help="Path to the vocab file if customized",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--ref_audio",
|
|
default=str(files(__package__ or "f5_tts").joinpath("infer/examples/basic/basic_ref_en.wav")),
|
|
help="Reference audio to provide model with speaker characteristics",
|
|
)
|
|
parser.add_argument(
|
|
"--ref_text",
|
|
default="",
|
|
help="Reference audio subtitle, leave empty to auto-transcribe",
|
|
)
|
|
|
|
parser.add_argument("--device", default=None, help="Device to run the model on")
|
|
parser.add_argument("--dtype", default=torch.float32, help="Data type to use for model inference")
|
|
|
|
args = parser.parse_args()
|
|
|
|
try:
|
|
# Initialize the processor with the model and vocoder
|
|
processor = TTSStreamingProcessor(
|
|
model=args.model,
|
|
ckpt_file=args.ckpt_file,
|
|
vocab_file=args.vocab_file,
|
|
ref_audio=args.ref_audio,
|
|
ref_text=args.ref_text,
|
|
device=args.device,
|
|
dtype=args.dtype,
|
|
)
|
|
|
|
# Start the server
|
|
start_server(args.host, args.port, processor)
|
|
|
|
except KeyboardInterrupt:
|
|
gc.collect()
|