        y0 = pad_sequence(y0, padding_value=0, batch_first=True)

        # Setup time steps
        t = torch.linspace(0, 1, teacher_steps + 1, device=device, dtype=step_cond.dtype)
        if sway_sampling_coef is not None:
            t = t + sway_sampling_coef * (torch.cos(torch.pi / 2 * t) - 1 + t)
        t = t[:(t > teacher_stopping_time).float().argmax() + 2]
        t = t[:-1]

        # Solve ODE
        trajectory = odeint(fn, y0, t, method="euler")

        if teacher_stopping_time < 1.0:
            # If early stopping, compute final step
            pred = fn(t[-1], trajectory[-1])
            test_out = trajectory[-1] + (1 - t[-1]) * pred
            return test_out
        else:
            return trajectory[-1]

    def _student_sampling(self, x1, cond, text, student_start_step, verbose, sway_coeff = -1):
        """Perform student model sampling."""
        steps = torch.Tensor([0, 0.25, 0.5, 0.75])
        steps = steps + sway_coeff * (torch.cos(torch.pi / 2 * steps) - 1 + steps)
        steps = steps[student_start_step:]

        for step in steps:
            time = torch.Tensor([step]).to(x1.device)

            x0 = torch.randn_like(x1)
            t = time.unsqueeze(-1).unsqueeze(-1)
            phi = (1 - t) * x0 + t * x1

            if verbose:
                print(f'current step: {step}')
            with torch.no_grad():
                pred = self.generator(
                    x=phi,
                    cond=cond,
                    text=text,
                    time=time,
                    drop_audio_cond=False,
                    drop_text=False
                )

                # Predicted mel spectrogram
                output = phi + (1 - t) * pred

            x1 = output

        return x1


fn

        # Pad conditioning
        cond = F.pad(cond, (0, 0, 0, max_duration - cond_seq_len), value=0.0)
        cond_mask = F.pad(cond_mask, (0, max_duration - cond_mask.shape[-1]), value=False)
        cond_mask = cond_mask.unsqueeze(-1)
        step_cond = torch.where(cond_mask, cond, torch.zeros_like(cond))

        return cond, text, step_cond, cond_mask, max_duration, duration

    def _teacher_sampling(self, step_cond, text, cond_mask, max_duration, duration,
                         teacher_steps, teacher_stopping_time, eta, cfg_strength, verbose, sway_sampling_coef = -1):
        """Perform teacher model sampling."""
        device = step_cond.device

        # Pre-generate noise sequence for stochastic sampling
        noise_seq = None
        if eta > 0:
            noise_seq = [torch.randn(1, max_duration, 100, device=device)
                        for _ in range(teacher_steps)]

        def fn(t, x):
            with torch.inference_mode():
                with torch.autocast(device_type="cuda", dtype=torch.float16):
                    if verbose:
                        print(f'current t: {t}')
                    step_frac = 1.0 - t.item()
                    step_idx = min(int(step_frac * len(noise_seq)), len(noise_seq) - 1) if noise_seq else 0

                    # Predict flow
                    pred = self.teacher(
                        x=x, cond=step_cond, text=text, time=t, mask=None,
                        drop_audio_cond=False, drop_text=False
                    )

                    if cfg_strength > 1e-5:
                        null_pred = self.teacher(
                            x=x, cond=step_cond, text=text, time=t, mask=None,
                            drop_audio_cond=True, drop_text=True
                        )
                        pred = pred + (pred - null_pred) * cfg_strength

                    # Add stochasticity if eta > 0
                    if eta > 0 and noise_seq is not None:
                        alpha_t = 1.0 - t.item()
                        sigma_t = t.item()
                        noise_scale = torch.sqrt(torch.tensor(
                            (sigma_t**2) / (alpha_t**2 + sigma_t**2) * eta,
                            device=device
                        ))
                        return pred + noise_scale * noise_seq[step_idx]
                    else:
                        return pred

        # Initialize noise
        y0 = []
        for dur in duration:
            y0.append(torch.randn(dur, 100, device=device, dtype=step_cond.dtype))
        y0 = pad_sequence(y0, padding_value=0, batch_first=True)

        # Setup time steps
        t = torch.linspace(0, 1, teacher_steps + 1, device=device, dtype=step_cond.dtype)
        if sway_sampling_coef is not None:
            t = t + sway_sampling_coef * (torch.cos(torch.pi / 2 * t) - 1 + t)
        t = t[:(t > teacher_stopping_time).float().argmax() + 2]
        t = t[:-1]

        # Solve ODE
        trajectory = odeint(fn, y0, t, method="euler")

        if teacher_stopping_time < 1.0:
            # If early stopping, compute final step
            pred = fn(t[-1], trajectory[-1])


# part of the code is borrowed from https://github.com/lawlict/ECAPA-TDNN

from __future__ import annotations

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchaudio.transforms as trans
from pathlib import Path
from torchaudio.models import Conformer

from f5_tts.model.utils import (
    default,
    exists,
    list_str_to_idx,
    list_str_to_tensor,
    lens_to_mask,
    mask_from_frac_lengths,
)

class ResBlock(nn.Module):
    def __init__(self, hidden_dim, n_conv=3, dropout_p=0.2):
        super().__init__()
        self._n_groups = 8
        self.blocks = nn.ModuleList([
            self._get_conv(hidden_dim, dilation=3**i, dropout_p=dropout_p)
            for i in range(n_conv)])

    def forward(self, x):
        for block in self.blocks:
            res = x
            x = block(x)
            x += res
        return x

    def _get_conv(self, hidden_dim, dilation, dropout_p=0.2):
        layers = [
            nn.Conv1d(hidden_dim, hidden_dim, kernel_size=3, padding=dilation, dilation=dilation),
            nn.ReLU(),
            nn.GroupNorm(num_groups=self._n_groups, num_channels=hidden_dim),
            nn.Dropout(p=dropout_p),
            nn.Conv1d(hidden_dim, hidden_dim, kernel_size=3, padding=1, dilation=1),
            nn.ReLU(),
            nn.Dropout(p=dropout_p)
        ]
        return nn.Sequential(*layers)

class ConformerDiscirminator(nn.Module):
    def __init__(self, input_dim, channels=512, num_layers=3, num_heads=8, depthwise_conv_kernel_size=15, use_group_norm=True):
        super().__init__()

        self.input_layer = nn.Conv1d(input_dim, channels, kernel_size=3, padding=1)

        self.resblock1 = nn.Sequential(
                ResBlock(channels),
                nn.GroupNorm(num_groups=1, num_channels=channels)
            )

        self.resblock2 = nn.Sequential(
                ResBlock(channels),
                nn.GroupNorm(num_groups=1, num_channels=channels)
            )

        self.conformer1 = Conformer(**{"input_dim": channels,
                "num_heads": num_heads,
                "ffn_dim": channels * 2,
                "num_layers": 1,
                "depthwise_conv_kernel_size": depthwise_conv_kernel_size // 2,
                "use_group_norm": use_group_norm})

        self.conformer2 = Conformer(**{"input_dim": channels,
                "num_heads": num_heads,
                "ffn_dim": channels * 2,
                "num_layers": num_layers - 1,
                "depthwise_conv_kernel_size": depthwise_conv_kernel_size,
                "use_group_norm": use_group_norm})

        self.linear = nn.Conv1d(channels, 1, kernel_size=1)

    def forward(self, x):
        # x = torch.stack(x, dim=1).transpose(-1, -2).flatten(start_dim=1, end_dim=2)
        x = torch.cat(x, dim=-1)
        x = x.transpose(1, 2)

        x = self.input_layer(x)

        x = self.resblock1(x)
        x = nn.functional.avg_pool1d(x, 2)
        x = self.resblock2(x)
        x = nn.functional.avg_pool1d(x, 2)

        # Transpose to (B, T, C) for the conformer.
        x = x.transpose(1, 2)
        batch_size, time_steps, _ = x.shape
        # Create a dummy lengths tensor (all sequences are assumed to be full length).
        lengths = torch.full((batch_size,), time_steps, device=x.device, dtype=torch.int64)
        # The built-in Conformer returns (output, output_lengths); we discard lengths.

        x, _ = self.conformer1(x, lengths)
        x, _ = self.conformer2(x, lengths)
        # Transpose back to (B, C, T).
        x = x.transpose(1, 2)

        # out = self.bn(self.pooling(out))
        out = self.linear(x).squeeze(1)

        return out

