| 1 | import argparse |
| 2 | import gc |
| 3 | import logging |
| 4 | import queue |
| 5 | import socket |
| 6 | import struct |
| 7 | import threading |
| 8 | import traceback |
| 9 | import wave |
| 10 | from importlib.resources import files |
| 11 | |
| 12 | import numpy as np |
| 13 | import torch |
| 14 | import torchaudio |
| 15 | from huggingface_hub import hf_hub_download |
| 16 | from hydra.utils import get_class |
| 17 | from omegaconf import OmegaConf |
| 18 | |
| 19 | from f5_tts.infer.utils_infer import ( |
| 20 | chunk_text, |
| 21 | infer_batch_process, |
| 22 | load_model, |
| 23 | load_vocoder, |
| 24 | preprocess_ref_audio_text, |
| 25 | ) |
| 26 | |
| 27 | |
| 28 | logging.basicConfig(level=logging.INFO) |
| 29 | logger = logging.getLogger(__name__) |
| 30 | |
| 31 | |
| 32 | class AudioFileWriterThread(threading.Thread): |
| 33 | """Threaded file writer to avoid blocking the TTS streaming process.""" |
| 34 | |
| 35 | def __init__(self, output_file, sampling_rate): |
| 36 | super().__init__() |
| 37 | self.output_file = output_file |
| 38 | self.sampling_rate = sampling_rate |
| 39 | self.queue = queue.Queue() |
| 40 | self.stop_event = threading.Event() |
| 41 | self.audio_data = [] |
| 42 | |
| 43 | def run(self): |
| 44 | """Process queued audio data and write it to a file.""" |
| 45 | logger.info("AudioFileWriterThread started.") |
| 46 | with wave.open(self.output_file, "wb") as wf: |
| 47 | wf.setnchannels(1) |
| 48 | wf.setsampwidth(2) |
| 49 | wf.setframerate(self.sampling_rate) |
| 50 | |
| 51 | while not self.stop_event.is_set() or not self.queue.empty(): |
| 52 | try: |
| 53 | chunk = self.queue.get(timeout=0.1) |
| 54 | if chunk is not None: |
| 55 | chunk = np.int16(chunk * 32767) |
| 56 | self.audio_data.append(chunk) |
| 57 | wf.writeframes(chunk.tobytes()) |
| 58 | except queue.Empty: |
| 59 | continue |
| 60 | |
| 61 | def add_chunk(self, chunk): |
| 62 | """Add a new chunk to the queue.""" |
| 63 | self.queue.put(chunk) |
| 64 | |
| 65 | def stop(self): |
| 66 | """Stop writing and ensure all queued data is written.""" |
| 67 | self.stop_event.set() |
| 68 | self.join() |
| 69 | logger.info("Audio writing completed.") |
| 70 | |
| 71 | |
| 72 | class TTSStreamingProcessor: |
| 73 | def __init__(self, model, ckpt_file, vocab_file, ref_audio, ref_text, device=None, dtype=torch.float32): |
| 74 | self.device = device or ( |
| 75 | "cuda" |
| 76 | if torch.cuda.is_available() |
| 77 | else "xpu" |
| 78 | if torch.xpu.is_available() |
| 79 | else "mps" |
| 80 | if torch.backends.mps.is_available() |
| 81 | else "cpu" |
| 82 | ) |
| 83 | model_cfg = OmegaConf.load(str(files("f5_tts").joinpath(f"configs/{model}.yaml"))) |
| 84 | self.model_cls = get_class(f"f5_tts.model.{model_cfg.model.backbone}") |
| 85 | self.model_arc = model_cfg.model.arch |
| 86 | self.mel_spec_type = model_cfg.model.mel_spec.mel_spec_type |
| 87 | self.sampling_rate = model_cfg.model.mel_spec.target_sample_rate |
| 88 | |
| 89 | self.model = self.load_ema_model(ckpt_file, vocab_file, dtype) |
| 90 | self.vocoder = self.load_vocoder_model() |
| 91 | |
| 92 | self.update_reference(ref_audio, ref_text) |
| 93 | self._warm_up() |
| 94 | self.file_writer_thread = None |
| 95 | self.first_package = True |
| 96 | |
| 97 | def load_ema_model(self, ckpt_file, vocab_file, dtype): |
| 98 | return load_model( |
| 99 | self.model_cls, |
| 100 | self.model_arc, |
| 101 | ckpt_path=ckpt_file, |
| 102 | mel_spec_type=self.mel_spec_type, |
| 103 | vocab_file=vocab_file, |
| 104 | ode_method="euler", |
| 105 | use_ema=True, |
| 106 | device=self.device, |
| 107 | ).to(self.device, dtype=dtype) |
| 108 | |
| 109 | def load_vocoder_model(self): |
| 110 | return load_vocoder(vocoder_name=self.mel_spec_type, is_local=False, local_path=None, device=self.device) |
| 111 | |
| 112 | def update_reference(self, ref_audio, ref_text): |
| 113 | self.ref_audio, self.ref_text = preprocess_ref_audio_text(ref_audio, ref_text) |
| 114 | self.audio, self.sr = torchaudio.load(self.ref_audio) |
| 115 | |
| 116 | ref_audio_duration = self.audio.shape[-1] / self.sr |
| 117 | ref_text_byte_len = len(self.ref_text.encode("utf-8")) |
| 118 | self.max_chars = int(ref_text_byte_len / (ref_audio_duration) * (25 - ref_audio_duration)) |
| 119 | self.few_chars = int(ref_text_byte_len / (ref_audio_duration) * (25 - ref_audio_duration) / 2) |
| 120 | self.min_chars = int(ref_text_byte_len / (ref_audio_duration) * (25 - ref_audio_duration) / 4) |
| 121 | |
| 122 | def _warm_up(self): |
| 123 | logger.info("Warming up the model...") |
| 124 | gen_text = "Warm-up text for the model." |
| 125 | for _ in infer_batch_process( |
| 126 | (self.audio, self.sr), |
| 127 | self.ref_text, |
| 128 | [gen_text], |
| 129 | self.model, |
| 130 | self.vocoder, |
| 131 | progress=None, |
| 132 | device=self.device, |
| 133 | streaming=True, |
| 134 | ): |
| 135 | pass |
| 136 | logger.info("Warm-up completed.") |
| 137 | |
| 138 | def generate_stream(self, text, conn): |
| 139 | text_batches = chunk_text(text, max_chars=self.max_chars) |
| 140 | if self.first_package: |
| 141 | text_batches = chunk_text(text_batches[0], max_chars=self.few_chars) + text_batches[1:] |
| 142 | text_batches = chunk_text(text_batches[0], max_chars=self.min_chars) + text_batches[1:] |
| 143 | self.first_package = False |
| 144 | |
| 145 | audio_stream = infer_batch_process( |
| 146 | (self.audio, self.sr), |
| 147 | self.ref_text, |
| 148 | text_batches, |
| 149 | self.model, |
| 150 | self.vocoder, |
| 151 | progress=None, |
| 152 | device=self.device, |
| 153 | streaming=True, |
| 154 | chunk_size=2048, |
| 155 | ) |
| 156 | |
| 157 | # Reset the file writer thread |
| 158 | if self.file_writer_thread is not None: |
| 159 | self.file_writer_thread.stop() |
| 160 | self.file_writer_thread = AudioFileWriterThread("output.wav", self.sampling_rate) |
| 161 | self.file_writer_thread.start() |
| 162 | |
| 163 | for audio_chunk, _ in audio_stream: |
| 164 | if len(audio_chunk) > 0: |
| 165 | logger.info(f"Generated audio chunk of size: {len(audio_chunk)}") |
| 166 | |
| 167 | # Send audio chunk via socket |
| 168 | conn.sendall(struct.pack(f"{len(audio_chunk)}f", *audio_chunk)) |
| 169 | |
| 170 | # Write to file asynchronously |
| 171 | self.file_writer_thread.add_chunk(audio_chunk) |
| 172 | |
| 173 | logger.info("Finished sending audio stream.") |
| 174 | conn.sendall(b"END") # Send end signal |
| 175 | |
| 176 | # Ensure all audio data is written before exiting |
| 177 | self.file_writer_thread.stop() |
| 178 | |
| 179 | |
| 180 | def handle_client(conn, processor): |
| 181 | try: |
| 182 | with conn: |
| 183 | conn.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) |
| 184 | while True: |
| 185 | data = conn.recv(1024) |
| 186 | if not data: |
| 187 | processor.first_package = True |
| 188 | break |
| 189 | data_str = data.decode("utf-8").strip() |
| 190 | logger.info(f"Received text: {data_str}") |
| 191 | |
| 192 | try: |
| 193 | processor.generate_stream(data_str, conn) |
| 194 | except Exception as inner_e: |
| 195 | logger.error(f"Error during processing: {inner_e}") |
| 196 | traceback.print_exc() |
| 197 | break |
| 198 | except Exception as e: |
| 199 | logger.error(f"Error handling client: {e}") |
| 200 | traceback.print_exc() |
| 201 | |
| 202 | |
| 203 | def start_server(host, port, processor): |
| 204 | with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: |
| 205 | s.bind((host, port)) |
| 206 | s.listen() |
| 207 | logger.info(f"Server started on {host}:{port}") |
| 208 | while True: |
| 209 | conn, addr = s.accept() |
| 210 | logger.info(f"Connected by {addr}") |
| 211 | handle_client(conn, processor) |
| 212 | |
| 213 | |
| 214 | if __name__ == "__main__": |
| 215 | parser = argparse.ArgumentParser() |
| 216 | |
| 217 | parser.add_argument("--host", default="0.0.0.0") |
| 218 | parser.add_argument("--port", default=9998) |
| 219 | |
| 220 | parser.add_argument( |
| 221 | "--model", |
| 222 | default="F5TTS_v1_Base", |
| 223 | help="The model name, e.g. F5TTS_v1_Base", |
| 224 | ) |
| 225 | parser.add_argument( |
| 226 | "--ckpt_file", |
| 227 | default=str(hf_hub_download(repo_id="SWivid/F5-TTS", filename="F5TTS_v1_Base/model_1250000.safetensors")), |
| 228 | help="Path to the model checkpoint file", |
| 229 | ) |
| 230 | parser.add_argument( |
| 231 | "--vocab_file", |
| 232 | default="", |
| 233 | help="Path to the vocab file if customized", |
| 234 | ) |
| 235 | |
| 236 | parser.add_argument( |
| 237 | "--ref_audio", |
| 238 | default=str(files("f5_tts").joinpath("infer/examples/basic/basic_ref_en.wav")), |
| 239 | help="Reference audio to provide model with speaker characteristics", |
| 240 | ) |
| 241 | parser.add_argument( |
| 242 | "--ref_text", |
| 243 | default="", |
| 244 | help="Reference audio subtitle, leave empty to auto-transcribe", |
| 245 | ) |
| 246 | |
| 247 | parser.add_argument("--device", default=None, help="Device to run the model on") |
| 248 | parser.add_argument("--dtype", default=torch.float32, help="Data type to use for model inference") |
| 249 | |
| 250 | args = parser.parse_args() |
| 251 | |
| 252 | try: |
| 253 | # Initialize the processor with the model and vocoder |
| 254 | processor = TTSStreamingProcessor( |
| 255 | model=args.model, |
| 256 | ckpt_file=args.ckpt_file, |
| 257 | vocab_file=args.vocab_file, |
| 258 | ref_audio=args.ref_audio, |
| 259 | ref_text=args.ref_text, |
| 260 | device=args.device, |
| 261 | dtype=args.dtype, |
| 262 | ) |
| 263 | |
| 264 | # Start the server |
| 265 | start_server(args.host, args.port, processor) |
| 266 | |
| 267 | except KeyboardInterrupt: |
| 268 | gc.collect() |
| 269 |