Aditya02/IndicF5

1import socket2import struct3import torch4import torchaudio5from threading import Thread6 7 8import gc9import traceback10 11 12from infer.utils_infer import infer_batch_process, preprocess_ref_audio_text, load_vocoder, load_model13from model.backbones.dit import DiT14 15 16class TTSStreamingProcessor:17    def __init__(self, ckpt_file, vocab_file, ref_audio, ref_text, device=None, dtype=torch.float32):18        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")19 20        # Load the model using the provided checkpoint and vocab files21        self.model = load_model(22            model_cls=DiT,23            model_cfg=dict(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4),24            ckpt_path=ckpt_file,25            mel_spec_type="vocos",  # or "bigvgan" depending on vocoder26            vocab_file=vocab_file,27            ode_method="euler",28            use_ema=True,29            device=self.device,30        ).to(self.device, dtype=dtype)31 32        # Load the vocoder33        self.vocoder = load_vocoder(is_local=False)34 35        # Set sampling rate for streaming36        self.sampling_rate = 24000  # Consistency with client37 38        # Set reference audio and text39        self.ref_audio = ref_audio40        self.ref_text = ref_text41 42        # Warm up the model43        self._warm_up()44 45    def _warm_up(self):46        """Warm up the model with a dummy input to ensure it's ready for real-time processing."""47        print("Warming up the model...")48        ref_audio, ref_text = preprocess_ref_audio_text(self.ref_audio, self.ref_text)49        audio, sr = torchaudio.load(ref_audio)50        gen_text = "Warm-up text for the model."51 52        # Pass the vocoder as an argument here53        infer_batch_process((audio, sr), ref_text, [gen_text], self.model, self.vocoder, device=self.device)54        print("Warm-up completed.")55 56    def generate_stream(self, text, play_steps_in_s=0.5):57        """Generate audio in chunks and yield them in real-time."""58        # Preprocess the reference audio and text59        ref_audio, ref_text = preprocess_ref_audio_text(self.ref_audio, self.ref_text)60 61        # Load reference audio62        audio, sr = torchaudio.load(ref_audio)63 64        # Run inference for the input text65        audio_chunk, final_sample_rate, _ = infer_batch_process(66            (audio, sr),67            ref_text,68            [text],69            self.model,70            self.vocoder,71            device=self.device,  # Pass vocoder here72        )73 74        # Break the generated audio into chunks and send them75        chunk_size = int(final_sample_rate * play_steps_in_s)76 77        if len(audio_chunk) < chunk_size:78            packed_audio = struct.pack(f"{len(audio_chunk)}f", *audio_chunk)79            yield packed_audio80            return81 82        for i in range(0, len(audio_chunk), chunk_size):83            chunk = audio_chunk[i : i + chunk_size]84 85            # Check if it's the final chunk86            if i + chunk_size >= len(audio_chunk):87                chunk = audio_chunk[i:]88 89            # Send the chunk if it is not empty90            if len(chunk) > 0:91                packed_audio = struct.pack(f"{len(chunk)}f", *chunk)92                yield packed_audio93 94 95def handle_client(client_socket, processor):96    try:97        while True:98            # Receive data from the client99            data = client_socket.recv(1024).decode("utf-8")100            if not data:101                break102 103            try:104                # The client sends the text input105                text = data.strip()106 107                # Generate and stream audio chunks108                for audio_chunk in processor.generate_stream(text):109                    client_socket.sendall(audio_chunk)110 111                # Send end-of-audio signal112                client_socket.sendall(b"END_OF_AUDIO")113 114            except Exception as inner_e:115                print(f"Error during processing: {inner_e}")116                traceback.print_exc()  # Print the full traceback to diagnose the issue117                break118 119    except Exception as e:120        print(f"Error handling client: {e}")121        traceback.print_exc()122    finally:123        client_socket.close()124 125 126def start_server(host, port, processor):127    server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)128    server.bind((host, port))129    server.listen(5)130    print(f"Server listening on {host}:{port}")131 132    while True:133        client_socket, addr = server.accept()134        print(f"Accepted connection from {addr}")135        client_handler = Thread(target=handle_client, args=(client_socket, processor))136        client_handler.start()137 138 139if __name__ == "__main__":140    try:141        # Load the model and vocoder using the provided files142        ckpt_file = ""  # pointing your checkpoint "ckpts/model/model_1096.pt"143        vocab_file = ""  # Add vocab file path if needed144        ref_audio = ""  # add ref audio"./tests/ref_audio/reference.wav"145        ref_text = ""146 147        # Initialize the processor with the model and vocoder148        processor = TTSStreamingProcessor(149            ckpt_file=ckpt_file,150            vocab_file=vocab_file,151            ref_audio=ref_audio,152            ref_text=ref_text,153            dtype=torch.float32,154        )155 156        # Start the server157        start_server("0.0.0.0", 9998, processor)158    except KeyboardInterrupt:159        gc.collect()160