if __name__ == "__main__":
    from f5_tts.model.utils import get_tokenizer
    from f5_tts.model import DiT

    bsz = 2

    tokenizer = "pinyin"  # 'pinyin', 'char', or 'custom'
    tokenizer_path = None  # if tokenizer = 'custom', define the path to the tokenizer you want to use (should be vocab.txt)
    dataset_name = "Emilia_ZH_EN"
    if tokenizer == "custom":
        tokenizer_path = tokenizer_path
    else:
        tokenizer_path = dataset_name
    vocab_char_map, vocab_size = get_tokenizer(tokenizer_path, tokenizer)

    fake_unet = DiT(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4, text_num_embeds=vocab_size, mel_dim=80)

    fake_unet = fake_unet.cuda()

    text = ["hello world"] * bsz
    lens = torch.randint(1, 1000, (bsz,)).cuda()
    inp = torch.randn(bsz, lens.max(), 80).cuda()

    batch, seq_len, dtype, device = *inp.shape[:2], inp.dtype, inp.device

    batch, seq_len, dtype, device = *inp.shape[:2], inp.dtype, inp.device

    # handle text as string
    if isinstance(text, list):
        if exists(vocab_char_map):
            text = list_str_to_idx(text, vocab_char_map).to(device)
        else:
            text = list_str_to_tensor(text).to(device)
        assert text.shape[0] == batch

    # lens and mask
    if not exists(lens):
        lens = torch.full((batch,), seq_len, device=device)

    mask = lens_to_mask(lens, length=seq_len)  # useless here, as collate_fn will pad to max length in batch
    frac_lengths_mask = (0.7, 1.0)

    # get a random span to mask out for training conditionally
    frac_lengths = torch.zeros((batch,), device=device).float().uniform_(*frac_lengths_mask)
    rand_span_mask = mask_from_frac_lengths(lens, frac_lengths)

    if exists(mask):
        rand_span_mask &= mask

    # Sample a time
    time = torch.rand((batch,), dtype=dtype, device=device)

    x1 = inp
    x0 = torch.randn_like(x1)
    t = time.unsqueeze(-1).unsqueeze(-1)

    phi = (1 - t) * x0 + t * x1
    flow = x1 - x0
    cond = torch.where(rand_span_mask[..., None], torch.zeros_like(x1), x1)

    layers = fake_unet(
        x=phi,
        cond=cond,
        text=text,
        time=time,
        drop_audio_cond=False,
        drop_text=False,
        classify_mode=True
    )

    # layers = torch.stack(layers, dim=1).transpose(-1, -2).flatten(start_dim=1, end_dim=2)
    # print(layers.shape)

    from ctcmodel import ConformerCTC
    ctcmodel = ConformerCTC(vocab_size=vocab_size, mel_dim=80, num_heads=8, d_hid=512, nlayers=6).cuda()
    real_out, layer = ctcmodel(inp)
    layer = layer[-3:] # only use the last 3 layers
    layer = [F.interpolate(l, mode='nearest', scale_factor=4).transpose(-1, -2) for l in layer]
    if layer[0].size(1) < layers[0].size(1):
        layer = [F.pad(l, (0, 0, 0, layers[0].size(1) - l.size(1))) for l in layer]

    layers = layer + layers

    model = ConformerDiscirminator(input_dim=23 * 1024 + 3 * 512,
                            channels=512
                            )

    model = model.cuda()
    print(model)
    out = model(layers)
    print(out.shape)


ResBlock

# part of the code is borrowed from https://github.com/lawlict/ECAPA-TDNN

from __future__ import annotations

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchaudio.transforms as trans
from pathlib import Path
from torchaudio.models import Conformer

from f5_tts.model.utils import (
    default,
    exists,
    list_str_to_idx,
    list_str_to_tensor,
    lens_to_mask,
    mask_from_frac_lengths,
)

class ResBlock(nn.Module):
    def __init__(self, hidden_dim, n_conv=3, dropout_p=0.2):
        super().__init__()
        self._n_groups = 8
        self.blocks = nn.ModuleList([
            self._get_conv(hidden_dim, dilation=3**i, dropout_p=dropout_p)
            for i in range(n_conv)])

    def forward(self, x):
        for block in self.blocks:
            res = x
            x = block(x)
            x += res
        return x

    def _get_conv(self, hidden_dim, dilation, dropout_p=0.2):
        layers = [
            nn.Conv1d(hidden_dim, hidden_dim, kernel_size=3, padding=dilation, dilation=dilation),
            nn.ReLU(),
            nn.GroupNorm(num_groups=self._n_groups, num_channels=hidden_dim),
            nn.Dropout(p=dropout_p),
            nn.Conv1d(hidden_dim, hidden_dim, kernel_size=3, padding=1, dilation=1),
            nn.ReLU(),
            nn.Dropout(p=dropout_p)
        ]
        return nn.Sequential(*layers)

