Delete src/f5-tts

src/f5-tts/api.py

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

src/f5-tts/model/__init__.py

@@ -1,10 +0,0 @@
-from f5_tts.model.cfm import CFM
-
-from f5_tts.model.backbones.unett import UNetT
-from f5_tts.model.backbones.dit import DiT
-from f5_tts.model.backbones.mmdit import MMDiT
-
-from f5_tts.model.trainer import Trainer
-
-
-__all__ = ["CFM", "UNetT", "DiT", "MMDiT", "Trainer"]

src/f5-tts/model/backbones/README.md

@@ -1,20 +0,0 @@
-## Backbones quick introduction
-
-
-### unett.py
-- flat unet transformer
-- structure same as in e2-tts & voicebox paper except using rotary pos emb
-- update: allow possible abs pos emb & convnextv2 blocks for embedded text before concat
-
-### dit.py
-- adaln-zero dit
-- embedded timestep as condition
-- concatted noised_input + masked_cond + embedded_text, linear proj in
-- possible abs pos emb & convnextv2 blocks for embedded text before concat
-- possible long skip connection (first layer to last layer)
-
-### mmdit.py
-- sd3 structure
-- timestep as condition
-- left stream: text embedded and applied a abs pos emb
-- right stream: masked_cond & noised_input concatted and with same conv pos emb as unett

src/f5-tts/model/backbones/dit.py

@@ -1,163 +0,0 @@
-"""
-ein notation:
-b - batch
-n - sequence
-nt - text sequence
-nw - raw wave length
-d - dimension
-"""
-
-from __future__ import annotations
-
-import torch
-from torch import nn
-import torch.nn.functional as F
-
-from x_transformers.x_transformers import RotaryEmbedding
-
-from f5_tts.model.modules import (
-    TimestepEmbedding,
-    ConvNeXtV2Block,
-    ConvPositionEmbedding,
-    DiTBlock,
-    AdaLayerNormZero_Final,
-    precompute_freqs_cis,
-    get_pos_embed_indices,
-)
-
-
-# Text embedding
-
-
-class TextEmbedding(nn.Module):
-    def __init__(self, text_num_embeds, text_dim, conv_layers=0, conv_mult=2):
-        super().__init__()
-        self.text_embed = nn.Embedding(text_num_embeds + 1, text_dim)  # use 0 as filler token
-
-        if conv_layers > 0:
-            self.extra_modeling = True
-            self.precompute_max_pos = 4096  # ~44s of 24khz audio
-            self.register_buffer("freqs_cis", precompute_freqs_cis(text_dim, self.precompute_max_pos), persistent=False)
-            self.text_blocks = nn.Sequential(
-                *[ConvNeXtV2Block(text_dim, text_dim * conv_mult) for _ in range(conv_layers)]
-            )
-        else:
-            self.extra_modeling = False
-
-    def forward(self, text: int["b nt"], seq_len, drop_text=False):  # noqa: F722
-        text = text + 1  # use 0 as filler token. preprocess of batch pad -1, see list_str_to_idx()
-        text = text[:, :seq_len]  # curtail if character tokens are more than the mel spec tokens
-        batch, text_len = text.shape[0], text.shape[1]
-        text = F.pad(text, (0, seq_len - text_len), value=0)
-
-        if drop_text:  # cfg for text
-            text = torch.zeros_like(text)
-
-        text = self.text_embed(text)  # b n -> b n d
-
-        # possible extra modeling
-        if self.extra_modeling:
-            # sinus pos emb
-            batch_start = torch.zeros((batch,), dtype=torch.long)
-            pos_idx = get_pos_embed_indices(batch_start, seq_len, max_pos=self.precompute_max_pos)
-            text_pos_embed = self.freqs_cis[pos_idx]
-            text = text + text_pos_embed
-
-            # convnextv2 blocks
-            text = self.text_blocks(text)
-
-        return text
-
-
-# noised input audio and context mixing embedding
-
-
-class InputEmbedding(nn.Module):
-    def __init__(self, mel_dim, text_dim, out_dim):
-        super().__init__()
-        self.proj = nn.Linear(mel_dim * 2 + text_dim, out_dim)
-        self.conv_pos_embed = ConvPositionEmbedding(dim=out_dim)
-
-    def forward(self, x: float["b n d"], cond: float["b n d"], text_embed: float["b n d"], drop_audio_cond=False):  # noqa: F722
-        if drop_audio_cond:  # cfg for cond audio
-            cond = torch.zeros_like(cond)
-
-        x = self.proj(torch.cat((x, cond, text_embed), dim=-1))
-        x = self.conv_pos_embed(x) + x
-        return x
-
-
-# Transformer backbone using DiT blocks
-
-
-class DiT(nn.Module):
-    def __init__(
-        self,
-        *,
-        dim,
-        depth=8,
-        heads=8,
-        dim_head=64,
-        dropout=0.1,
-        ff_mult=4,
-        mel_dim=100,
-        text_num_embeds=256,
-        text_dim=None,
-        conv_layers=0,
-        long_skip_connection=False,
-    ):
-        super().__init__()
-
-        self.time_embed = TimestepEmbedding(dim)
-        if text_dim is None:
-            text_dim = mel_dim
-        self.text_embed = TextEmbedding(text_num_embeds, text_dim, conv_layers=conv_layers)
-        self.input_embed = InputEmbedding(mel_dim, text_dim, dim)
-
-        self.rotary_embed = RotaryEmbedding(dim_head)
-
-        self.dim = dim
-        self.depth = depth
-
-        self.transformer_blocks = nn.ModuleList(
-            [DiTBlock(dim=dim, heads=heads, dim_head=dim_head, ff_mult=ff_mult, dropout=dropout) for _ in range(depth)]
-        )
-        self.long_skip_connection = nn.Linear(dim * 2, dim, bias=False) if long_skip_connection else None
-
-        self.norm_out = AdaLayerNormZero_Final(dim)  # final modulation
-        self.proj_out = nn.Linear(dim, mel_dim)
-
-    def forward(
-        self,
-        x: float["b n d"],  # nosied input audio  # noqa: F722
-        cond: float["b n d"],  # masked cond audio  # noqa: F722
-        text: int["b nt"],  # text  # noqa: F722
-        time: float["b"] | float[""],  # time step  # noqa: F821 F722
-        drop_audio_cond,  # cfg for cond audio
-        drop_text,  # cfg for text
-        mask: bool["b n"] | None = None,  # noqa: F722
-    ):
-        batch, seq_len = x.shape[0], x.shape[1]
-        if time.ndim == 0:
-            time = time.repeat(batch)
-
-        # t: conditioning time, c: context (text + masked cond audio), x: noised input audio
-        t = self.time_embed(time)
-        text_embed = self.text_embed(text, seq_len, drop_text=drop_text)
-        x = self.input_embed(x, cond, text_embed, drop_audio_cond=drop_audio_cond)
-
-        rope = self.rotary_embed.forward_from_seq_len(seq_len)
-
-        if self.long_skip_connection is not None:
-            residual = x
-
-        for block in self.transformer_blocks:
-            x = block(x, t, mask=mask, rope=rope)
-
-        if self.long_skip_connection is not None:
-            x = self.long_skip_connection(torch.cat((x, residual), dim=-1))
-
-        x = self.norm_out(x, t)
-        output = self.proj_out(x)
-
-        return output

src/f5-tts/model/backbones/mmdit.py

@@ -1,146 +0,0 @@
-"""
-ein notation:
-b - batch
-n - sequence
-nt - text sequence
-nw - raw wave length
-d - dimension
-"""
-
-from __future__ import annotations
-
-import torch
-from torch import nn
-
-from x_transformers.x_transformers import RotaryEmbedding
-
-from f5_tts.model.modules import (
-    TimestepEmbedding,
-    ConvPositionEmbedding,
-    MMDiTBlock,
-    AdaLayerNormZero_Final,
-    precompute_freqs_cis,
-    get_pos_embed_indices,
-)
-
-
-# text embedding
-
-
-class TextEmbedding(nn.Module):
-    def __init__(self, out_dim, text_num_embeds):
-        super().__init__()
-        self.text_embed = nn.Embedding(text_num_embeds + 1, out_dim)  # will use 0 as filler token
-
-        self.precompute_max_pos = 1024
-        self.register_buffer("freqs_cis", precompute_freqs_cis(out_dim, self.precompute_max_pos), persistent=False)
-
-    def forward(self, text: int["b nt"], drop_text=False) -> int["b nt d"]:  # noqa: F722
-        text = text + 1
-        if drop_text:
-            text = torch.zeros_like(text)
-        text = self.text_embed(text)
-
-        # sinus pos emb
-        batch_start = torch.zeros((text.shape[0],), dtype=torch.long)
-        batch_text_len = text.shape[1]
-        pos_idx = get_pos_embed_indices(batch_start, batch_text_len, max_pos=self.precompute_max_pos)
-        text_pos_embed = self.freqs_cis[pos_idx]
-
-        text = text + text_pos_embed
-
-        return text
-
-
-# noised input & masked cond audio embedding
-
-
-class AudioEmbedding(nn.Module):
-    def __init__(self, in_dim, out_dim):
-        super().__init__()
-        self.linear = nn.Linear(2 * in_dim, out_dim)
-        self.conv_pos_embed = ConvPositionEmbedding(out_dim)
-
-    def forward(self, x: float["b n d"], cond: float["b n d"], drop_audio_cond=False):  # noqa: F722
-        if drop_audio_cond:
-            cond = torch.zeros_like(cond)
-        x = torch.cat((x, cond), dim=-1)
-        x = self.linear(x)
-        x = self.conv_pos_embed(x) + x
-        return x
-
-
-# Transformer backbone using MM-DiT blocks
-
-
-class MMDiT(nn.Module):
-    def __init__(
-        self,
-        *,
-        dim,
-        depth=8,
-        heads=8,
-        dim_head=64,
-        dropout=0.1,
-        ff_mult=4,
-        text_num_embeds=256,
-        mel_dim=100,
-    ):
-        super().__init__()
-
-        self.time_embed = TimestepEmbedding(dim)
-        self.text_embed = TextEmbedding(dim, text_num_embeds)
-        self.audio_embed = AudioEmbedding(mel_dim, dim)
-
-        self.rotary_embed = RotaryEmbedding(dim_head)
-
-        self.dim = dim
-        self.depth = depth
-
-        self.transformer_blocks = nn.ModuleList(
-            [
-                MMDiTBlock(
-                    dim=dim,
-                    heads=heads,
-                    dim_head=dim_head,
-                    dropout=dropout,
-                    ff_mult=ff_mult,
-                    context_pre_only=i == depth - 1,
-                )
-                for i in range(depth)
-            ]
-        )
-        self.norm_out = AdaLayerNormZero_Final(dim)  # final modulation
-        self.proj_out = nn.Linear(dim, mel_dim)
-
-    def forward(
-        self,
-        x: float["b n d"],  # nosied input audio  # noqa: F722
-        cond: float["b n d"],  # masked cond audio  # noqa: F722
-        text: int["b nt"],  # text  # noqa: F722
-        time: float["b"] | float[""],  # time step  # noqa: F821 F722
-        drop_audio_cond,  # cfg for cond audio
-        drop_text,  # cfg for text
-        mask: bool["b n"] | None = None,  # noqa: F722
-    ):
-        batch = x.shape[0]
-        if time.ndim == 0:
-            time = time.repeat(batch)
-
-        # t: conditioning (time), c: context (text + masked cond audio), x: noised input audio
-        t = self.time_embed(time)
-        c = self.text_embed(text, drop_text=drop_text)
-        x = self.audio_embed(x, cond, drop_audio_cond=drop_audio_cond)
-
-        seq_len = x.shape[1]
-        text_len = text.shape[1]
-        rope_audio = self.rotary_embed.forward_from_seq_len(seq_len)
-        rope_text = self.rotary_embed.forward_from_seq_len(text_len)
-
-        for block in self.transformer_blocks:
-            c, x = block(x, c, t, mask=mask, rope=rope_audio, c_rope=rope_text)
-
-        x = self.norm_out(x, t)
-        output = self.proj_out(x)
-
-        return output

