Upload 2 files

src/f5-tts/api.py

@@ -0,0 +1,151 @@
+import random
+import sys
+from importlib.resources import files
+
+import soundfile as sf
+import torch
+import tqdm
+from cached_path import cached_path
+
+from f5_tts.infer.utils_infer import (
+    hop_length,
+    infer_process,
+    load_model,
+    load_vocoder,
+    preprocess_ref_audio_text,
+    remove_silence_for_generated_wav,
+    save_spectrogram,
+    target_sample_rate,
+)
+from f5_tts.model import DiT, UNetT
+from f5_tts.model.utils import seed_everything
+
+
+class F5TTS:
+    def __init__(
+        self,
+        model_type="F5-TTS",
+        ckpt_file="",
+        vocab_file="",
+        ode_method="euler",
+        use_ema=True,
+        vocoder_name="vocos",
+        local_path=None,
+        device=None,
+    ):
+        # Initialize parameters
+        self.final_wave = None
+        self.target_sample_rate = target_sample_rate
+        self.hop_length = hop_length
+        self.seed = -1
+        self.mel_spec_type = vocoder_name
+
+        # Set device
+        self.device = device or (
+            "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
+        )
+
+        # Load models
+        self.load_vocoder_model(vocoder_name, local_path)
+        self.load_ema_model(model_type, ckpt_file, vocoder_name, vocab_file, ode_method, use_ema)
+
+    def load_vocoder_model(self, vocoder_name, local_path):
+        self.vocoder = load_vocoder(vocoder_name, local_path is not None, local_path, self.device)
+
+    def load_ema_model(self, model_type, ckpt_file, mel_spec_type, vocab_file, ode_method, use_ema):
+        if model_type == "F5-TTS":
+            if not ckpt_file:
+                if mel_spec_type == "vocos":
+                    ckpt_file = str(cached_path("hf://SWivid/F5-TTS/F5TTS_Base/model_1200000.safetensors"))
+                elif mel_spec_type == "bigvgan":
+                    ckpt_file = str(cached_path("hf://SWivid/F5-TTS/F5TTS_Base_bigvgan/model_1250000.pt"))
+            model_cfg = dict(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4)
+            model_cls = DiT
+        elif model_type == "E2-TTS":
+            if not ckpt_file:
+                ckpt_file = str(cached_path("hf://SWivid/E2-TTS/E2TTS_Base/model_1200000.safetensors"))
+            model_cfg = dict(dim=1024, depth=24, heads=16, ff_mult=4)
+            model_cls = UNetT
+        else:
+            raise ValueError(f"Unknown model type: {model_type}")
+
+        self.ema_model = load_model(
+            model_cls, model_cfg, ckpt_file, mel_spec_type, vocab_file, ode_method, use_ema, self.device
+        )
+
+    def export_wav(self, wav, file_wave, remove_silence=False):
+        sf.write(file_wave, wav, self.target_sample_rate)
+
+        if remove_silence:
+            remove_silence_for_generated_wav(file_wave)
+
+    def export_spectrogram(self, spect, file_spect):
+        save_spectrogram(spect, file_spect)
+
+    def infer(
+        self,
+        ref_file,
+        ref_text,
+        gen_text,
+        show_info=print,
+        progress=tqdm,
+        target_rms=0.1,
+        cross_fade_duration=0.15,
+        sway_sampling_coef=-1,
+        cfg_strength=2,
+        nfe_step=32,
+        speed=1.0,
+        fix_duration=None,
+        remove_silence=False,
+        file_wave=None,
+        file_spect=None,
+        seed=-1,
+    ):
+        if seed == -1:
+            seed = random.randint(0, sys.maxsize)
+        seed_everything(seed)
+        self.seed = seed
+
+        ref_file, ref_text = preprocess_ref_audio_text(ref_file, ref_text, device=self.device)
+
+        wav, sr, spect = infer_process(
+            ref_file,
+            ref_text,
+            gen_text,
+            self.ema_model,
+            self.vocoder,
+            self.mel_spec_type,
+            show_info=show_info,
+            progress=progress,
+            target_rms=target_rms,
+            cross_fade_duration=cross_fade_duration,
+            nfe_step=nfe_step,
+            cfg_strength=cfg_strength,
+            sway_sampling_coef=sway_sampling_coef,
+            speed=speed,
+            fix_duration=fix_duration,
+            device=self.device,
+        )
+
+        if file_wave is not None:
+            self.export_wav(wav, file_wave, remove_silence)
+
+        if file_spect is not None:
+            self.export_spectrogram(spect, file_spect)
+
+        return wav, sr, spect
+
+
+if __name__ == "__main__":
+    f5tts = F5TTS()
+
+    wav, sr, spect = f5tts.infer(
+        ref_file=str(files("f5_tts").joinpath("infer/examples/basic/basic_ref_en.wav")),
+        ref_text="some call me nature, others call me mother nature.",
+        gen_text="""I don't really care what you call me. I've been a silent spectator, watching species evolve, empires rise and fall. But always remember, I am mighty and enduring. Respect me and I'll nurture you; ignore me and you shall face the consequences.""",
+        file_wave=str(files("f5_tts").joinpath("../../tests/api_out.wav")),
+        file_spect=str(files("f5_tts").joinpath("../../tests/api_out.png")),
+        seed=-1,  # random seed = -1
+    )
+
+    print("seed :", f5tts.seed)