class ConformerDiscirminator(nn.Module):
    def __init__(self, input_dim, channels=512, num_layers=3, num_heads=8, depthwise_conv_kernel_size=15, use_group_norm=True):
        super().__init__()

        self.input_layer = nn.Conv1d(input_dim, channels, kernel_size=3, padding=1)

        self.resblock1 = nn.Sequential(
                ResBlock(channels),
                nn.GroupNorm(num_groups=1, num_channels=channels)
            )

        self.resblock2 = nn.Sequential(
                ResBlock(channels),
                nn.GroupNorm(num_groups=1, num_channels=channels)
            )

        self.conformer1 = Conformer(**{"input_dim": channels,
                "num_heads": num_heads,
                "ffn_dim": channels * 2,


ConformerDiscirminator

    def forward(self, x):
        for block in self.blocks:
            res = x
            x = block(x)
            x += res
        return x

    def _get_conv(self, hidden_dim, dilation, dropout_p=0.2):
        layers = [
            nn.Conv1d(hidden_dim, hidden_dim, kernel_size=3, padding=dilation, dilation=dilation),
            nn.ReLU(),
            nn.GroupNorm(num_groups=self._n_groups, num_channels=hidden_dim),
            nn.Dropout(p=dropout_p),
            nn.Conv1d(hidden_dim, hidden_dim, kernel_size=3, padding=1, dilation=1),
            nn.ReLU(),
            nn.Dropout(p=dropout_p)
        ]
        return nn.Sequential(*layers)

class ConformerDiscirminator(nn.Module):
    def __init__(self, input_dim, channels=512, num_layers=3, num_heads=8, depthwise_conv_kernel_size=15, use_group_norm=True):
        super().__init__()

        self.input_layer = nn.Conv1d(input_dim, channels, kernel_size=3, padding=1)

        self.resblock1 = nn.Sequential(
                ResBlock(channels),
                nn.GroupNorm(num_groups=1, num_channels=channels)
            )

        self.resblock2 = nn.Sequential(
                ResBlock(channels),
                nn.GroupNorm(num_groups=1, num_channels=channels)
            )

        self.conformer1 = Conformer(**{"input_dim": channels,
                "num_heads": num_heads,
                "ffn_dim": channels * 2,
                "num_layers": 1,
                "depthwise_conv_kernel_size": depthwise_conv_kernel_size // 2,
                "use_group_norm": use_group_norm})

        self.conformer2 = Conformer(**{"input_dim": channels,
                "num_heads": num_heads,
                "ffn_dim": channels * 2,
                "num_layers": num_layers - 1,
                "depthwise_conv_kernel_size": depthwise_conv_kernel_size,
                "use_group_norm": use_group_norm})

        self.linear = nn.Conv1d(channels, 1, kernel_size=1)

    def forward(self, x):
        # x = torch.stack(x, dim=1).transpose(-1, -2).flatten(start_dim=1, end_dim=2)
        x = torch.cat(x, dim=-1)
        x = x.transpose(1, 2)

        x = self.input_layer(x)

        x = self.resblock1(x)
        x = nn.functional.avg_pool1d(x, 2)
        x = self.resblock2(x)
        x = nn.functional.avg_pool1d(x, 2)

        # Transpose to (B, T, C) for the conformer.
        x = x.transpose(1, 2)
        batch_size, time_steps, _ = x.shape
        # Create a dummy lengths tensor (all sequences are assumed to be full length).
        lengths = torch.full((batch_size,), time_steps, device=x.device, dtype=torch.int64)
        # The built-in Conformer returns (output, output_lengths); we discard lengths.

        x, _ = self.conformer1(x, lengths)
        x, _ = self.conformer2(x, lengths)
        # Transpose back to (B, C, T).
        x = x.transpose(1, 2)

        # out = self.bn(self.pooling(out))
        out = self.linear(x).squeeze(1)

        return out

if __name__ == "__main__":
    from f5_tts.model.utils import get_tokenizer
    from f5_tts.model import DiT

    bsz = 2

    tokenizer = "pinyin"  # 'pinyin', 'char', or 'custom'
    tokenizer_path = None  # if tokenizer = 'custom', define the path to the tokenizer you want to use (should be vocab.txt)
    dataset_name = "Emilia_ZH_EN"
    if tokenizer == "custom":
        tokenizer_path = tokenizer_path
    else:
        tokenizer_path = dataset_name
    vocab_char_map, vocab_size = get_tokenizer(tokenizer_path, tokenizer)

    fake_unet = DiT(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4, text_num_embeds=vocab_size, mel_dim=80)

    fake_unet = fake_unet.cuda()


    def forward(self, x):
        for block in self.blocks:
            res = x
            x = block(x)
            x += res
        return x

    def _get_conv(self, hidden_dim, dilation, dropout_p=0.2):
        layers = [
            nn.Conv1d(hidden_dim, hidden_dim, kernel_size=3, padding=dilation, dilation=dilation),
            nn.ReLU(),
            nn.GroupNorm(num_groups=self._n_groups, num_channels=hidden_dim),
            nn.Dropout(p=dropout_p),
            nn.Conv1d(hidden_dim, hidden_dim, kernel_size=3, padding=1, dilation=1),
            nn.ReLU(),
            nn.Dropout(p=dropout_p)
        ]
        return nn.Sequential(*layers)

class ConformerDiscirminator(nn.Module):
    def __init__(self, input_dim, channels=512, num_layers=3, num_heads=8, depthwise_conv_kernel_size=15, use_group_norm=True):
        super().__init__()

        self.input_layer = nn.Conv1d(input_dim, channels, kernel_size=3, padding=1)

        self.resblock1 = nn.Sequential(
                ResBlock(channels),
                nn.GroupNorm(num_groups=1, num_channels=channels)
            )

        self.resblock2 = nn.Sequential(
                ResBlock(channels),
                nn.GroupNorm(num_groups=1, num_channels=channels)
            )

        self.conformer1 = Conformer(**{"input_dim": channels,
                "num_heads": num_heads,
                "ffn_dim": channels * 2,
                "num_layers": 1,
                "depthwise_conv_kernel_size": depthwise_conv_kernel_size // 2,
                "use_group_norm": use_group_norm})

        self.conformer2 = Conformer(**{"input_dim": channels,
                "num_heads": num_heads,
                "ffn_dim": channels * 2,
                "num_layers": num_layers - 1,
                "depthwise_conv_kernel_size": depthwise_conv_kernel_size,
                "use_group_norm": use_group_norm})

        self.linear = nn.Conv1d(channels, 1, kernel_size=1)

    def forward(self, x):
        # x = torch.stack(x, dim=1).transpose(-1, -2).flatten(start_dim=1, end_dim=2)
        x = torch.cat(x, dim=-1)
        x = x.transpose(1, 2)

        x = self.input_layer(x)

        x = self.resblock1(x)
        x = nn.functional.avg_pool1d(x, 2)
        x = self.resblock2(x)
        x = nn.functional.avg_pool1d(x, 2)

        # Transpose to (B, T, C) for the conformer.
        x = x.transpose(1, 2)
        batch_size, time_steps, _ = x.shape
        # Create a dummy lengths tensor (all sequences are assumed to be full length).
        lengths = torch.full((batch_size,), time_steps, device=x.device, dtype=torch.int64)
        # The built-in Conformer returns (output, output_lengths); we discard lengths.


forward

                ResBlock(channels),
                nn.GroupNorm(num_groups=1, num_channels=channels)
            )

        self.conformer1 = Conformer(**{"input_dim": channels,
                "num_heads": num_heads,
                "ffn_dim": channels * 2,
                "num_layers": 1,
                "depthwise_conv_kernel_size": depthwise_conv_kernel_size // 2,
                "use_group_norm": use_group_norm})

        self.conformer2 = Conformer(**{"input_dim": channels,
                "num_heads": num_heads,
                "ffn_dim": channels * 2,
                "num_layers": num_layers - 1,
                "depthwise_conv_kernel_size": depthwise_conv_kernel_size,
                "use_group_norm": use_group_norm})

        self.linear = nn.Conv1d(channels, 1, kernel_size=1)

    def forward(self, x):
        # x = torch.stack(x, dim=1).transpose(-1, -2).flatten(start_dim=1, end_dim=2)
        x = torch.cat(x, dim=-1)
        x = x.transpose(1, 2)

        x = self.input_layer(x)

        x = self.resblock1(x)
        x = nn.functional.avg_pool1d(x, 2)
        x = self.resblock2(x)
        x = nn.functional.avg_pool1d(x, 2)

        # Transpose to (B, T, C) for the conformer.
        x = x.transpose(1, 2)
        batch_size, time_steps, _ = x.shape
        # Create a dummy lengths tensor (all sequences are assumed to be full length).
        lengths = torch.full((batch_size,), time_steps, device=x.device, dtype=torch.int64)
        # The built-in Conformer returns (output, output_lengths); we discard lengths.

        x, _ = self.conformer1(x, lengths)
        x, _ = self.conformer2(x, lengths)
        # Transpose back to (B, C, T).
        x = x.transpose(1, 2)

        # out = self.bn(self.pooling(out))
        out = self.linear(x).squeeze(1)

        return out

if __name__ == "__main__":
    from f5_tts.model.utils import get_tokenizer
    from f5_tts.model import DiT

    bsz = 2

    tokenizer = "pinyin"  # 'pinyin', 'char', or 'custom'
    tokenizer_path = None  # if tokenizer = 'custom', define the path to the tokenizer you want to use (should be vocab.txt)
    dataset_name = "Emilia_ZH_EN"
    if tokenizer == "custom":
        tokenizer_path = tokenizer_path
    else:
        tokenizer_path = dataset_name
    vocab_char_map, vocab_size = get_tokenizer(tokenizer_path, tokenizer)

    fake_unet = DiT(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4, text_num_embeds=vocab_size, mel_dim=80)

    fake_unet = fake_unet.cuda()


    lens_to_mask,
    mask_from_frac_lengths,
)