src/f5-tts/model/backbones/unett.py

@@ -1,219 +0,0 @@
-"""
-ein notation:
-b - batch
-n - sequence
-nt - text sequence
-nw - raw wave length
-d - dimension
-"""
-
-from __future__ import annotations
-from typing import Literal
-
-import torch
-from torch import nn
-import torch.nn.functional as F
-
-from x_transformers import RMSNorm
-from x_transformers.x_transformers import RotaryEmbedding
-
-from f5_tts.model.modules import (
-    TimestepEmbedding,
-    ConvNeXtV2Block,
-    ConvPositionEmbedding,
-    Attention,
-    AttnProcessor,
-    FeedForward,
-    precompute_freqs_cis,
-    get_pos_embed_indices,
-)
-
-
-# Text embedding
-
-
-class TextEmbedding(nn.Module):
-    def __init__(self, text_num_embeds, text_dim, conv_layers=0, conv_mult=2):
-        super().__init__()
-        self.text_embed = nn.Embedding(text_num_embeds + 1, text_dim)  # use 0 as filler token
-
-        if conv_layers > 0:
-            self.extra_modeling = True
-            self.precompute_max_pos = 4096  # ~44s of 24khz audio
-            self.register_buffer("freqs_cis", precompute_freqs_cis(text_dim, self.precompute_max_pos), persistent=False)
-            self.text_blocks = nn.Sequential(
-                *[ConvNeXtV2Block(text_dim, text_dim * conv_mult) for _ in range(conv_layers)]
-            )
-        else:
-            self.extra_modeling = False
-
-    def forward(self, text: int["b nt"], seq_len, drop_text=False):  # noqa: F722
-        text = text + 1  # use 0 as filler token. preprocess of batch pad -1, see list_str_to_idx()
-        text = text[:, :seq_len]  # curtail if character tokens are more than the mel spec tokens
-        batch, text_len = text.shape[0], text.shape[1]
-        text = F.pad(text, (0, seq_len - text_len), value=0)
-
-        if drop_text:  # cfg for text
-            text = torch.zeros_like(text)
-
-        text = self.text_embed(text)  # b n -> b n d
-
-        # possible extra modeling
-        if self.extra_modeling:
-            # sinus pos emb
-            batch_start = torch.zeros((batch,), dtype=torch.long)
-            pos_idx = get_pos_embed_indices(batch_start, seq_len, max_pos=self.precompute_max_pos)
-            text_pos_embed = self.freqs_cis[pos_idx]
-            text = text + text_pos_embed
-
-            # convnextv2 blocks
-            text = self.text_blocks(text)
-
-        return text
-
-
-# noised input audio and context mixing embedding
-
-
-class InputEmbedding(nn.Module):
-    def __init__(self, mel_dim, text_dim, out_dim):
-        super().__init__()
-        self.proj = nn.Linear(mel_dim * 2 + text_dim, out_dim)
-        self.conv_pos_embed = ConvPositionEmbedding(dim=out_dim)
-
-    def forward(self, x: float["b n d"], cond: float["b n d"], text_embed: float["b n d"], drop_audio_cond=False):  # noqa: F722
-        if drop_audio_cond:  # cfg for cond audio
-            cond = torch.zeros_like(cond)
-
-        x = self.proj(torch.cat((x, cond, text_embed), dim=-1))
-        x = self.conv_pos_embed(x) + x
-        return x
-
-
-# Flat UNet Transformer backbone
-
-
-class UNetT(nn.Module):
-    def __init__(
-        self,
-        *,
-        dim,
-        depth=8,
-        heads=8,
-        dim_head=64,
-        dropout=0.1,
-        ff_mult=4,
-        mel_dim=100,
-        text_num_embeds=256,
-        text_dim=None,
-        conv_layers=0,
-        skip_connect_type: Literal["add", "concat", "none"] = "concat",
-    ):
-        super().__init__()
-        assert depth % 2 == 0, "UNet-Transformer's depth should be even."
-
-        self.time_embed = TimestepEmbedding(dim)
-        if text_dim is None:
-            text_dim = mel_dim
-        self.text_embed = TextEmbedding(text_num_embeds, text_dim, conv_layers=conv_layers)
-        self.input_embed = InputEmbedding(mel_dim, text_dim, dim)
-
-        self.rotary_embed = RotaryEmbedding(dim_head)
-
-        # transformer layers & skip connections
-
-        self.dim = dim
-        self.skip_connect_type = skip_connect_type
-        needs_skip_proj = skip_connect_type == "concat"
-
-        self.depth = depth
-        self.layers = nn.ModuleList([])
-
-        for idx in range(depth):
-            is_later_half = idx >= (depth // 2)
-
-            attn_norm = RMSNorm(dim)
-            attn = Attention(
-                processor=AttnProcessor(),
-                dim=dim,
-                heads=heads,
-                dim_head=dim_head,
-                dropout=dropout,
-            )
-
-            ff_norm = RMSNorm(dim)
-            ff = FeedForward(dim=dim, mult=ff_mult, dropout=dropout, approximate="tanh")
-
-            skip_proj = nn.Linear(dim * 2, dim, bias=False) if needs_skip_proj and is_later_half else None
-
-            self.layers.append(
-                nn.ModuleList(
-                    [
-                        skip_proj,
-                        attn_norm,
-                        attn,
-                        ff_norm,
-                        ff,
-                    ]
-                )
-            )
-
-        self.norm_out = RMSNorm(dim)
-        self.proj_out = nn.Linear(dim, mel_dim)
-
-    def forward(
-        self,
-        x: float["b n d"],  # nosied input audio  # noqa: F722
-        cond: float["b n d"],  # masked cond audio  # noqa: F722
-        text: int["b nt"],  # text  # noqa: F722
-        time: float["b"] | float[""],  # time step  # noqa: F821 F722
-        drop_audio_cond,  # cfg for cond audio
-        drop_text,  # cfg for text
-        mask: bool["b n"] | None = None,  # noqa: F722
-    ):
-        batch, seq_len = x.shape[0], x.shape[1]
-        if time.ndim == 0:
-            time = time.repeat(batch)
-
-        # t: conditioning time, c: context (text + masked cond audio), x: noised input audio
-        t = self.time_embed(time)
-        text_embed = self.text_embed(text, seq_len, drop_text=drop_text)
-        x = self.input_embed(x, cond, text_embed, drop_audio_cond=drop_audio_cond)
-
-        # postfix time t to input x, [b n d] -> [b n+1 d]
-        x = torch.cat([t.unsqueeze(1), x], dim=1)  # pack t to x
-        if mask is not None:
-            mask = F.pad(mask, (1, 0), value=1)
-
-        rope = self.rotary_embed.forward_from_seq_len(seq_len + 1)
-
-        # flat unet transformer
-        skip_connect_type = self.skip_connect_type
-        skips = []
-        for idx, (maybe_skip_proj, attn_norm, attn, ff_norm, ff) in enumerate(self.layers):
-            layer = idx + 1
-
-            # skip connection logic
-            is_first_half = layer <= (self.depth // 2)
-            is_later_half = not is_first_half
-
-            if is_first_half:
-                skips.append(x)
-
-            if is_later_half:
-                skip = skips.pop()
-                if skip_connect_type == "concat":
-                    x = torch.cat((x, skip), dim=-1)
-                    x = maybe_skip_proj(x)
-                elif skip_connect_type == "add":
-                    x = x + skip
-
-            # attention and feedforward blocks
-            x = attn(attn_norm(x), rope=rope, mask=mask) + x
-            x = ff(ff_norm(x)) + x
-
-        assert len(skips) == 0
-
-        x = self.norm_out(x)[:, 1:, :]  # unpack t from x
-
-        return self.proj_out(x)

src/f5-tts/model/cfm.py

@@ -1,285 +0,0 @@
-"""
-ein notation:
-b - batch
-n - sequence
-nt - text sequence
-nw - raw wave length
-d - dimension
-"""
-
-from __future__ import annotations
-
-from random import random
-from typing import Callable
-
-import torch
-import torch.nn.functional as F
-from torch import nn
-from torch.nn.utils.rnn import pad_sequence
-from torchdiffeq import odeint
-
-from f5_tts.model.modules import MelSpec
-from f5_tts.model.utils import (
-    default,
-    exists,
-    lens_to_mask,
-    list_str_to_idx,
-    list_str_to_tensor,
-    mask_from_frac_lengths,
-)
-
-
-class CFM(nn.Module):
-    def __init__(
-        self,
-        transformer: nn.Module,
-        sigma=0.0,
-        odeint_kwargs: dict = dict(
-            # atol = 1e-5,
-            # rtol = 1e-5,
-            method="euler"  # 'midpoint'
-        ),
-        audio_drop_prob=0.3,
-        cond_drop_prob=0.2,
-        num_channels=None,
-        mel_spec_module: nn.Module | None = None,
-        mel_spec_kwargs: dict = dict(),
-        frac_lengths_mask: tuple[float, float] = (0.7, 1.0),
-        vocab_char_map: dict[str:int] | None = None,
-    ):
-        super().__init__()
-
-        self.frac_lengths_mask = frac_lengths_mask
-
-        # mel spec
-        self.mel_spec = default(mel_spec_module, MelSpec(**mel_spec_kwargs))
-        num_channels = default(num_channels, self.mel_spec.n_mel_channels)
-        self.num_channels = num_channels
-
-        # classifier-free guidance
-        self.audio_drop_prob = audio_drop_prob
-        self.cond_drop_prob = cond_drop_prob
-
-        # transformer
-        self.transformer = transformer
-        dim = transformer.dim
-        self.dim = dim
-
-        # conditional flow related
-        self.sigma = sigma
-
-        # sampling related
-        self.odeint_kwargs = odeint_kwargs
-
-        # vocab map for tokenization
-        self.vocab_char_map = vocab_char_map
-
-    @property
-    def device(self):
-        return next(self.parameters()).device
-
-    @torch.no_grad()
-    def sample(
-        self,
-        cond: float["b n d"] | float["b nw"],  # noqa: F722
-        text: int["b nt"] | list[str],  # noqa: F722
-        duration: int | int["b"],  # noqa: F821
-        *,
-        lens: int["b"] | None = None,  # noqa: F821
-        steps=32,
-        cfg_strength=1.0,
-        sway_sampling_coef=None,
-        seed: int | None = None,
-        max_duration=4096,
-        vocoder: Callable[[float["b d n"]], float["b nw"]] | None = None,  # noqa: F722
-        no_ref_audio=False,
-        duplicate_test=False,
-        t_inter=0.1,
-        edit_mask=None,
-    ):
-        self.eval()
-        # raw wave
-
-        if cond.ndim == 2:
-            cond = self.mel_spec(cond)
-            cond = cond.permute(0, 2, 1)
-            assert cond.shape[-1] == self.num_channels
-
-        cond = cond.to(next(self.parameters()).dtype)
-
-        batch, cond_seq_len, device = *cond.shape[:2], cond.device
-        if not exists(lens):
-            lens = torch.full((batch,), cond_seq_len, device=device, dtype=torch.long)
-
-        # text
-
-        if isinstance(text, list):
-            if exists(self.vocab_char_map):
-                text = list_str_to_idx(text, self.vocab_char_map).to(device)
-            else:
-                text = list_str_to_tensor(text).to(device)
-            assert text.shape[0] == batch
-
-        if exists(text):
-            text_lens = (text != -1).sum(dim=-1)
-            lens = torch.maximum(text_lens, lens)  # make sure lengths are at least those of the text characters
-
-        # duration
-
-        cond_mask = lens_to_mask(lens)
-        if edit_mask is not None:
-            cond_mask = cond_mask & edit_mask
-
-        if isinstance(duration, int):
-            duration = torch.full((batch,), duration, device=device, dtype=torch.long)
-
-        duration = torch.maximum(lens + 1, duration)  # just add one token so something is generated
-        duration = duration.clamp(max=max_duration)
-        max_duration = duration.amax()
-
-        # duplicate test corner for inner time step oberservation
-        if duplicate_test:
-            test_cond = F.pad(cond, (0, 0, cond_seq_len, max_duration - 2 * cond_seq_len), value=0.0)
-
-        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)
-        )  # allow direct control (cut cond audio) with lens passed in
-
-        if batch > 1:
-            mask = lens_to_mask(duration)
-        else:  # save memory and speed up, as single inference need no mask currently
-            mask = None
-
-        # test for no ref audio
-        if no_ref_audio:
-            cond = torch.zeros_like(cond)
-
-        # neural ode
-
-        def fn(t, x):
-            # at each step, conditioning is fixed
-            # step_cond = torch.where(cond_mask, cond, torch.zeros_like(cond))
-
-            # predict flow
-            pred = self.transformer(
-                x=x, cond=step_cond, text=text, time=t, mask=mask, drop_audio_cond=False, drop_text=False
-            )
-            if cfg_strength < 1e-5:
-                return pred
-
-            null_pred = self.transformer(
-                x=x, cond=step_cond, text=text, time=t, mask=mask, drop_audio_cond=True, drop_text=True
-            )
-            return pred + (pred - null_pred) * cfg_strength
-
-        # noise input
-        # to make sure batch inference result is same with different batch size, and for sure single inference
-        # still some difference maybe due to convolutional layers
-        y0 = []
-        for dur in duration:
-            if exists(seed):
-                torch.manual_seed(seed)
-            y0.append(torch.randn(dur, self.num_channels, device=self.device, dtype=step_cond.dtype))
-        y0 = pad_sequence(y0, padding_value=0, batch_first=True)
-
-        t_start = 0
-
-        # duplicate test corner for inner time step oberservation
-        if duplicate_test:
-            t_start = t_inter
-            y0 = (1 - t_start) * y0 + t_start * test_cond
-            steps = int(steps * (1 - t_start))
-
-        t = torch.linspace(t_start, 1, steps, device=self.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)
-
-        trajectory = odeint(fn, y0, t, **self.odeint_kwargs)
-
-        sampled = trajectory[-1]
-        out = sampled
-        out = torch.where(cond_mask, cond, out)
-
-        if exists(vocoder):
-            out = out.permute(0, 2, 1)
-            out = vocoder(out)
-
-        return out, trajectory
-
-    def forward(
-        self,
-        inp: float["b n d"] | float["b nw"],  # mel or raw wave  # noqa: F722
-        text: int["b nt"] | list[str],  # noqa: F722
-        *,
-        lens: int["b"] | None = None,  # noqa: F821
-        noise_scheduler: str | None = None,
-    ):
-        # handle raw wave
-        if inp.ndim == 2:
-            inp = self.mel_spec(inp)
-            inp = inp.permute(0, 2, 1)
-            assert inp.shape[-1] == self.num_channels
-
-        batch, seq_len, dtype, device, _σ1 = *inp.shape[:2], inp.dtype, self.device, self.sigma
-
-        # handle text as string
-        if isinstance(text, list):
-            if exists(self.vocab_char_map):
-                text = list_str_to_idx(text, self.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
-
-        # get a random span to mask out for training conditionally
-        frac_lengths = torch.zeros((batch,), device=self.device).float().uniform_(*self.frac_lengths_mask)
-        rand_span_mask = mask_from_frac_lengths(lens, frac_lengths)
-
-        if exists(mask):
-            rand_span_mask &= mask
-
-        # mel is x1
-        x1 = inp
-
-        # x0 is gaussian noise
-        x0 = torch.randn_like(x1)
-
-        # time step
-        time = torch.rand((batch,), dtype=dtype, device=self.device)
-        # TODO. noise_scheduler
-
-        # sample xt (φ_t(x) in the paper)
-        t = time.unsqueeze(-1).unsqueeze(-1)
-        φ = (1 - t) * x0 + t * x1
-        flow = x1 - x0
-
-        # only predict what is within the random mask span for infilling
-        cond = torch.where(rand_span_mask[..., None], torch.zeros_like(x1), x1)
-
-        # transformer and cfg training with a drop rate
-        drop_audio_cond = random() < self.audio_drop_prob  # p_drop in voicebox paper
-        if random() < self.cond_drop_prob:  # p_uncond in voicebox paper
-            drop_audio_cond = True
-            drop_text = True
-        else:
-            drop_text = False
-
-        # if want rigourously mask out padding, record in collate_fn in dataset.py, and pass in here
-        # adding mask will use more memory, thus also need to adjust batchsampler with scaled down threshold for long sequences
-        pred = self.transformer(
-            x=φ, cond=cond, text=text, time=time, drop_audio_cond=drop_audio_cond, drop_text=drop_text
-        )
-
-        # flow matching loss
-        loss = F.mse_loss(pred, flow, reduction="none")
-        loss = loss[rand_span_mask]
-
-        return loss.mean(), cond, pred

src/f5-tts/model/dataset.py

@@ -1,314 +0,0 @@
-import json
-import random
-from importlib.resources import files
-
-import torch
-import torch.nn.functional as F
-import torchaudio
-from datasets import Dataset as Dataset_
-from datasets import load_from_disk
-from torch import nn
-from torch.utils.data import Dataset, Sampler
-from tqdm import tqdm
-
-from f5_tts.model.modules import MelSpec
-from f5_tts.model.utils import default
-
-
-class HFDataset(Dataset):
-    def __init__(
-        self,
-        hf_dataset: Dataset,
-        target_sample_rate=24_000,
-        n_mel_channels=100,
-        hop_length=256,
-        n_fft=1024,
-        win_length=1024,
-        mel_spec_type="vocos",
-    ):
-        self.data = hf_dataset
-        self.target_sample_rate = target_sample_rate
-        self.hop_length = hop_length
-
-        self.mel_spectrogram = MelSpec(
-            n_fft=n_fft,
-            hop_length=hop_length,
-            win_length=win_length,
-            n_mel_channels=n_mel_channels,
-            target_sample_rate=target_sample_rate,
-            mel_spec_type=mel_spec_type,
-        )
-
-    def get_frame_len(self, index):
-        row = self.data[index]
-        audio = row["audio"]["array"]
-        sample_rate = row["audio"]["sampling_rate"]
-        return audio.shape[-1] / sample_rate * self.target_sample_rate / self.hop_length
-
-    def __len__(self):
-        return len(self.data)
-
-    def __getitem__(self, index):
-        row = self.data[index]
-        audio = row["audio"]["array"]
-
-        # logger.info(f"Audio shape: {audio.shape}")
-
-        sample_rate = row["audio"]["sampling_rate"]
-        duration = audio.shape[-1] / sample_rate
-
-        if duration > 30 or duration < 0.3:
-            return self.__getitem__((index + 1) % len(self.data))
-
-        audio_tensor = torch.from_numpy(audio).float()
-
-        if sample_rate != self.target_sample_rate:
-            resampler = torchaudio.transforms.Resample(sample_rate, self.target_sample_rate)
-            audio_tensor = resampler(audio_tensor)
-
-        audio_tensor = audio_tensor.unsqueeze(0)  # 't -> 1 t')
-
-        mel_spec = self.mel_spectrogram(audio_tensor)
-
-        mel_spec = mel_spec.squeeze(0)  # '1 d t -> d t'
-
-        text = row["text"]
-
-        return dict(
-            mel_spec=mel_spec,
-            text=text,
-        )
-
-
-class CustomDataset(Dataset):
-    def __init__(
-        self,
-        custom_dataset: Dataset,
-        durations=None,
-        target_sample_rate=24_000,
-        hop_length=256,
-        n_mel_channels=100,
-        n_fft=1024,
-        win_length=1024,
-        mel_spec_type="vocos",
-        preprocessed_mel=False,
-        mel_spec_module: nn.Module | None = None,
-    ):
-        self.data = custom_dataset
-        self.durations = durations
-        self.target_sample_rate = target_sample_rate
-        self.hop_length = hop_length
-        self.n_fft = n_fft
-        self.win_length = win_length
-        self.mel_spec_type = mel_spec_type
-        self.preprocessed_mel = preprocessed_mel
-
-        if not preprocessed_mel:
-            self.mel_spectrogram = default(
-                mel_spec_module,
-                MelSpec(
-                    n_fft=n_fft,
-                    hop_length=hop_length,
-                    win_length=win_length,
-                    n_mel_channels=n_mel_channels,
-                    target_sample_rate=target_sample_rate,
-                    mel_spec_type=mel_spec_type,
-                ),
-            )
-
-    def get_frame_len(self, index):
-        if (
-            self.durations is not None
-        ):  # Please make sure the separately provided durations are correct, otherwise 99.99% OOM
-            return self.durations[index] * self.target_sample_rate / self.hop_length
-        return self.data[index]["duration"] * self.target_sample_rate / self.hop_length
-
-    def __len__(self):
-        return len(self.data)
-
-    def __getitem__(self, index):
-        row = self.data[index]
-        audio_path = row["audio_path"]
-        text = row["text"]
-        duration = row["duration"]
-
-        if self.preprocessed_mel:
-            mel_spec = torch.tensor(row["mel_spec"])
-
-        else:
-            audio, source_sample_rate = torchaudio.load(audio_path)
-            if audio.shape[0] > 1:
-                audio = torch.mean(audio, dim=0, keepdim=True)
-
-            if duration > 30 or duration < 0.3:
-                return self.__getitem__((index + 1) % len(self.data))
-
-            if source_sample_rate != self.target_sample_rate:
-                resampler = torchaudio.transforms.Resample(source_sample_rate, self.target_sample_rate)
-                audio = resampler(audio)
-
-            mel_spec = self.mel_spectrogram(audio)
-            mel_spec = mel_spec.squeeze(0)  # '1 d t -> d t')
-
-        return dict(
-            mel_spec=mel_spec,
-            text=text,
-        )
-
-
-# Dynamic Batch Sampler
-
-
-class DynamicBatchSampler(Sampler[list[int]]):
-    """Extension of Sampler that will do the following:
-    1.  Change the batch size (essentially number of sequences)
-        in a batch to ensure that the total number of frames are less
-        than a certain threshold.
-    2.  Make sure the padding efficiency in the batch is high.
-    """
-
-    def __init__(
-        self, sampler: Sampler[int], frames_threshold: int, max_samples=0, random_seed=None, drop_last: bool = False
-    ):
-        self.sampler = sampler
-        self.frames_threshold = frames_threshold
-        self.max_samples = max_samples
-
-        indices, batches = [], []
-        data_source = self.sampler.data_source
-
-        for idx in tqdm(
-            self.sampler, desc="Sorting with sampler... if slow, check whether dataset is provided with duration"
-        ):
-            indices.append((idx, data_source.get_frame_len(idx)))
-        indices.sort(key=lambda elem: elem[1])
-
-        batch = []
-        batch_frames = 0
-        for idx, frame_len in tqdm(
-            indices, desc=f"Creating dynamic batches with {frames_threshold} audio frames per gpu"
-        ):
-            if batch_frames + frame_len <= self.frames_threshold and (max_samples == 0 or len(batch) < max_samples):
-                batch.append(idx)
-                batch_frames += frame_len
-            else:
-                if len(batch) > 0:
-                    batches.append(batch)
-                if frame_len <= self.frames_threshold:
-                    batch = [idx]
-                    batch_frames = frame_len
-                else:
-                    batch = []
-                    batch_frames = 0
-
-        if not drop_last and len(batch) > 0:
-            batches.append(batch)
-
-        del indices
-
-        # if want to have different batches between epochs, may just set a seed and log it in ckpt
-        # cuz during multi-gpu training, although the batch on per gpu not change between epochs, the formed general minibatch is different
-        # e.g. for epoch n, use (random_seed + n)
-        random.seed(random_seed)
-        random.shuffle(batches)
-
-        self.batches = batches
-
-    def __iter__(self):
-        return iter(self.batches)
-
-    def __len__(self):
-        return len(self.batches)
-
-
-# Load dataset
-
-
-def load_dataset(
-    dataset_name: str,
-    tokenizer: str = "pinyin",
-    dataset_type: str = "CustomDataset",
-    audio_type: str = "raw",
-    mel_spec_module: nn.Module | None = None,
-    mel_spec_kwargs: dict = dict(),
-) -> CustomDataset | HFDataset:
-    """
-    dataset_type    - "CustomDataset" if you want to use tokenizer name and default data path to load for train_dataset
-                    - "CustomDatasetPath" if you just want to pass the full path to a preprocessed dataset without relying on tokenizer
-    """
-
-    print("Loading dataset ...")
-
-    if dataset_type == "CustomDataset":
-        rel_data_path = str(files("f5_tts").joinpath(f"../../data/{dataset_name}_{tokenizer}"))
-        if audio_type == "raw":
-            try:
-                train_dataset = load_from_disk(f"{rel_data_path}/raw")
-            except:  # noqa: E722
-                train_dataset = Dataset_.from_file(f"{rel_data_path}/raw.arrow")
-            preprocessed_mel = False
-        elif audio_type == "mel":
-            train_dataset = Dataset_.from_file(f"{rel_data_path}/mel.arrow")
-            preprocessed_mel = True
-        with open(f"{rel_data_path}/duration.json", "r", encoding="utf-8") as f:
-            data_dict = json.load(f)
-        durations = data_dict["duration"]
-        train_dataset = CustomDataset(
-            train_dataset,
-            durations=durations,
-            preprocessed_mel=preprocessed_mel,
-            mel_spec_module=mel_spec_module,
-            **mel_spec_kwargs,
-        )
-
-    elif dataset_type == "CustomDatasetPath":
-        try:
-            train_dataset = load_from_disk(f"{dataset_name}/raw")
-        except:  # noqa: E722
-            train_dataset = Dataset_.from_file(f"{dataset_name}/raw.arrow")
-
-        with open(f"{dataset_name}/duration.json", "r", encoding="utf-8") as f:
-            data_dict = json.load(f)
-        durations = data_dict["duration"]
-        train_dataset = CustomDataset(
-            train_dataset, durations=durations, preprocessed_mel=preprocessed_mel, **mel_spec_kwargs
-        )
-
-    elif dataset_type == "HFDataset":
-        print(
-            "Should manually modify the path of huggingface dataset to your need.\n"
-            + "May also the corresponding script cuz different dataset may have different format."
-        )
-        pre, post = dataset_name.split("_")
-        train_dataset = HFDataset(
-            load_dataset(f"{pre}/{pre}", split=f"train.{post}", cache_dir=str(files("f5_tts").joinpath("../../data"))),
-        )
-
-    return train_dataset
-
-
-# collation
-
-
-def collate_fn(batch):
-    mel_specs = [item["mel_spec"].squeeze(0) for item in batch]
-    mel_lengths = torch.LongTensor([spec.shape[-1] for spec in mel_specs])
-    max_mel_length = mel_lengths.amax()
-
-    padded_mel_specs = []
-    for spec in mel_specs:  # TODO. maybe records mask for attention here
-        padding = (0, max_mel_length - spec.size(-1))
-        padded_spec = F.pad(spec, padding, value=0)
-        padded_mel_specs.append(padded_spec)
-
-    mel_specs = torch.stack(padded_mel_specs)
-
-    text = [item["text"] for item in batch]
-    text_lengths = torch.LongTensor([len(item) for item in text])
-
-    return dict(
-        mel=mel_specs,
-        mel_lengths=mel_lengths,
-        text=text,
-        text_lengths=text_lengths,
-    )

src/f5-tts/model/modules.py

@@ -1,658 +0,0 @@
-"""
-ein notation:
-b - batch
-n - sequence
-nt - text sequence
-nw - raw wave length
-d - dimension
-"""
-
-from __future__ import annotations
-
-import math
-from typing import Optional
-
-import torch
-import torch.nn.functional as F
-import torchaudio
-from librosa.filters import mel as librosa_mel_fn
-from torch import nn
-from x_transformers.x_transformers import apply_rotary_pos_emb
-
-
-# raw wav to mel spec
-
-
-mel_basis_cache = {}
-hann_window_cache = {}
-
-
-def get_bigvgan_mel_spectrogram(
-    waveform,
-    n_fft=1024,
-    n_mel_channels=100,
-    target_sample_rate=24000,
-    hop_length=256,
-    win_length=1024,
-    fmin=0,
-    fmax=None,
-    center=False,
-):  # Copy from https://github.com/NVIDIA/BigVGAN/tree/main
-    device = waveform.device
-    key = f"{n_fft}_{n_mel_channels}_{target_sample_rate}_{hop_length}_{win_length}_{fmin}_{fmax}_{device}"
-
-    if key not in mel_basis_cache:
-        mel = librosa_mel_fn(sr=target_sample_rate, n_fft=n_fft, n_mels=n_mel_channels, fmin=fmin, fmax=fmax)
-        mel_basis_cache[key] = torch.from_numpy(mel).float().to(device)  # TODO: why they need .float()?
-        hann_window_cache[key] = torch.hann_window(win_length).to(device)
-
-    mel_basis = mel_basis_cache[key]
-    hann_window = hann_window_cache[key]
-
-    padding = (n_fft - hop_length) // 2
-    waveform = torch.nn.functional.pad(waveform.unsqueeze(1), (padding, padding), mode="reflect").squeeze(1)
-
-    spec = torch.stft(
-        waveform,
-        n_fft,
-        hop_length=hop_length,
-        win_length=win_length,
-        window=hann_window,
-        center=center,
-        pad_mode="reflect",
-        normalized=False,
-        onesided=True,
-        return_complex=True,
-    )
-    spec = torch.sqrt(torch.view_as_real(spec).pow(2).sum(-1) + 1e-9)
-
-    mel_spec = torch.matmul(mel_basis, spec)
-    mel_spec = torch.log(torch.clamp(mel_spec, min=1e-5))
-
-    return mel_spec
-
-
-def get_vocos_mel_spectrogram(
-    waveform,
-    n_fft=1024,
-    n_mel_channels=100,
-    target_sample_rate=24000,
-    hop_length=256,
-    win_length=1024,
-):
-    mel_stft = torchaudio.transforms.MelSpectrogram(
-        sample_rate=target_sample_rate,
-        n_fft=n_fft,
-        win_length=win_length,
-        hop_length=hop_length,
-        n_mels=n_mel_channels,
-        power=1,
-        center=True,
-        normalized=False,
-        norm=None,
-    ).to(waveform.device)
-    if len(waveform.shape) == 3:
-        waveform = waveform.squeeze(1)  # 'b 1 nw -> b nw'
-
-    assert len(waveform.shape) == 2
-
-    mel = mel_stft(waveform)
-    mel = mel.clamp(min=1e-5).log()
-    return mel
-
-
-class MelSpec(nn.Module):
-    def __init__(
-        self,
-        n_fft=1024,
-        hop_length=256,
-        win_length=1024,
-        n_mel_channels=100,
-        target_sample_rate=24_000,
-        mel_spec_type="vocos",
-    ):
-        super().__init__()
-        assert mel_spec_type in ["vocos", "bigvgan"], print("We only support two extract mel backend: vocos or bigvgan")
-
-        self.n_fft = n_fft
-        self.hop_length = hop_length
-        self.win_length = win_length
-        self.n_mel_channels = n_mel_channels
-        self.target_sample_rate = target_sample_rate
-
-        if mel_spec_type == "vocos":
-            self.extractor = get_vocos_mel_spectrogram
-        elif mel_spec_type == "bigvgan":
-            self.extractor = get_bigvgan_mel_spectrogram
-
-        self.register_buffer("dummy", torch.tensor(0), persistent=False)
-
-    def forward(self, wav):
-        if self.dummy.device != wav.device:
-            self.to(wav.device)
-
-        mel = self.extractor(
-            waveform=wav,
-            n_fft=self.n_fft,
-            n_mel_channels=self.n_mel_channels,
-            target_sample_rate=self.target_sample_rate,
-            hop_length=self.hop_length,
-            win_length=self.win_length,
-        )
-
-        return mel
-
-
-# sinusoidal position embedding
-
-
-class SinusPositionEmbedding(nn.Module):
-    def __init__(self, dim):
-        super().__init__()
-        self.dim = dim
-
-    def forward(self, x, scale=1000):
-        device = x.device
-        half_dim = self.dim // 2
-        emb = math.log(10000) / (half_dim - 1)
-        emb = torch.exp(torch.arange(half_dim, device=device).float() * -emb)
-        emb = scale * x.unsqueeze(1) * emb.unsqueeze(0)
-        emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
-        return emb
-
-
-# convolutional position embedding
-
-
-class ConvPositionEmbedding(nn.Module):
-    def __init__(self, dim, kernel_size=31, groups=16):
-        super().__init__()
-        assert kernel_size % 2 != 0
-        self.conv1d = nn.Sequential(
-            nn.Conv1d(dim, dim, kernel_size, groups=groups, padding=kernel_size // 2),
-            nn.Mish(),
-            nn.Conv1d(dim, dim, kernel_size, groups=groups, padding=kernel_size // 2),
-            nn.Mish(),
-        )
-
-    def forward(self, x: float["b n d"], mask: bool["b n"] | None = None):  # noqa: F722
-        if mask is not None:
-            mask = mask[..., None]
-            x = x.masked_fill(~mask, 0.0)
-
-        x = x.permute(0, 2, 1)
-        x = self.conv1d(x)
-        out = x.permute(0, 2, 1)
-
-        if mask is not None:
-            out = out.masked_fill(~mask, 0.0)
-
-        return out
-
-
-# rotary positional embedding related
-
-
-def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0, theta_rescale_factor=1.0):
-    # proposed by reddit user bloc97, to rescale rotary embeddings to longer sequence length without fine-tuning
-    # has some connection to NTK literature
-    # https://www.reddit.com/r/LocalLLaMA/comments/14lz7j5/ntkaware_scaled_rope_allows_llama_models_to_have/
-    # https://github.com/lucidrains/rotary-embedding-torch/blob/main/rotary_embedding_torch/rotary_embedding_torch.py
-    theta *= theta_rescale_factor ** (dim / (dim - 2))
-    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
-    t = torch.arange(end, device=freqs.device)  # type: ignore
-    freqs = torch.outer(t, freqs).float()  # type: ignore
-    freqs_cos = torch.cos(freqs)  # real part
-    freqs_sin = torch.sin(freqs)  # imaginary part
-    return torch.cat([freqs_cos, freqs_sin], dim=-1)
-
-
-def get_pos_embed_indices(start, length, max_pos, scale=1.0):
-    # length = length if isinstance(length, int) else length.max()
-    scale = scale * torch.ones_like(start, dtype=torch.float32)  # in case scale is a scalar
-    pos = (
-        start.unsqueeze(1)
-        + (torch.arange(length, device=start.device, dtype=torch.float32).unsqueeze(0) * scale.unsqueeze(1)).long()
-    )
-    # avoid extra long error.
-    pos = torch.where(pos < max_pos, pos, max_pos - 1)
-    return pos
-
-
-# Global Response Normalization layer (Instance Normalization ?)
-
-
-class GRN(nn.Module):
-    def __init__(self, dim):
-        super().__init__()
-        self.gamma = nn.Parameter(torch.zeros(1, 1, dim))
-        self.beta = nn.Parameter(torch.zeros(1, 1, dim))
-
-    def forward(self, x):
-        Gx = torch.norm(x, p=2, dim=1, keepdim=True)
-        Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6)
-        return self.gamma * (x * Nx) + self.beta + x
-
-
-# ConvNeXt-V2 Block https://github.com/facebookresearch/ConvNeXt-V2/blob/main/models/convnextv2.py
-# ref: https://github.com/bfs18/e2_tts/blob/main/rfwave/modules.py#L108
-
-
-class ConvNeXtV2Block(nn.Module):
-    def __init__(
-        self,
-        dim: int,
-        intermediate_dim: int,
-        dilation: int = 1,
-    ):
-        super().__init__()
-        padding = (dilation * (7 - 1)) // 2
-        self.dwconv = nn.Conv1d(
-            dim, dim, kernel_size=7, padding=padding, groups=dim, dilation=dilation
-        )  # depthwise conv
-        self.norm = nn.LayerNorm(dim, eps=1e-6)
-        self.pwconv1 = nn.Linear(dim, intermediate_dim)  # pointwise/1x1 convs, implemented with linear layers
-        self.act = nn.GELU()
-        self.grn = GRN(intermediate_dim)
-        self.pwconv2 = nn.Linear(intermediate_dim, dim)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        residual = x
-        x = x.transpose(1, 2)  # b n d -> b d n
-        x = self.dwconv(x)
-        x = x.transpose(1, 2)  # b d n -> b n d
-        x = self.norm(x)
-        x = self.pwconv1(x)
-        x = self.act(x)
-        x = self.grn(x)
-        x = self.pwconv2(x)
-        return residual + x
-
-
-# AdaLayerNormZero
-# return with modulated x for attn input, and params for later mlp modulation
-
-
-class AdaLayerNormZero(nn.Module):
-    def __init__(self, dim):
-        super().__init__()
-
-        self.silu = nn.SiLU()
-        self.linear = nn.Linear(dim, dim * 6)
-
-        self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
-
-    def forward(self, x, emb=None):
-        emb = self.linear(self.silu(emb))
-        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = torch.chunk(emb, 6, dim=1)
-
-        x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]
-        return x, gate_msa, shift_mlp, scale_mlp, gate_mlp
-
-
-# AdaLayerNormZero for final layer
-# return only with modulated x for attn input, cuz no more mlp modulation
-
-
-class AdaLayerNormZero_Final(nn.Module):
-    def __init__(self, dim):
-        super().__init__()
-
-        self.silu = nn.SiLU()
-        self.linear = nn.Linear(dim, dim * 2)
-
-        self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
-
-    def forward(self, x, emb):
-        emb = self.linear(self.silu(emb))
-        scale, shift = torch.chunk(emb, 2, dim=1)
-
-        x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :]
-        return x
-
-
-# FeedForward
-
-
-class FeedForward(nn.Module):
-    def __init__(self, dim, dim_out=None, mult=4, dropout=0.0, approximate: str = "none"):
-        super().__init__()
-        inner_dim = int(dim * mult)
-        dim_out = dim_out if dim_out is not None else dim
-
-        activation = nn.GELU(approximate=approximate)
-        project_in = nn.Sequential(nn.Linear(dim, inner_dim), activation)
-        self.ff = nn.Sequential(project_in, nn.Dropout(dropout), nn.Linear(inner_dim, dim_out))
-
-    def forward(self, x):
-        return self.ff(x)
-
-
-# Attention with possible joint part
-# modified from diffusers/src/diffusers/models/attention_processor.py
-
-
-class Attention(nn.Module):
-    def __init__(
-        self,
-        processor: JointAttnProcessor | AttnProcessor,
-        dim: int,
-        heads: int = 8,
-        dim_head: int = 64,
-        dropout: float = 0.0,
-        context_dim: Optional[int] = None,  # if not None -> joint attention
-        context_pre_only=None,
-    ):
-        super().__init__()
-
-        if not hasattr(F, "scaled_dot_product_attention"):
-            raise ImportError("Attention equires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
-
-        self.processor = processor
-
-        self.dim = dim
-        self.heads = heads
-        self.inner_dim = dim_head * heads
-        self.dropout = dropout
-
-        self.context_dim = context_dim
-        self.context_pre_only = context_pre_only
-
-        self.to_q = nn.Linear(dim, self.inner_dim)
-        self.to_k = nn.Linear(dim, self.inner_dim)
-        self.to_v = nn.Linear(dim, self.inner_dim)
-
-        if self.context_dim is not None:
-            self.to_k_c = nn.Linear(context_dim, self.inner_dim)
-            self.to_v_c = nn.Linear(context_dim, self.inner_dim)
-            if self.context_pre_only is not None:
-                self.to_q_c = nn.Linear(context_dim, self.inner_dim)
-
-        self.to_out = nn.ModuleList([])
-        self.to_out.append(nn.Linear(self.inner_dim, dim))
-        self.to_out.append(nn.Dropout(dropout))
-
-        if self.context_pre_only is not None and not self.context_pre_only:
-            self.to_out_c = nn.Linear(self.inner_dim, dim)
-
-    def forward(
-        self,
-        x: float["b n d"],  # noised input x  # noqa: F722
-        c: float["b n d"] = None,  # context c  # noqa: F722
-        mask: bool["b n"] | None = None,  # noqa: F722
-        rope=None,  # rotary position embedding for x
-        c_rope=None,  # rotary position embedding for c
-    ) -> torch.Tensor:
-        if c is not None:
-            return self.processor(self, x, c=c, mask=mask, rope=rope, c_rope=c_rope)
-        else:
-            return self.processor(self, x, mask=mask, rope=rope)
-
-
-# Attention processor
-
-
-class AttnProcessor:
-    def __init__(self):
-        pass
-
-    def __call__(
-        self,
-        attn: Attention,
-        x: float["b n d"],  # noised input x  # noqa: F722
-        mask: bool["b n"] | None = None,  # noqa: F722
-        rope=None,  # rotary position embedding
-    ) -> torch.FloatTensor:
-        batch_size = x.shape[0]
-
-        # `sample` projections.
-        query = attn.to_q(x)
-        key = attn.to_k(x)
-        value = attn.to_v(x)
-
-        # apply rotary position embedding
-        if rope is not None:
-            freqs, xpos_scale = rope
-            q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0)
-
-            query = apply_rotary_pos_emb(query, freqs, q_xpos_scale)
-            key = apply_rotary_pos_emb(key, freqs, k_xpos_scale)
-
-        # attention
-        inner_dim = key.shape[-1]
-        head_dim = inner_dim // attn.heads
-        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
-        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
-        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
-
-        # mask. e.g. inference got a batch with different target durations, mask out the padding
-        if mask is not None:
-            attn_mask = mask
-            attn_mask = attn_mask.unsqueeze(1).unsqueeze(1)  # 'b n -> b 1 1 n'
-            attn_mask = attn_mask.expand(batch_size, attn.heads, query.shape[-2], key.shape[-2])
-        else:
-            attn_mask = None
-
-        x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False)
-        x = x.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
-        x = x.to(query.dtype)
-
-        # linear proj
-        x = attn.to_out[0](x)
-        # dropout
-        x = attn.to_out[1](x)
-
-        if mask is not None:
-            mask = mask.unsqueeze(-1)
-            x = x.masked_fill(~mask, 0.0)
-
-        return x
-
-
-# Joint Attention processor for MM-DiT
-# modified from diffusers/src/diffusers/models/attention_processor.py
-
-
-class JointAttnProcessor:
-    def __init__(self):
-        pass
-
-    def __call__(
-        self,
-        attn: Attention,
-        x: float["b n d"],  # noised input x  # noqa: F722
-        c: float["b nt d"] = None,  # context c, here text # noqa: F722
-        mask: bool["b n"] | None = None,  # noqa: F722
-        rope=None,  # rotary position embedding for x
-        c_rope=None,  # rotary position embedding for c
-    ) -> torch.FloatTensor:
-        residual = x
-
-        batch_size = c.shape[0]
-
-        # `sample` projections.
-        query = attn.to_q(x)
-        key = attn.to_k(x)
-        value = attn.to_v(x)
-
-        # `context` projections.
-        c_query = attn.to_q_c(c)
-        c_key = attn.to_k_c(c)
-        c_value = attn.to_v_c(c)
-
-        # apply rope for context and noised input independently
-        if rope is not None:
-            freqs, xpos_scale = rope
-            q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0)
-            query = apply_rotary_pos_emb(query, freqs, q_xpos_scale)
-            key = apply_rotary_pos_emb(key, freqs, k_xpos_scale)
-        if c_rope is not None:
-            freqs, xpos_scale = c_rope
-            q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0)
-            c_query = apply_rotary_pos_emb(c_query, freqs, q_xpos_scale)
-            c_key = apply_rotary_pos_emb(c_key, freqs, k_xpos_scale)
-
-        # attention
-        query = torch.cat([query, c_query], dim=1)
-        key = torch.cat([key, c_key], dim=1)
-        value = torch.cat([value, c_value], dim=1)
-
-        inner_dim = key.shape[-1]
-        head_dim = inner_dim // attn.heads
-        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
-        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
-        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
-
-        # mask. e.g. inference got a batch with different target durations, mask out the padding
-        if mask is not None:
-            attn_mask = F.pad(mask, (0, c.shape[1]), value=True)  # no mask for c (text)
-            attn_mask = attn_mask.unsqueeze(1).unsqueeze(1)  # 'b n -> b 1 1 n'
-            attn_mask = attn_mask.expand(batch_size, attn.heads, query.shape[-2], key.shape[-2])
-        else:
-            attn_mask = None
-
-        x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False)
-        x = x.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
-        x = x.to(query.dtype)
-
-        # Split the attention outputs.
-        x, c = (
-            x[:, : residual.shape[1]],
-            x[:, residual.shape[1] :],
-        )
-
-        # linear proj
-        x = attn.to_out[0](x)
-        # dropout
-        x = attn.to_out[1](x)
-        if not attn.context_pre_only:
-            c = attn.to_out_c(c)
-
-        if mask is not None:
-            mask = mask.unsqueeze(-1)
-            x = x.masked_fill(~mask, 0.0)
-            # c = c.masked_fill(~mask, 0.)  # no mask for c (text)
-
-        return x, c
-
-
-# DiT Block
-
-
-class DiTBlock(nn.Module):
-    def __init__(self, dim, heads, dim_head, ff_mult=4, dropout=0.1):
-        super().__init__()
-
-        self.attn_norm = AdaLayerNormZero(dim)
-        self.attn = Attention(
-            processor=AttnProcessor(),
-            dim=dim,
-            heads=heads,
-            dim_head=dim_head,
-            dropout=dropout,
-        )
-
-        self.ff_norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
-        self.ff = FeedForward(dim=dim, mult=ff_mult, dropout=dropout, approximate="tanh")
-
-    def forward(self, x, t, mask=None, rope=None):  # x: noised input, t: time embedding
-        # pre-norm & modulation for attention input
-        norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(x, emb=t)
-
-        # attention
-        attn_output = self.attn(x=norm, mask=mask, rope=rope)
-
-        # process attention output for input x
-        x = x + gate_msa.unsqueeze(1) * attn_output
-
-        norm = self.ff_norm(x) * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
-        ff_output = self.ff(norm)
-        x = x + gate_mlp.unsqueeze(1) * ff_output
-
-        return x
-
-
-# MMDiT Block https://arxiv.org/abs/2403.03206
-
-
-class MMDiTBlock(nn.Module):
-    r"""
-    modified from diffusers/src/diffusers/models/attention.py
-
-    notes.
-    _c: context related. text, cond, etc. (left part in sd3 fig2.b)
-    _x: noised input related. (right part)
-    context_pre_only: last layer only do prenorm + modulation cuz no more ffn
-    """
-
-    def __init__(self, dim, heads, dim_head, ff_mult=4, dropout=0.1, context_pre_only=False):
-        super().__init__()
-
-        self.context_pre_only = context_pre_only
-
-        self.attn_norm_c = AdaLayerNormZero_Final(dim) if context_pre_only else AdaLayerNormZero(dim)
-        self.attn_norm_x = AdaLayerNormZero(dim)
-        self.attn = Attention(
-            processor=JointAttnProcessor(),
-            dim=dim,
-            heads=heads,
-            dim_head=dim_head,
-            dropout=dropout,
-            context_dim=dim,
-            context_pre_only=context_pre_only,
-        )
-
-        if not context_pre_only:
-            self.ff_norm_c = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
-            self.ff_c = FeedForward(dim=dim, mult=ff_mult, dropout=dropout, approximate="tanh")
-        else:
-            self.ff_norm_c = None
-            self.ff_c = None
-        self.ff_norm_x = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
-        self.ff_x = FeedForward(dim=dim, mult=ff_mult, dropout=dropout, approximate="tanh")
-
-    def forward(self, x, c, t, mask=None, rope=None, c_rope=None):  # x: noised input, c: context, t: time embedding
-        # pre-norm & modulation for attention input
-        if self.context_pre_only:
-            norm_c = self.attn_norm_c(c, t)
-        else:
-            norm_c, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.attn_norm_c(c, emb=t)
-        norm_x, x_gate_msa, x_shift_mlp, x_scale_mlp, x_gate_mlp = self.attn_norm_x(x, emb=t)
-
-        # attention
-        x_attn_output, c_attn_output = self.attn(x=norm_x, c=norm_c, mask=mask, rope=rope, c_rope=c_rope)
-
-        # process attention output for context c
-        if self.context_pre_only:
-            c = None
-        else:  # if not last layer
-            c = c + c_gate_msa.unsqueeze(1) * c_attn_output
-
-            norm_c = self.ff_norm_c(c) * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None]
-            c_ff_output = self.ff_c(norm_c)
-            c = c + c_gate_mlp.unsqueeze(1) * c_ff_output
-
-        # process attention output for input x
-        x = x + x_gate_msa.unsqueeze(1) * x_attn_output
-
-        norm_x = self.ff_norm_x(x) * (1 + x_scale_mlp[:, None]) + x_shift_mlp[:, None]
-        x_ff_output = self.ff_x(norm_x)
-        x = x + x_gate_mlp.unsqueeze(1) * x_ff_output
-
-        return c, x
-
-
-# time step conditioning embedding
-
-
-class TimestepEmbedding(nn.Module):
-    def __init__(self, dim, freq_embed_dim=256):
-        super().__init__()
-        self.time_embed = SinusPositionEmbedding(freq_embed_dim)
-        self.time_mlp = nn.Sequential(nn.Linear(freq_embed_dim, dim), nn.SiLU(), nn.Linear(dim, dim))
-
-    def forward(self, timestep: float["b"]):  # noqa: F821
-        time_hidden = self.time_embed(timestep)
-        time_hidden = time_hidden.to(timestep.dtype)
-        time = self.time_mlp(time_hidden)  # b d
-        return time