src/f5-tts/socket.py

@@ -0,0 +1,159 @@
+import socket
+import struct
+import torch
+import torchaudio
+from threading import Thread
+
+
+import gc
+import traceback
+
+
+from infer.utils_infer import infer_batch_process, preprocess_ref_audio_text, load_vocoder, load_model
+from model.backbones.dit import DiT
+
+
+class TTSStreamingProcessor:
+    def __init__(self, ckpt_file, vocab_file, ref_audio, ref_text, device=None, dtype=torch.float32):
+        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+
+        # Load the model using the provided checkpoint and vocab files
+        self.model = load_model(
+            DiT,
+            dict(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4),
+            ckpt_file,
+            vocab_file,
+        ).to(self.device, dtype=dtype)
+
+        # Load the vocoder
+        self.vocoder = load_vocoder(is_local=False)
+
+        # Set sampling rate for streaming
+        self.sampling_rate = 24000  # Consistency with client
+
+        # Set reference audio and text
+        self.ref_audio = ref_audio
+        self.ref_text = ref_text
+
+        # Warm up the model
+        self._warm_up()
+
+    def _warm_up(self):
+        """Warm up the model with a dummy input to ensure it's ready for real-time processing."""
+        print("Warming up the model...")
+        ref_audio, ref_text = preprocess_ref_audio_text(self.ref_audio, self.ref_text)
+        audio, sr = torchaudio.load(ref_audio)
+        gen_text = "Warm-up text for the model."
+
+        # Pass the vocoder as an argument here
+        infer_batch_process((audio, sr), ref_text, [gen_text], self.model, self.vocoder, device=self.device)
+        print("Warm-up completed.")
+
+    def generate_stream(self, text, play_steps_in_s=0.5):
+        """Generate audio in chunks and yield them in real-time."""
+        # Preprocess the reference audio and text
+        ref_audio, ref_text = preprocess_ref_audio_text(self.ref_audio, self.ref_text)
+
+        # Load reference audio
+        audio, sr = torchaudio.load(ref_audio)
+
+        # Run inference for the input text
+        audio_chunk, final_sample_rate, _ = infer_batch_process(
+            (audio, sr),
+            ref_text,
+            [text],
+            self.model,
+            self.vocoder,
+            device=self.device,  # Pass vocoder here
+        )
+
+        # Break the generated audio into chunks and send them
+        chunk_size = int(final_sample_rate * play_steps_in_s)
+
+        for i in range(0, len(audio_chunk), chunk_size):
+            chunk = audio_chunk[i : i + chunk_size]
+
+            # Check if it's the final chunk
+            if i + chunk_size >= len(audio_chunk):
+                chunk = audio_chunk[i:]
+
+            # Avoid sending empty or repeated chunks
+            if len(chunk) == 0:
+                break
+
+            # Pack and send the audio chunk
+            packed_audio = struct.pack(f"{len(chunk)}f", *chunk)
+            yield packed_audio
+
+        # Ensure that no final word is repeated by not resending partial chunks
+        if len(audio_chunk) % chunk_size != 0:
+            remaining_chunk = audio_chunk[-(len(audio_chunk) % chunk_size) :]
+            packed_audio = struct.pack(f"{len(remaining_chunk)}f", *remaining_chunk)
+            yield packed_audio
+
+
+def handle_client(client_socket, processor):
+    try:
+        while True:
+            # Receive data from the client
+            data = client_socket.recv(1024).decode("utf-8")
+            if not data:
+                break
+
+            try:
+                # The client sends the text input
+                text = data.strip()
+
+                # Generate and stream audio chunks
+                for audio_chunk in processor.generate_stream(text):
+                    client_socket.sendall(audio_chunk)
+
+                # Send end-of-audio signal
+                client_socket.sendall(b"END_OF_AUDIO")
+
+            except Exception as inner_e:
+                print(f"Error during processing: {inner_e}")
+                traceback.print_exc()  # Print the full traceback to diagnose the issue
+                break
+
+    except Exception as e:
+        print(f"Error handling client: {e}")
+        traceback.print_exc()
+    finally:
+        client_socket.close()
+
+
+def start_server(host, port, processor):
+    server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    server.bind((host, port))
+    server.listen(5)
+    print(f"Server listening on {host}:{port}")
+
+    while True:
+        client_socket, addr = server.accept()
+        print(f"Accepted connection from {addr}")
+        client_handler = Thread(target=handle_client, args=(client_socket, processor))
+        client_handler.start()
+
+
+if __name__ == "__main__":
+    try:
+        # Load the model and vocoder using the provided files
+        ckpt_file = ""  # pointing your checkpoint "ckpts/model/model_1096.pt"
+        vocab_file = ""  # Add vocab file path if needed
+        ref_audio = ""  # add ref audio"./tests/ref_audio/reference.wav"
+        ref_text = ""
+
+        # Initialize the processor with the model and vocoder
+        processor = TTSStreamingProcessor(
+            ckpt_file=ckpt_file,
+            vocab_file=vocab_file,
+            ref_audio=ref_audio,
+            ref_text=ref_text,
+            dtype=torch.float32,
+        )
+
+        # Start the server
+        start_server("0.0.0.0", 9998, processor)
+    except KeyboardInterrupt:
+        gc.collect()