class ResBlock(nn.Module):
    def __init__(self, hidden_dim, n_conv=3, dropout_p=0.2):
        super().__init__()
        self._n_groups = 8
        self.blocks = nn.ModuleList([
            self._get_conv(hidden_dim, dilation=3**i, dropout_p=dropout_p)
            for i in range(n_conv)])

    def forward(self, x):
        for block in self.blocks:
            res = x
            x = block(x)
            x += res
        return x

    def _get_conv(self, hidden_dim, dilation, dropout_p=0.2):
        layers = [
            nn.Conv1d(hidden_dim, hidden_dim, kernel_size=3, padding=dilation, dilation=dilation),
            nn.ReLU(),
            nn.GroupNorm(num_groups=self._n_groups, num_channels=hidden_dim),
            nn.Dropout(p=dropout_p),
            nn.Conv1d(hidden_dim, hidden_dim, kernel_size=3, padding=1, dilation=1),
            nn.ReLU(),
            nn.Dropout(p=dropout_p)
        ]
        return nn.Sequential(*layers)

class ConformerDiscirminator(nn.Module):
    def __init__(self, input_dim, channels=512, num_layers=3, num_heads=8, depthwise_conv_kernel_size=15, use_group_norm=True):
        super().__init__()

        self.input_layer = nn.Conv1d(input_dim, channels, kernel_size=3, padding=1)

        self.resblock1 = nn.Sequential(
                ResBlock(channels),
                nn.GroupNorm(num_groups=1, num_channels=channels)
            )

        self.resblock2 = nn.Sequential(
                ResBlock(channels),
                nn.GroupNorm(num_groups=1, num_channels=channels)
            )

        self.conformer1 = Conformer(**{"input_dim": channels,
                "num_heads": num_heads,
                "ffn_dim": channels * 2,


src.ctcmodel

from torch import nn
import torch
import copy

from pathlib import Path
from torchaudio.models import Conformer

from f5_tts.model.utils import default
from f5_tts.model.utils import exists
from f5_tts.model.utils import list_str_to_idx
from f5_tts.model.utils import list_str_to_tensor
from f5_tts.model.utils import lens_to_mask
from f5_tts.model.utils import mask_from_frac_lengths

from f5_tts.model.utils import (
    default,
    exists,
    list_str_to_idx,
    list_str_to_tensor,
    lens_to_mask,
    mask_from_frac_lengths,
)

class ResBlock(nn.Module):
    def __init__(self, hidden_dim, n_conv=3, dropout_p=0.2):
        super().__init__()
        self._n_groups = 8
        self.blocks = nn.ModuleList([
            self._get_conv(hidden_dim, dilation=3**i, dropout_p=dropout_p)
            for i in range(n_conv)])

    def forward(self, x):
        for block in self.blocks:
            res = x
            x = block(x)
            x += res
        return x

    def _get_conv(self, hidden_dim, dilation, dropout_p=0.2):
        layers = [
            nn.Conv1d(hidden_dim, hidden_dim, kernel_size=3, padding=dilation, dilation=dilation),
            nn.ReLU(),
            nn.GroupNorm(num_groups=self._n_groups, num_channels=hidden_dim),
            nn.Dropout(p=dropout_p),
            nn.Conv1d(hidden_dim, hidden_dim, kernel_size=3, padding=1, dilation=1),
            nn.ReLU(),
            nn.Dropout(p=dropout_p)
        ]
        return nn.Sequential(*layers)

class ConformerCTC(nn.Module):
    def __init__(self,
                 vocab_size,
                 mel_dim=100,
                 num_heads=8,
                 d_hid=512,
                 nlayers=6):
        super().__init__()

        self.mel_proj = nn.Conv1d(mel_dim, d_hid, kernel_size=3, padding=1)

        self.d_hid = d_hid

        self.resblock1 = nn.Sequential(
                ResBlock(d_hid),
                nn.GroupNorm(num_groups=1, num_channels=d_hid)
            )

        self.resblock2 = nn.Sequential(
                ResBlock(d_hid),
                nn.GroupNorm(num_groups=1, num_channels=d_hid)
            )

        self.conf_pre = torch.nn.ModuleList(
            [Conformer(
             input_dim=d_hid,
             num_heads=num_heads,
             ffn_dim=d_hid * 2,
             num_layers=1,
             depthwise_conv_kernel_size=15,
             use_group_norm=True,)
                for _ in range(nlayers // 2)
            ]
        )

        self.conf_after = torch.nn.ModuleList(
            [Conformer(
             input_dim=d_hid,
             num_heads=num_heads,
             ffn_dim=d_hid * 2,
             num_layers=1,
             depthwise_conv_kernel_size=7,
             use_group_norm=True,)
                for _ in range(nlayers // 2)
            ]
        )

        self.out = nn.Linear(d_hid, 1 + vocab_size) # 1 for blank

        self.ctc_loss = nn.CTCLoss(blank=vocab_size, zero_infinity=True).cuda()

    def forward(self, latent, text=None, text_lens=None):
        layers = []

        x = self.mel_proj(latent.transpose(-1, -2)).transpose(-1, -2)

        x = x.transpose(1, 2)
        layers.append(nn.functional.avg_pool1d(x, 4))
        # x = x.transpose(1, 2)

        x = self.resblock1(x)
        x = nn.functional.avg_pool1d(x, 2)
        layers.append(nn.functional.avg_pool1d(x, 2))
        x = self.resblock2(x)
        x = nn.functional.avg_pool1d(x, 2)
        layers.append(x)

        x = x.transpose(1, 2)

        batch_size, time_steps, _ = x.shape
        # Create a dummy lengths tensor (all sequences are assumed to be full length).
        input_lengths = torch.full((batch_size,), time_steps, device=x.device, dtype=torch.int64)

        for layer in (self.conf_pre):
            x, _ = layer(x, input_lengths)
            layers.append(x.transpose(1, 2))

        for layer in (self.conf_after):
            x, _ = layer(x, input_lengths)
            layers.append(x.transpose(1, 2))

        x = self.out(x)

        if text_lens is not None and text is not None:
            loss = self.ctc_loss(x.log_softmax(dim=2).transpose(0, 1), text, input_lengths, text_lens)
            return x, layers, loss
        else:
            return x, layers

if __name__ == "__main__":
    from f5_tts.model.utils import get_tokenizer

    bsz = 16

    tokenizer = "pinyin"  # 'pinyin', 'char', or 'custom'
    tokenizer_path = None  # if tokenizer = 'custom', define the path to the tokenizer you want to use (should be vocab.txt)
    dataset_name = "Emilia_ZH_EN"
    if tokenizer == "custom":
        tokenizer_path = tokenizer_path
    else:
        tokenizer_path = dataset_name
    vocab_char_map, vocab_size = get_tokenizer(tokenizer_path, tokenizer)

    model = ConformerCTC(vocab_size, mel_dim=80, num_heads=8, d_hid=512, nlayers=6).cuda()

    text = ["hello world"] * bsz
    lens = torch.randint(1, 1000, (bsz,)).cuda()
    inp = torch.randn(bsz, lens.max(), 80).cuda()

    batch, seq_len, dtype, device = *inp.shape[:2], inp.dtype, inp.device

    # handle text as string
    text_lens = torch.tensor([len(t) for t in text], device=device)
    if isinstance(text, list):
        if exists(vocab_char_map):
            text = list_str_to_idx(text, vocab_char_map).to(device)
        else:
            text = list_str_to_tensor(text).to(device)
        assert text.shape[0] == batch

    # lens and mask
    if not exists(lens):
        lens = torch.full((batch,), seq_len, device=device)

    out, layers, loss = model(inp, text_lens)

    print(out.shape)
    print(out)
    print(len(layers))
    print(torch.stack(layers, axis=1).shape)
    print(loss)

    probs = out.softmax(dim=2)  # Convert logits to probabilities

    # Greedy decoding
    best_path = torch.argmax(probs, dim=2)

    decoded_sequences = []
    blank_idx = vocab_size

    char_vocab_map = list(vocab_char_map.keys())

    for batch in best_path:
        decoded_sequence = []
        previous_token = None

        for token in batch:
            if token != previous_token:  # Collapse repeated tokens
                if token != blank_idx:  # Ignore blank tokens
                    decoded_sequence.append(token.item())
            previous_token = token

        decoded_sequences.append(decoded_sequence)

    # Convert token indices to characters
    decoded_texts = [''.join([char_vocab_map[token] for token in sequence]) for sequence in decoded_sequences]
    gt_texts = []
    for i in range(text_lens.size(0)):
        gt_texts.append(''.join([char_vocab_map[token] for token in text[i, :text_lens[i]]]))

    print(decoded_texts)
    print(gt_texts)


ResBlock

from torchaudio.models import Conformer

from f5_tts.model.utils import default
from f5_tts.model.utils import exists
from f5_tts.model.utils import list_str_to_idx
from f5_tts.model.utils import list_str_to_tensor
from f5_tts.model.utils import lens_to_mask
from f5_tts.model.utils import mask_from_frac_lengths

from f5_tts.model.utils import (
    default,
    exists,
    list_str_to_idx,
    list_str_to_tensor,
    lens_to_mask,
    mask_from_frac_lengths,
)

class ResBlock(nn.Module):
    def __init__(self, hidden_dim, n_conv=3, dropout_p=0.2):
        super().__init__()
        self._n_groups = 8
        self.blocks = nn.ModuleList([
            self._get_conv(hidden_dim, dilation=3**i, dropout_p=dropout_p)
            for i in range(n_conv)])

    def forward(self, x):
        for block in self.blocks:
            res = x
            x = block(x)
            x += res
        return x

    def _get_conv(self, hidden_dim, dilation, dropout_p=0.2):
        layers = [
            nn.Conv1d(hidden_dim, hidden_dim, kernel_size=3, padding=dilation, dilation=dilation),
            nn.ReLU(),
            nn.GroupNorm(num_groups=self._n_groups, num_channels=hidden_dim),
            nn.Dropout(p=dropout_p),
            nn.Conv1d(hidden_dim, hidden_dim, kernel_size=3, padding=1, dilation=1),
            nn.ReLU(),
            nn.Dropout(p=dropout_p)
        ]
        return nn.Sequential(*layers)

class ConformerCTC(nn.Module):
    def __init__(self,
                 vocab_size,
                 mel_dim=100,
                 num_heads=8,
                 d_hid=512,
                 nlayers=6):
        super().__init__()

        self.mel_proj = nn.Conv1d(mel_dim, d_hid, kernel_size=3, padding=1)

        self.d_hid = d_hid

        self.resblock1 = nn.Sequential(
                ResBlock(d_hid),
                nn.GroupNorm(num_groups=1, num_channels=d_hid)
            )


ConformerCTC

    def forward(self, x):
        for block in self.blocks:
            res = x
            x = block(x)
            x += res
        return x

    def _get_conv(self, hidden_dim, dilation, dropout_p=0.2):
        layers = [
            nn.Conv1d(hidden_dim, hidden_dim, kernel_size=3, padding=dilation, dilation=dilation),
            nn.ReLU(),
            nn.GroupNorm(num_groups=self._n_groups, num_channels=hidden_dim),
            nn.Dropout(p=dropout_p),
            nn.Conv1d(hidden_dim, hidden_dim, kernel_size=3, padding=1, dilation=1),
            nn.ReLU(),
            nn.Dropout(p=dropout_p)
        ]
        return nn.Sequential(*layers)

class ConformerCTC(nn.Module):
    def __init__(self,
                 vocab_size,
                 mel_dim=100,
                 num_heads=8,
                 d_hid=512,
                 nlayers=6):
        super().__init__()

        self.mel_proj = nn.Conv1d(mel_dim, d_hid, kernel_size=3, padding=1)

        self.d_hid = d_hid

        self.resblock1 = nn.Sequential(
                ResBlock(d_hid),
                nn.GroupNorm(num_groups=1, num_channels=d_hid)
            )

        self.resblock2 = nn.Sequential(
                ResBlock(d_hid),
                nn.GroupNorm(num_groups=1, num_channels=d_hid)
            )

        self.conf_pre = torch.nn.ModuleList(
            [Conformer(
             input_dim=d_hid,
             num_heads=num_heads,
             ffn_dim=d_hid * 2,
             num_layers=1,
             depthwise_conv_kernel_size=15,
             use_group_norm=True,)
                for _ in range(nlayers // 2)
            ]
        )

        self.conf_after = torch.nn.ModuleList(
            [Conformer(
             input_dim=d_hid,
             num_heads=num_heads,
             ffn_dim=d_hid * 2,
             num_layers=1,
             depthwise_conv_kernel_size=7,
             use_group_norm=True,)
                for _ in range(nlayers // 2)
            ]
        )

        self.out = nn.Linear(d_hid, 1 + vocab_size) # 1 for blank

        self.ctc_loss = nn.CTCLoss(blank=vocab_size, zero_infinity=True).cuda()

    def forward(self, latent, text=None, text_lens=None):
        layers = []

        x = self.mel_proj(latent.transpose(-1, -2)).transpose(-1, -2)

        x = x.transpose(1, 2)
        layers.append(nn.functional.avg_pool1d(x, 4))
        # x = x.transpose(1, 2)

        x = self.resblock1(x)
        x = nn.functional.avg_pool1d(x, 2)
        layers.append(nn.functional.avg_pool1d(x, 2))
        x = self.resblock2(x)
        x = nn.functional.avg_pool1d(x, 2)
        layers.append(x)

        x = x.transpose(1, 2)

        batch_size, time_steps, _ = x.shape
        # Create a dummy lengths tensor (all sequences are assumed to be full length).
        input_lengths = torch.full((batch_size,), time_steps, device=x.device, dtype=torch.int64)

        for layer in (self.conf_pre):
            x, _ = layer(x, input_lengths)
            layers.append(x.transpose(1, 2))

        for layer in (self.conf_after):
            x, _ = layer(x, input_lengths)
            layers.append(x.transpose(1, 2))

        x = self.out(x)

        if text_lens is not None and text is not None:
            loss = self.ctc_loss(x.log_softmax(dim=2).transpose(0, 1), text, input_lengths, text_lens)
            return x, layers, loss
        else:
            return x, layers

if __name__ == "__main__":
    from f5_tts.model.utils import get_tokenizer

    bsz = 16

    tokenizer = "pinyin"  # 'pinyin', 'char', or 'custom'
    tokenizer_path = None  # if tokenizer = 'custom', define the path to the tokenizer you want to use (should be vocab.txt)
    dataset_name = "Emilia_ZH_EN"
    if tokenizer == "custom":
        tokenizer_path = tokenizer_path
    else:
        tokenizer_path = dataset_name
    vocab_char_map, vocab_size = get_tokenizer(tokenizer_path, tokenizer)

    model = ConformerCTC(vocab_size, mel_dim=80, num_heads=8, d_hid=512, nlayers=6).cuda()

    text = ["hello world"] * bsz


        for block in self.blocks:
            res = x
            x = block(x)
            x += res
        return x

    def _get_conv(self, hidden_dim, dilation, dropout_p=0.2):
        layers = [
            nn.Conv1d(hidden_dim, hidden_dim, kernel_size=3, padding=dilation, dilation=dilation),
            nn.ReLU(),
            nn.GroupNorm(num_groups=self._n_groups, num_channels=hidden_dim),
            nn.Dropout(p=dropout_p),
            nn.Conv1d(hidden_dim, hidden_dim, kernel_size=3, padding=1, dilation=1),
            nn.ReLU(),
            nn.Dropout(p=dropout_p)
        ]
        return nn.Sequential(*layers)

class ConformerCTC(nn.Module):
    def __init__(self,
                 vocab_size,
                 mel_dim=100,
                 num_heads=8,
                 d_hid=512,
                 nlayers=6):
        super().__init__()

        self.mel_proj = nn.Conv1d(mel_dim, d_hid, kernel_size=3, padding=1)

        self.d_hid = d_hid

        self.resblock1 = nn.Sequential(
                ResBlock(d_hid),
                nn.GroupNorm(num_groups=1, num_channels=d_hid)
            )

        self.resblock2 = nn.Sequential(
                ResBlock(d_hid),
                nn.GroupNorm(num_groups=1, num_channels=d_hid)
            )

        self.conf_pre = torch.nn.ModuleList(
            [Conformer(
             input_dim=d_hid,
             num_heads=num_heads,
             ffn_dim=d_hid * 2,
             num_layers=1,
             depthwise_conv_kernel_size=15,
             use_group_norm=True,)
                for _ in range(nlayers // 2)
            ]
        )

        self.conf_after = torch.nn.ModuleList(
            [Conformer(
             input_dim=d_hid,
             num_heads=num_heads,
             ffn_dim=d_hid * 2,
             num_layers=1,
             depthwise_conv_kernel_size=7,
             use_group_norm=True,)
                for _ in range(nlayers // 2)
            ]
        )

        self.out = nn.Linear(d_hid, 1 + vocab_size) # 1 for blank

        self.ctc_loss = nn.CTCLoss(blank=vocab_size, zero_infinity=True).cuda()

    def forward(self, latent, text=None, text_lens=None):
        layers = []

        x = self.mel_proj(latent.transpose(-1, -2)).transpose(-1, -2)

        x = x.transpose(1, 2)
        layers.append(nn.functional.avg_pool1d(x, 4))
        # x = x.transpose(1, 2)

        x = self.resblock1(x)
        x = nn.functional.avg_pool1d(x, 2)
        layers.append(nn.functional.avg_pool1d(x, 2))
        x = self.resblock2(x)
        x = nn.functional.avg_pool1d(x, 2)
        layers.append(x)

        x = x.transpose(1, 2)


forward

            ]
        )

        self.conf_after = torch.nn.ModuleList(
            [Conformer(
             input_dim=d_hid,
             num_heads=num_heads,
             ffn_dim=d_hid * 2,
             num_layers=1,
             depthwise_conv_kernel_size=7,
             use_group_norm=True,)
                for _ in range(nlayers // 2)
            ]
        )

        self.out = nn.Linear(d_hid, 1 + vocab_size) # 1 for blank

        self.ctc_loss = nn.CTCLoss(blank=vocab_size, zero_infinity=True).cuda()

    def forward(self, latent, text=None, text_lens=None):
        layers = []

        x = self.mel_proj(latent.transpose(-1, -2)).transpose(-1, -2)

        x = x.transpose(1, 2)
        layers.append(nn.functional.avg_pool1d(x, 4))
        # x = x.transpose(1, 2)

        x = self.resblock1(x)
        x = nn.functional.avg_pool1d(x, 2)
        layers.append(nn.functional.avg_pool1d(x, 2))
        x = self.resblock2(x)
        x = nn.functional.avg_pool1d(x, 2)
        layers.append(x)

        x = x.transpose(1, 2)

        batch_size, time_steps, _ = x.shape
        # Create a dummy lengths tensor (all sequences are assumed to be full length).
        input_lengths = torch.full((batch_size,), time_steps, device=x.device, dtype=torch.int64)

        for layer in (self.conf_pre):
            x, _ = layer(x, input_lengths)
            layers.append(x.transpose(1, 2))

        for layer in (self.conf_after):
            x, _ = layer(x, input_lengths)
            layers.append(x.transpose(1, 2))

        x = self.out(x)

        if text_lens is not None and text is not None:
            loss = self.ctc_loss(x.log_softmax(dim=2).transpose(0, 1), text, input_lengths, text_lens)
            return x, layers, loss
        else:
            return x, layers

if __name__ == "__main__":
    from f5_tts.model.utils import get_tokenizer

    bsz = 16

    tokenizer = "pinyin"  # 'pinyin', 'char', or 'custom'
    tokenizer_path = None  # if tokenizer = 'custom', define the path to the tokenizer you want to use (should be vocab.txt)
    dataset_name = "Emilia_ZH_EN"
    if tokenizer == "custom":
        tokenizer_path = tokenizer_path
    else:
        tokenizer_path = dataset_name
    vocab_char_map, vocab_size = get_tokenizer(tokenizer_path, tokenizer)

    model = ConformerCTC(vocab_size, mel_dim=80, num_heads=8, d_hid=512, nlayers=6).cuda()

    text = ["hello world"] * bsz


    lens_to_mask,
    mask_from_frac_lengths,
)

class ResBlock(nn.Module):
    def __init__(self, hidden_dim, n_conv=3, dropout_p=0.2):
        super().__init__()
        self._n_groups = 8
        self.blocks = nn.ModuleList([
            self._get_conv(hidden_dim, dilation=3**i, dropout_p=dropout_p)
            for i in range(n_conv)])

    def forward(self, x):
        for block in self.blocks:
            res = x
            x = block(x)
            x += res
        return x

    def _get_conv(self, hidden_dim, dilation, dropout_p=0.2):
        layers = [
            nn.Conv1d(hidden_dim, hidden_dim, kernel_size=3, padding=dilation, dilation=dilation),
            nn.ReLU(),
            nn.GroupNorm(num_groups=self._n_groups, num_channels=hidden_dim),
            nn.Dropout(p=dropout_p),
            nn.Conv1d(hidden_dim, hidden_dim, kernel_size=3, padding=1, dilation=1),
            nn.ReLU(),
            nn.Dropout(p=dropout_p)
        ]
        return nn.Sequential(*layers)

class ConformerCTC(nn.Module):
    def __init__(self,
                 vocab_size,
                 mel_dim=100,
                 num_heads=8,
                 d_hid=512,
                 nlayers=6):
        super().__init__()

        self.mel_proj = nn.Conv1d(mel_dim, d_hid, kernel_size=3, padding=1)

        self.d_hid = d_hid

        self.resblock1 = nn.Sequential(
                ResBlock(d_hid),
                nn.GroupNorm(num_groups=1, num_channels=d_hid)
            )


import random
import sys
from importlib.resources import files

import soundfile as sf
import tqdm
from cached_path import cached_path
from hydra.utils import get_class
from omegaconf import OmegaConf

from f5_tts.infer.utils_infer import (
    infer_process,
    load_model,
    load_vocoder,
    preprocess_ref_audio_text,
    remove_silence_for_generated_wav,
    save_spectrogram,
    transcribe,
)
from f5_tts.model.utils import seed_everything

class F5TTS:
    def __init__(
        self,
        model="F5TTS_v1_Base",
        ckpt_file="",
        vocab_file="",
        ode_method="euler",
        use_ema=True,
        vocoder_local_path=None,
        device=None,
        hf_cache_dir=None,
    ):
        model_cfg = OmegaConf.load(str(files("f5_tts").joinpath(f"configs/{model}.yaml")))
        model_cls = get_class(f"f5_tts.model.{model_cfg.model.backbone}")
        model_arc = model_cfg.model.arch

        self.mel_spec_type = model_cfg.model.mel_spec.mel_spec_type
        self.target_sample_rate = model_cfg.model.mel_spec.target_sample_rate

        self.ode_method = ode_method
        self.use_ema = use_ema

        if device is not None:
            self.device = device
        else:
            import torch

            self.device = (
                "cuda"
                if torch.cuda.is_available()
                else "xpu"
                if torch.xpu.is_available()
                else "mps"
                if torch.backends.mps.is_available()
                else "cpu"
            )

        # Load models
        self.vocoder = load_vocoder(
            self.mel_spec_type, vocoder_local_path is not None, vocoder_local_path, self.device, hf_cache_dir
        )

        repo_name, ckpt_step, ckpt_type = "F5-TTS", 1250000, "safetensors"

        # override for previous models
        if model == "F5TTS_Base":
            if self.mel_spec_type == "vocos":
                ckpt_step = 1200000
            elif self.mel_spec_type == "bigvgan":
                model = "F5TTS_Base_bigvgan"
                ckpt_type = "pt"
        elif model == "E2TTS_Base":
            repo_name = "E2-TTS"
            ckpt_step = 1200000

        if not ckpt_file:
            ckpt_file = str(
                cached_path(f"hf://SWivid/{repo_name}/{model}/model_{ckpt_step}.{ckpt_type}", cache_dir=hf_cache_dir)
            )
        self.ema_model = load_model(
            model_cls, model_arc, ckpt_file, self.mel_spec_type, vocab_file, self.ode_method, self.use_ema, self.device
        )

    def transcribe(self, ref_audio, language=None):
        return transcribe(ref_audio, language)

    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, spec, file_spec):
        save_spectrogram(spec, file_spec)

    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_spec=None,
        seed=None,
    ):
        if seed is None:
            seed = random.randint(0, sys.maxsize)
        seed_everything(seed)
        self.seed = seed

        ref_file, ref_text = preprocess_ref_audio_text(ref_file, ref_text)

        wav, sr, spec = 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_spec is not None:
            self.export_spectrogram(spec, file_spec)

        return wav, sr, spec

if __name__ == "__main__":
    f5tts = F5TTS()

    wav, sr, spec = 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_spec=str(files("f5_tts").joinpath("../../tests/api_out.png")),
        seed=None,
    )

    print("seed :", f5tts.seed)


from importlib.resources import files

import soundfile as sf
import tqdm
from cached_path import cached_path
from hydra.utils import get_class
from omegaconf import OmegaConf

from f5_tts.infer.utils_infer import (
    infer_process,
    load_model,
    load_vocoder,
    preprocess_ref_audio_text,
    remove_silence_for_generated_wav,
    save_spectrogram,
    transcribe,
)
from f5_tts.model.utils import seed_everything

class F5TTS:
    def __init__(
        self,
        model="F5TTS_v1_Base",
        ckpt_file="",
        vocab_file="",
        ode_method="euler",
        use_ema=True,
        vocoder_local_path=None,
        device=None,
        hf_cache_dir=None,
    ):
        model_cfg = OmegaConf.load(str(files("f5_tts").joinpath(f"configs/{model}.yaml")))
        model_cls = get_class(f"f5_tts.model.{model_cfg.model.backbone}")
        model_arc = model_cfg.model.arch

        self.mel_spec_type = model_cfg.model.mel_spec.mel_spec_type
        self.target_sample_rate = model_cfg.model.mel_spec.target_sample_rate

        self.ode_method = ode_method
        self.use_ema = use_ema

        if device is not None:
            self.device = device
        else:
            import torch

            self.device = (
                "cuda"
                if torch.cuda.is_available()
                else "xpu"
                if torch.xpu.is_available()
                else "mps"
                if torch.backends.mps.is_available()
                else "cpu"
            )

        # Load models
        self.vocoder = load_vocoder(
            self.mel_spec_type, vocoder_local_path is not None, vocoder_local_path, self.device, hf_cache_dir
        )

        repo_name, ckpt_step, ckpt_type = "F5-TTS", 1250000, "safetensors"

        # override for previous models
        if model == "F5TTS_Base":
            if self.mel_spec_type == "vocos":
                ckpt_step = 1200000
            elif self.mel_spec_type == "bigvgan":
                model = "F5TTS_Base_bigvgan"
                ckpt_type = "pt"
        elif model == "E2TTS_Base":
            repo_name = "E2-TTS"
            ckpt_step = 1200000

        if not ckpt_file:
            ckpt_file = str(
                cached_path(f"hf://SWivid/{repo_name}/{model}/model_{ckpt_step}.{ckpt_type}", cache_dir=hf_cache_dir)
            )
        self.ema_model = load_model(
            model_cls, model_arc, ckpt_file, self.mel_spec_type, vocab_file, self.ode_method, self.use_ema, self.device
        )

    def transcribe(self, ref_audio, language=None):
        return transcribe(ref_audio, language)

    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, spec, file_spec):
        save_spectrogram(spec, file_spec)

    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_spec=None,
        seed=None,
    ):
        if seed is None:
            seed = random.randint(0, sys.maxsize)
        seed_everything(seed)
        self.seed = seed

        ref_file, ref_text = preprocess_ref_audio_text(ref_file, ref_text)

        wav, sr, spec = 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_spec is not None:
            self.export_spectrogram(spec, file_spec)

        return wav, sr, spec

if __name__ == "__main__":
    f5tts = F5TTS()

    wav, sr, spec = 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_spec=str(files("f5_tts").joinpath("../../tests/api_out.png")),
        seed=None,
    )

    print("seed :", f5tts.seed)