src/f5-tts/model/trainer.py

@@ -1,353 +0,0 @@
-from __future__ import annotations
-
-import gc
-import os
-
-import torch
-import torchaudio
-import wandb
-from accelerate import Accelerator
-from accelerate.utils import DistributedDataParallelKwargs
-from ema_pytorch import EMA
-from torch.optim import AdamW
-from torch.optim.lr_scheduler import LinearLR, SequentialLR
-from torch.utils.data import DataLoader, Dataset, SequentialSampler
-from tqdm import tqdm
-
-from f5_tts.model import CFM
-from f5_tts.model.dataset import DynamicBatchSampler, collate_fn
-from f5_tts.model.utils import default, exists
-
-# trainer
-
-
-class Trainer:
-    def __init__(
-        self,
-        model: CFM,
-        epochs,
-        learning_rate,
-        num_warmup_updates=20000,
-        save_per_updates=1000,
-        checkpoint_path=None,
-        batch_size=32,
-        batch_size_type: str = "sample",
-        max_samples=32,
-        grad_accumulation_steps=1,
-        max_grad_norm=1.0,
-        noise_scheduler: str | None = None,
-        duration_predictor: torch.nn.Module | None = None,
-        logger: str | None = "wandb",  # "wandb" | "tensorboard" | None
-        wandb_project="test_e2-tts",
-        wandb_run_name="test_run",
-        wandb_resume_id: str = None,
-        log_samples: bool = False,
-        last_per_steps=None,
-        accelerate_kwargs: dict = dict(),
-        ema_kwargs: dict = dict(),
-        bnb_optimizer: bool = False,
-        mel_spec_type: str = "vocos",  # "vocos" | "bigvgan"
-    ):
-        ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
-
-        if logger == "wandb" and not wandb.api.api_key:
-            logger = None
-        print(f"Using logger: {logger}")
-        self.log_samples = log_samples
-
-        self.accelerator = Accelerator(
-            log_with=logger if logger == "wandb" else None,
-            kwargs_handlers=[ddp_kwargs],
-            gradient_accumulation_steps=grad_accumulation_steps,
-            **accelerate_kwargs,
-        )
-
-        self.logger = logger
-        if self.logger == "wandb":
-            if exists(wandb_resume_id):
-                init_kwargs = {"wandb": {"resume": "allow", "name": wandb_run_name, "id": wandb_resume_id}}
-            else:
-                init_kwargs = {"wandb": {"resume": "allow", "name": wandb_run_name}}
-
-            self.accelerator.init_trackers(
-                project_name=wandb_project,
-                init_kwargs=init_kwargs,
-                config={
-                    "epochs": epochs,
-                    "learning_rate": learning_rate,
-                    "num_warmup_updates": num_warmup_updates,
-                    "batch_size": batch_size,
-                    "batch_size_type": batch_size_type,
-                    "max_samples": max_samples,
-                    "grad_accumulation_steps": grad_accumulation_steps,
-                    "max_grad_norm": max_grad_norm,
-                    "gpus": self.accelerator.num_processes,
-                    "noise_scheduler": noise_scheduler,
-                },
-            )
-
-        elif self.logger == "tensorboard":
-            from torch.utils.tensorboard import SummaryWriter
-
-            self.writer = SummaryWriter(log_dir=f"runs/{wandb_run_name}")
-
-        self.model = model
-
-        if self.is_main:
-            self.ema_model = EMA(model, include_online_model=False, **ema_kwargs)
-            self.ema_model.to(self.accelerator.device)
-
-        self.epochs = epochs
-        self.num_warmup_updates = num_warmup_updates
-        self.save_per_updates = save_per_updates
-        self.last_per_steps = default(last_per_steps, save_per_updates * grad_accumulation_steps)
-        self.checkpoint_path = default(checkpoint_path, "ckpts/test_e2-tts")
-
-        self.batch_size = batch_size
-        self.batch_size_type = batch_size_type
-        self.max_samples = max_samples
-        self.grad_accumulation_steps = grad_accumulation_steps
-        self.max_grad_norm = max_grad_norm
-        self.vocoder_name = mel_spec_type
-
-        self.noise_scheduler = noise_scheduler
-
-        self.duration_predictor = duration_predictor
-
-        if bnb_optimizer:
-            import bitsandbytes as bnb
-
-            self.optimizer = bnb.optim.AdamW8bit(model.parameters(), lr=learning_rate)
-        else:
-            self.optimizer = AdamW(model.parameters(), lr=learning_rate)
-        self.model, self.optimizer = self.accelerator.prepare(self.model, self.optimizer)
-
-    @property
-    def is_main(self):
-        return self.accelerator.is_main_process
-
-    def save_checkpoint(self, step, last=False):
-        self.accelerator.wait_for_everyone()
-        if self.is_main:
-            checkpoint = dict(
-                model_state_dict=self.accelerator.unwrap_model(self.model).state_dict(),
-                optimizer_state_dict=self.accelerator.unwrap_model(self.optimizer).state_dict(),
-                ema_model_state_dict=self.ema_model.state_dict(),
-                scheduler_state_dict=self.scheduler.state_dict(),
-                step=step,
-            )
-            if not os.path.exists(self.checkpoint_path):
-                os.makedirs(self.checkpoint_path)
-            if last:
-                self.accelerator.save(checkpoint, f"{self.checkpoint_path}/model_last.pt")
-                print(f"Saved last checkpoint at step {step}")
-            else:
-                self.accelerator.save(checkpoint, f"{self.checkpoint_path}/model_{step}.pt")
-
-    def load_checkpoint(self):
-        if (
-            not exists(self.checkpoint_path)
-            or not os.path.exists(self.checkpoint_path)
-            or not os.listdir(self.checkpoint_path)
-        ):
-            return 0
-
-        self.accelerator.wait_for_everyone()
-        if "model_last.pt" in os.listdir(self.checkpoint_path):
-            latest_checkpoint = "model_last.pt"
-        else:
-            latest_checkpoint = sorted(
-                [f for f in os.listdir(self.checkpoint_path) if f.endswith(".pt")],
-                key=lambda x: int("".join(filter(str.isdigit, x))),
-            )[-1]
-        # checkpoint = torch.load(f"{self.checkpoint_path}/{latest_checkpoint}", map_location=self.accelerator.device)  # rather use accelerator.load_state ಥ_ಥ
-        checkpoint = torch.load(f"{self.checkpoint_path}/{latest_checkpoint}", weights_only=True, map_location="cpu")
-
-        # patch for backward compatibility, 305e3ea
-        for key in ["ema_model.mel_spec.mel_stft.mel_scale.fb", "ema_model.mel_spec.mel_stft.spectrogram.window"]:
-            if key in checkpoint["ema_model_state_dict"]:
-                del checkpoint["ema_model_state_dict"][key]
-
-        if self.is_main:
-            self.ema_model.load_state_dict(checkpoint["ema_model_state_dict"])
-
-        if "step" in checkpoint:
-            # patch for backward compatibility, 305e3ea
-            for key in ["mel_spec.mel_stft.mel_scale.fb", "mel_spec.mel_stft.spectrogram.window"]:
-                if key in checkpoint["model_state_dict"]:
-                    del checkpoint["model_state_dict"][key]
-
-            self.accelerator.unwrap_model(self.model).load_state_dict(checkpoint["model_state_dict"])
-            self.accelerator.unwrap_model(self.optimizer).load_state_dict(checkpoint["optimizer_state_dict"])
-            if self.scheduler:
-                self.scheduler.load_state_dict(checkpoint["scheduler_state_dict"])
-            step = checkpoint["step"]
-        else:
-            checkpoint["model_state_dict"] = {
-                k.replace("ema_model.", ""): v
-                for k, v in checkpoint["ema_model_state_dict"].items()
-                if k not in ["initted", "step"]
-            }
-            self.accelerator.unwrap_model(self.model).load_state_dict(checkpoint["model_state_dict"])
-            step = 0
-
-        del checkpoint
-        gc.collect()
-        return step
-
-    def train(self, train_dataset: Dataset, num_workers=16, resumable_with_seed: int = None):
-        if self.log_samples:
-            from f5_tts.infer.utils_infer import cfg_strength, load_vocoder, nfe_step, sway_sampling_coef
-
-            vocoder = load_vocoder(vocoder_name=self.vocoder_name)
-            target_sample_rate = self.accelerator.unwrap_model(self.model).mel_spec.target_sample_rate
-            log_samples_path = f"{self.checkpoint_path}/samples"
-            os.makedirs(log_samples_path, exist_ok=True)
-
-        if exists(resumable_with_seed):
-            generator = torch.Generator()
-            generator.manual_seed(resumable_with_seed)
-        else:
-            generator = None
-
-        if self.batch_size_type == "sample":
-            train_dataloader = DataLoader(
-                train_dataset,
-                collate_fn=collate_fn,
-                num_workers=num_workers,
-                pin_memory=True,
-                persistent_workers=True,
-                batch_size=self.batch_size,
-                shuffle=True,
-                generator=generator,
-            )
-        elif self.batch_size_type == "frame":
-            self.accelerator.even_batches = False
-            sampler = SequentialSampler(train_dataset)
-            batch_sampler = DynamicBatchSampler(
-                sampler, self.batch_size, max_samples=self.max_samples, random_seed=resumable_with_seed, drop_last=False
-            )
-            train_dataloader = DataLoader(
-                train_dataset,
-                collate_fn=collate_fn,
-                num_workers=num_workers,
-                pin_memory=True,
-                persistent_workers=True,
-                batch_sampler=batch_sampler,
-            )
-        else:
-            raise ValueError(f"batch_size_type must be either 'sample' or 'frame', but received {self.batch_size_type}")
-
-        #  accelerator.prepare() dispatches batches to devices;
-        #  which means the length of dataloader calculated before, should consider the number of devices
-        warmup_steps = (
-            self.num_warmup_updates * self.accelerator.num_processes
-        )  # consider a fixed warmup steps while using accelerate multi-gpu ddp
-        # otherwise by default with split_batches=False, warmup steps change with num_processes
-        total_steps = len(train_dataloader) * self.epochs / self.grad_accumulation_steps
-        decay_steps = total_steps - warmup_steps
-        warmup_scheduler = LinearLR(self.optimizer, start_factor=1e-8, end_factor=1.0, total_iters=warmup_steps)
-        decay_scheduler = LinearLR(self.optimizer, start_factor=1.0, end_factor=1e-8, total_iters=decay_steps)
-        self.scheduler = SequentialLR(
-            self.optimizer, schedulers=[warmup_scheduler, decay_scheduler], milestones=[warmup_steps]
-        )
-        train_dataloader, self.scheduler = self.accelerator.prepare(
-            train_dataloader, self.scheduler
-        )  # actual steps = 1 gpu steps / gpus
-        start_step = self.load_checkpoint()
-        global_step = start_step
-
-        if exists(resumable_with_seed):
-            orig_epoch_step = len(train_dataloader)
-            skipped_epoch = int(start_step // orig_epoch_step)
-            skipped_batch = start_step % orig_epoch_step
-            skipped_dataloader = self.accelerator.skip_first_batches(train_dataloader, num_batches=skipped_batch)
-        else:
-            skipped_epoch = 0
-
-        for epoch in range(skipped_epoch, self.epochs):
-            self.model.train()
-            if exists(resumable_with_seed) and epoch == skipped_epoch:
-                progress_bar = tqdm(
-                    skipped_dataloader,
-                    desc=f"Epoch {epoch+1}/{self.epochs}",
-                    unit="step",
-                    disable=not self.accelerator.is_local_main_process,
-                    initial=skipped_batch,
-                    total=orig_epoch_step,
-                )
-            else:
-                progress_bar = tqdm(
-                    train_dataloader,
-                    desc=f"Epoch {epoch+1}/{self.epochs}",
-                    unit="step",
-                    disable=not self.accelerator.is_local_main_process,
-                )
-
-            for batch in progress_bar:
-                with self.accelerator.accumulate(self.model):
-                    text_inputs = batch["text"]
-                    mel_spec = batch["mel"].permute(0, 2, 1)
-                    mel_lengths = batch["mel_lengths"]
-
-                    # TODO. add duration predictor training
-                    if self.duration_predictor is not None and self.accelerator.is_local_main_process:
-                        dur_loss = self.duration_predictor(mel_spec, lens=batch.get("durations"))
-                        self.accelerator.log({"duration loss": dur_loss.item()}, step=global_step)
-
-                    loss, cond, pred = self.model(
-                        mel_spec, text=text_inputs, lens=mel_lengths, noise_scheduler=self.noise_scheduler
-                    )
-                    self.accelerator.backward(loss)
-
-                    if self.max_grad_norm > 0 and self.accelerator.sync_gradients:
-                        self.accelerator.clip_grad_norm_(self.model.parameters(), self.max_grad_norm)
-
-                    self.optimizer.step()
-                    self.scheduler.step()
-                    self.optimizer.zero_grad()
-
-                if self.is_main:
-                    self.ema_model.update()
-
-                global_step += 1
-
-                if self.accelerator.is_local_main_process:
-                    self.accelerator.log({"loss": loss.item(), "lr": self.scheduler.get_last_lr()[0]}, step=global_step)
-                    if self.logger == "tensorboard":
-                        self.writer.add_scalar("loss", loss.item(), global_step)
-                        self.writer.add_scalar("lr", self.scheduler.get_last_lr()[0], global_step)
-
-                progress_bar.set_postfix(step=str(global_step), loss=loss.item())
-
-                if global_step % (self.save_per_updates * self.grad_accumulation_steps) == 0:
-                    self.save_checkpoint(global_step)
-
-                    if self.log_samples and self.accelerator.is_local_main_process:
-                        ref_audio, ref_audio_len = vocoder.decode(batch["mel"][0].unsqueeze(0)), mel_lengths[0]
-                        torchaudio.save(
-                            f"{log_samples_path}/step_{global_step}_ref.wav", ref_audio.cpu(), target_sample_rate
-                        )
-                        with torch.inference_mode():
-                            generated, _ = self.accelerator.unwrap_model(self.model).sample(
-                                cond=mel_spec[0][:ref_audio_len].unsqueeze(0),
-                                text=[text_inputs[0] + [" "] + text_inputs[0]],
-                                duration=ref_audio_len * 2,
-                                steps=nfe_step,
-                                cfg_strength=cfg_strength,
-                                sway_sampling_coef=sway_sampling_coef,
-                            )
-                        generated = generated.to(torch.float32)
-                        gen_audio = vocoder.decode(
-                            generated[:, ref_audio_len:, :].permute(0, 2, 1).to(self.accelerator.device)
-                        )
-                        torchaudio.save(
-                            f"{log_samples_path}/step_{global_step}_gen.wav", gen_audio.cpu(), target_sample_rate
-                        )
-
-                if global_step % self.last_per_steps == 0:
-                    self.save_checkpoint(global_step, last=True)
-
-        self.save_checkpoint(global_step, last=True)
-
-        self.accelerator.end_training()

src/f5-tts/model/utils.py

@@ -1,185 +0,0 @@
-from __future__ import annotations
-
-import os
-import random
-from collections import defaultdict
-from importlib.resources import files
-
-import torch
-from torch.nn.utils.rnn import pad_sequence
-
-import jieba
-from pypinyin import lazy_pinyin, Style
-
-
-# seed everything
-
-
-def seed_everything(seed=0):
-    random.seed(seed)
-    os.environ["PYTHONHASHSEED"] = str(seed)
-    torch.manual_seed(seed)
-    torch.cuda.manual_seed(seed)
-    torch.cuda.manual_seed_all(seed)
-    torch.backends.cudnn.deterministic = True
-    torch.backends.cudnn.benchmark = False
-
-
-# helpers
-
-
-def exists(v):
-    return v is not None
-
-
-def default(v, d):
-    return v if exists(v) else d
-
-
-# tensor helpers
-
-
-def lens_to_mask(t: int["b"], length: int | None = None) -> bool["b n"]:  # noqa: F722 F821
-    if not exists(length):
-        length = t.amax()
-
-    seq = torch.arange(length, device=t.device)
-    return seq[None, :] < t[:, None]
-
-
-def mask_from_start_end_indices(seq_len: int["b"], start: int["b"], end: int["b"]):  # noqa: F722 F821
-    max_seq_len = seq_len.max().item()
-    seq = torch.arange(max_seq_len, device=start.device).long()
-    start_mask = seq[None, :] >= start[:, None]
-    end_mask = seq[None, :] < end[:, None]
-    return start_mask & end_mask
-
-
-def mask_from_frac_lengths(seq_len: int["b"], frac_lengths: float["b"]):  # noqa: F722 F821
-    lengths = (frac_lengths * seq_len).long()
-    max_start = seq_len - lengths
-
-    rand = torch.rand_like(frac_lengths)
-    start = (max_start * rand).long().clamp(min=0)
-    end = start + lengths
-
-    return mask_from_start_end_indices(seq_len, start, end)
-
-
-def maybe_masked_mean(t: float["b n d"], mask: bool["b n"] = None) -> float["b d"]:  # noqa: F722
-    if not exists(mask):
-        return t.mean(dim=1)
-
-    t = torch.where(mask[:, :, None], t, torch.tensor(0.0, device=t.device))
-    num = t.sum(dim=1)
-    den = mask.float().sum(dim=1)
-
-    return num / den.clamp(min=1.0)
-
-
-# simple utf-8 tokenizer, since paper went character based
-def list_str_to_tensor(text: list[str], padding_value=-1) -> int["b nt"]:  # noqa: F722
-    list_tensors = [torch.tensor([*bytes(t, "UTF-8")]) for t in text]  # ByT5 style
-    text = pad_sequence(list_tensors, padding_value=padding_value, batch_first=True)
-    return text
-
-
-# char tokenizer, based on custom dataset's extracted .txt file
-def list_str_to_idx(
-    text: list[str] | list[list[str]],
-    vocab_char_map: dict[str, int],  # {char: idx}
-    padding_value=-1,
-) -> int["b nt"]:  # noqa: F722
-    list_idx_tensors = [torch.tensor([vocab_char_map.get(c, 0) for c in t]) for t in text]  # pinyin or char style
-    text = pad_sequence(list_idx_tensors, padding_value=padding_value, batch_first=True)
-    return text
-
-
-# Get tokenizer
-
-
-def get_tokenizer(dataset_name, tokenizer: str = "pinyin"):
-    """
-    tokenizer   - "pinyin" do g2p for only chinese characters, need .txt vocab_file
-                - "char" for char-wise tokenizer, need .txt vocab_file
-                - "byte" for utf-8 tokenizer
-                - "custom" if you're directly passing in a path to the vocab.txt you want to use
-    vocab_size  - if use "pinyin", all available pinyin types, common alphabets (also those with accent) and symbols
-                - if use "char", derived from unfiltered character & symbol counts of custom dataset
-                - if use "byte", set to 256 (unicode byte range)
-    """
-    if tokenizer in ["pinyin", "char"]:
-        tokenizer_path = os.path.join(files("f5_tts").joinpath("../../data"), f"{dataset_name}_{tokenizer}/vocab.txt")
-        with open(tokenizer_path, "r", encoding="utf-8") as f:
-            vocab_char_map = {}
-            for i, char in enumerate(f):
-                vocab_char_map[char[:-1]] = i
-        vocab_size = len(vocab_char_map)
-        assert vocab_char_map[" "] == 0, "make sure space is of idx 0 in vocab.txt, cuz 0 is used for unknown char"
-
-    elif tokenizer == "byte":
-        vocab_char_map = None
-        vocab_size = 256
-
-    elif tokenizer == "custom":
-        with open(dataset_name, "r", encoding="utf-8") as f:
-            vocab_char_map = {}
-            for i, char in enumerate(f):
-                vocab_char_map[char[:-1]] = i
-        vocab_size = len(vocab_char_map)
-
-    return vocab_char_map, vocab_size
-
-
-# convert char to pinyin
-
-
-def convert_char_to_pinyin(text_list, polyphone=True):
-    final_text_list = []
-    god_knows_why_en_testset_contains_zh_quote = str.maketrans(
-        {"“": '"', "”": '"', "‘": "'", "’": "'"}
-    )  # in case librispeech (orig no-pc) test-clean
-    custom_trans = str.maketrans({";": ","})  # add custom trans here, to address oov
-    for text in text_list:
-        char_list = []
-        text = text.translate(god_knows_why_en_testset_contains_zh_quote)
-        text = text.translate(custom_trans)
-        for seg in jieba.cut(text):
-            seg_byte_len = len(bytes(seg, "UTF-8"))
-            if seg_byte_len == len(seg):  # if pure alphabets and symbols
-                if char_list and seg_byte_len > 1 and char_list[-1] not in " :'\"":
-                    char_list.append(" ")
-                char_list.extend(seg)
-            elif polyphone and seg_byte_len == 3 * len(seg):  # if pure chinese characters
-                seg = lazy_pinyin(seg, style=Style.TONE3, tone_sandhi=True)
-                for c in seg:
-                    if c not in "。，、；：？！《》【】—…":
-                        char_list.append(" ")
-                    char_list.append(c)
-            else:  # if mixed chinese characters, alphabets and symbols
-                for c in seg:
-                    if ord(c) < 256:
-                        char_list.extend(c)
-                    else:
-                        if c not in "。，、；：？！《》【】—…":
-                            char_list.append(" ")
-                            char_list.extend(lazy_pinyin(c, style=Style.TONE3, tone_sandhi=True))
-                        else:  # if is zh punc
-                            char_list.append(c)
-        final_text_list.append(char_list)
-
-    return final_text_list
-
-
-# filter func for dirty data with many repetitions
-
-
-def repetition_found(text, length=2, tolerance=10):
-    pattern_count = defaultdict(int)
-    for i in range(len(text) - length + 1):
-        pattern = text[i : i + length]
-        pattern_count[pattern] += 1
-    for pattern, count in pattern_count.items():
-        if count > tolerance:
-            return True
-    return False

src/f5-tts/socket.py

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