import soundfile as sf
import tqdm
from cached_path import cached_path
from hydra.utils import get_class
from omegaconf import OmegaConf

from f5_tts.infer.utils_infer import (
    infer_process,
    load_model,
    load_vocoder,
    preprocess_ref_audio_text,
    remove_silence_for_generated_wav,
    save_spectrogram,
    transcribe,
)
from f5_tts.model.utils import seed_everything

class F5TTS:
    def __init__(
        self,
        model="F5TTS_v1_Base",
        ckpt_file="",
        vocab_file="",
        ode_method="euler",
        use_ema=True,
        vocoder_local_path=None,
        device=None,
        hf_cache_dir=None,
    ):
        model_cfg = OmegaConf.load(str(files("f5_tts").joinpath(f"configs/{model}.yaml")))
        model_cls = get_class(f"f5_tts.model.{model_cfg.model.backbone}")
        model_arc = model_cfg.model.arch

        self.mel_spec_type = model_cfg.model.mel_spec.mel_spec_type
        self.target_sample_rate = model_cfg.model.mel_spec.target_sample_rate

        self.ode_method = ode_method
        self.use_ema = use_ema

        if device is not None:
            self.device = device
        else:
            import torch

            self.device = (
                "cuda"
                if torch.cuda.is_available()
                else "xpu"
                if torch.xpu.is_available()
                else "mps"
                if torch.backends.mps.is_available()
                else "cpu"
            )

        # Load models
        self.vocoder = load_vocoder(
            self.mel_spec_type, vocoder_local_path is not None, vocoder_local_path, self.device, hf_cache_dir
        )

        repo_name, ckpt_step, ckpt_type = "F5-TTS", 1250000, "safetensors"

        # override for previous models
        if model == "F5TTS_Base":
            if self.mel_spec_type == "vocos":
                ckpt_step = 1200000
            elif self.mel_spec_type == "bigvgan":
                model = "F5TTS_Base_bigvgan"
                ckpt_type = "pt"
        elif model == "E2TTS_Base":
            repo_name = "E2-TTS"
            ckpt_step = 1200000

        if not ckpt_file:
            ckpt_file = str(
                cached_path(f"hf://SWivid/{repo_name}/{model}/model_{ckpt_step}.{ckpt_type}", cache_dir=hf_cache_dir)
            )
        self.ema_model = load_model(
            model_cls, model_arc, ckpt_file, self.mel_spec_type, vocab_file, self.ode_method, self.use_ema, self.device
        )

    def transcribe(self, ref_audio, language=None):
        return transcribe(ref_audio, language)

    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, spec, file_spec):
        save_spectrogram(spec, file_spec)

    def infer(
        self,
        ref_file,
        ref_text,
        gen_text,
        show_info=print,
        progress=tqdm,


transcribe

        # override for previous models
        if model == "F5TTS_Base":
            if self.mel_spec_type == "vocos":
                ckpt_step = 1200000
            elif self.mel_spec_type == "bigvgan":
                model = "F5TTS_Base_bigvgan"
                ckpt_type = "pt"
        elif model == "E2TTS_Base":
            repo_name = "E2-TTS"
            ckpt_step = 1200000

        if not ckpt_file:
            ckpt_file = str(
                cached_path(f"hf://SWivid/{repo_name}/{model}/model_{ckpt_step}.{ckpt_type}", cache_dir=hf_cache_dir)
            )
        self.ema_model = load_model(
            model_cls, model_arc, ckpt_file, self.mel_spec_type, vocab_file, self.ode_method, self.use_ema, self.device
        )

    def transcribe(self, ref_audio, language=None):
        return transcribe(ref_audio, language)

    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, spec, file_spec):
        save_spectrogram(spec, file_spec)

    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,


            if self.mel_spec_type == "vocos":
                ckpt_step = 1200000
            elif self.mel_spec_type == "bigvgan":
                model = "F5TTS_Base_bigvgan"
                ckpt_type = "pt"
        elif model == "E2TTS_Base":
            repo_name = "E2-TTS"
            ckpt_step = 1200000

        if not ckpt_file:
            ckpt_file = str(
                cached_path(f"hf://SWivid/{repo_name}/{model}/model_{ckpt_step}.{ckpt_type}", cache_dir=hf_cache_dir)
            )
        self.ema_model = load_model(
            model_cls, model_arc, ckpt_file, self.mel_spec_type, vocab_file, self.ode_method, self.use_ema, self.device
        )

    def transcribe(self, ref_audio, language=None):
        return transcribe(ref_audio, language)

    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, spec, file_spec):
        save_spectrogram(spec, file_spec)

    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,


            repo_name = "E2-TTS"
            ckpt_step = 1200000

        if not ckpt_file:
            ckpt_file = str(
                cached_path(f"hf://SWivid/{repo_name}/{model}/model_{ckpt_step}.{ckpt_type}", cache_dir=hf_cache_dir)
            )
        self.ema_model = load_model(
            model_cls, model_arc, ckpt_file, self.mel_spec_type, vocab_file, self.ode_method, self.use_ema, self.device
        )

    def transcribe(self, ref_audio, language=None):
        return transcribe(ref_audio, language)

    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, spec, file_spec):
        save_spectrogram(spec, file_spec)

    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_spec=None,
        seed=None,
    ):


infer

        if not ckpt_file:
            ckpt_file = str(
                cached_path(f"hf://SWivid/{repo_name}/{model}/model_{ckpt_step}.{ckpt_type}", cache_dir=hf_cache_dir)
            )
        self.ema_model = load_model(
            model_cls, model_arc, ckpt_file, self.mel_spec_type, vocab_file, self.ode_method, self.use_ema, self.device
        )

    def transcribe(self, ref_audio, language=None):
        return transcribe(ref_audio, language)

    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, spec, file_spec):
        save_spectrogram(spec, file_spec)

    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_spec=None,
        seed=None,
    ):
        if seed is None:
            seed = random.randint(0, sys.maxsize)
        seed_everything(seed)
        self.seed = seed

        ref_file, ref_text = preprocess_ref_audio_text(ref_file, ref_text)

        wav, sr, spec = 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_spec is not None:
            self.export_spectrogram(spec, file_spec)

        return wav, sr, spec

if __name__ == "__main__":
    f5tts = F5TTS()

    wav, sr, spec = 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_spec=str(files("f5_tts").joinpath("../../tests/api_out.png")),
        seed=None,
    )

    print("seed :", f5tts.seed)


import asyncio
import logging
import socket
import time

import numpy as np
import pyaudio

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

async def listen_to_F5TTS(text, server_ip="localhost", server_port=9998):
    client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    await asyncio.get_event_loop().run_in_executor(None, client_socket.connect, (server_ip, int(server_port)))

    start_time = time.time()
    first_chunk_time = None

    async def play_audio_stream():
        nonlocal first_chunk_time
        p = pyaudio.PyAudio()
        stream = p.open(format=pyaudio.paFloat32, channels=1, rate=24000, output=True, frames_per_buffer=2048)

        try:
            while True:
                data = await asyncio.get_event_loop().run_in_executor(None, client_socket.recv, 8192)
                if not data:
                    break
                if data == b"END":
                    logger.info("End of audio received.")
                    break

                audio_array = np.frombuffer(data, dtype=np.float32)
                stream.write(audio_array.tobytes())

                if first_chunk_time is None:
                    first_chunk_time = time.time()

        finally:
            stream.stop_stream()
            stream.close()
            p.terminate()

        logger.info(f"Total time taken: {time.time() - start_time:.4f} seconds")

    try:
        data_to_send = f"{text}".encode("utf-8")
        await asyncio.get_event_loop().run_in_executor(None, client_socket.sendall, data_to_send)
        await play_audio_stream()

    except Exception as e:
        logger.error(f"Error in listen_to_F5TTS: {e}")

    finally:
        client_socket.close()

if __name__ == "__main__":
    text_to_send = "As a Reader assistant, I'm familiar with new technology. which are key to its improved performance in terms of both training speed and inference efficiency. Let's break down the components"

    asyncio.run(listen_to_F5TTS(text_to_send))


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 f5_tts.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("f5_tts").joinpath(f"configs/{model}.yaml")))
        self.model_cls = get_class(f"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("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()


AudioFileWriterThread

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 f5_tts.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("f5_tts").joinpath(f"configs/{model}.yaml")))
        self.model_cls = get_class(f"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)


TTSStreamingProcessor

                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("f5_tts").joinpath(f"configs/{model}.yaml")))
        self.model_cls = get_class(f"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


        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(


            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",


                    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("f5_tts").joinpath(f"configs/{model}.yaml")))
        self.model_cls = get_class(f"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)


run

    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"


        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("f5_tts").joinpath(f"configs/{model}.yaml")))


stop

        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("f5_tts").joinpath(f"configs/{model}.yaml")))
        self.model_cls = get_class(f"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)


            else "xpu"
            if torch.xpu.is_available()
            else "mps"
            if torch.backends.mps.is_available()
            else "cpu"
        )
        model_cfg = OmegaConf.load(str(files("f5_tts").joinpath(f"configs/{model}.yaml")))
        self.model_cls = get_class(f"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,


        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,


        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: