Sync from GitHub repo

This Space is synced from the GitHub repo: https://github.com/SWivid/F5-TTS. Please submit contributions to the Space there

app.py

@@ -4,7 +4,6 @@ import torchaudio
 import gradio as gr
 import numpy as np
 import tempfile
-from einops import rearrange
 from vocos import Vocos
 from pydub import AudioSegment, silence
 from model import CFM, UNetT, DiT, MMDiT
@@ -175,7 +174,7 @@ def infer_batch(ref_audio, ref_text, gen_text_batches, exp_name, remove_silence,
 
         generated = generated.to(torch.float32)
         generated = generated[:, ref_audio_len:, :]
-        generated_mel_spec = rearrange(generated, "1 n d -> 1 d n")
+        generated_mel_spec = generated.permute(0, 2, 1)
         generated_wave = vocos.decode(generated_mel_spec.cpu())
         if rms < target_rms:
             generated_wave = generated_wave * rms / target_rms

inference-cli.py

@@ -11,7 +11,6 @@ import torch
 import torchaudio
 import tqdm
 from cached_path import cached_path
-from einops import rearrange
 from pydub import AudioSegment, silence
 from transformers import pipeline
 from vocos import Vocos
@@ -274,7 +273,7 @@ def infer_batch(ref_audio, ref_text, gen_text_batches, model, remove_silence, cr
 
         generated = generated.to(torch.float32)
         generated = generated[:, ref_audio_len:, :]
-        generated_mel_spec = rearrange(generated, "1 n d -> 1 d n")
+        generated_mel_spec = generated.permute(0, 2, 1)
         generated_wave = vocos.decode(generated_mel_spec.cpu())
         if rms < target_rms:
             generated_wave = generated_wave * rms / target_rms
@@ -427,4 +426,4 @@ def process(ref_audio, ref_text, text_gen, model, remove_silence):
             print(f.name)
 
 
-process(ref_audio, ref_text, gen_text, model, remove_silence)
+process(ref_audio, ref_text, gen_text, model, remove_silence)

model/backbones/dit.py

@@ -13,8 +13,6 @@ import torch
 from torch import nn
 import torch.nn.functional as F
 
-from einops import repeat
-
 from x_transformers.x_transformers import RotaryEmbedding
 
 from model.modules import (
@@ -134,7 +132,7 @@ class DiT(nn.Module):
     ):
         batch, seq_len = x.shape[0], x.shape[1]
         if time.ndim == 0:
-            time = repeat(time, ' -> b', b = batch)
+            time = time.repeat(batch)
         
         # t: conditioning time, c: context (text + masked cond audio), x: noised input audio
         t = self.time_embed(time)

model/backbones/mmdit.py

@@ -12,8 +12,6 @@ from __future__ import annotations
 import torch
 from torch import nn
 
-from einops import repeat
-
 from x_transformers.x_transformers import RotaryEmbedding
 
 from model.modules import (
@@ -115,7 +113,7 @@ class MMDiT(nn.Module):
     ):
         batch = x.shape[0]
         if time.ndim == 0:
-            time = repeat(time, ' -> b', b = batch)
+            time = time.repeat(batch)
 
         # t: conditioning (time), c: context (text + masked cond audio), x: noised input audio
         t = self.time_embed(time)

model/backbones/unett.py

@@ -14,8 +14,6 @@ import torch
 from torch import nn
 import torch.nn.functional as F
 
-from einops import repeat, pack, unpack
-
 from x_transformers import RMSNorm
 from x_transformers.x_transformers import RotaryEmbedding
 
@@ -155,7 +153,7 @@ class UNetT(nn.Module):
     ):
         batch, seq_len = x.shape[0], x.shape[1]
         if time.ndim == 0:
-            time = repeat(time, ' -> b', b = batch)
+            time = time.repeat(batch)
         
         # t: conditioning time, c: context (text + masked cond audio), x: noised input audio
         t = self.time_embed(time)
@@ -163,7 +161,7 @@ class UNetT(nn.Module):
         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, ps = pack((t, x), 'b * 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)
         
@@ -196,6 +194,6 @@ class UNetT(nn.Module):
 
         assert len(skips) == 0
 
-        _, x = unpack(self.norm_out(x), ps, 'b * d')
+        x = self.norm_out(x)[:, 1:, :]  # unpack t from x
 
         return self.proj_out(x)

model/cfm.py

@@ -18,10 +18,7 @@ from torch.nn.utils.rnn import pad_sequence
 
 from torchdiffeq import odeint
 
-from einops import rearrange
-
 from model.modules import MelSpec
-
 from model.utils import (
     default, exists, 
     list_str_to_idx, list_str_to_tensor, 
@@ -105,7 +102,7 @@ class CFM(nn.Module):
 
         if cond.ndim == 2:
             cond = self.mel_spec(cond)
-            cond = rearrange(cond, 'b d n -> b n d')
+            cond = cond.permute(0, 2, 1)
             assert cond.shape[-1] == self.num_channels
 
         batch, cond_seq_len, device = *cond.shape[:2], cond.device
@@ -144,7 +141,7 @@ class CFM(nn.Module):
             
         cond = F.pad(cond, (0, 0, 0, max_duration - cond_seq_len), value = 0.)
         cond_mask = F.pad(cond_mask, (0, max_duration - cond_mask.shape[-1]), value = False)
-        cond_mask = rearrange(cond_mask, '... -> ... 1')
+        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:
@@ -199,7 +196,7 @@ class CFM(nn.Module):
         out = torch.where(cond_mask, cond, out)
 
         if exists(vocoder):
-            out = rearrange(out, 'b n d -> b d n')
+            out = out.permute(0, 2, 1)
             out = vocoder(out)
 
         return out, trajectory
@@ -215,7 +212,7 @@ class CFM(nn.Module):
         # handle raw wave
         if inp.ndim == 2:
             inp = self.mel_spec(inp)
-            inp = rearrange(inp, 'b d n -> b n d')
+            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
@@ -252,7 +249,7 @@ class CFM(nn.Module):
         # TODO. noise_scheduler
 
         # sample xt (φ_t(x) in the paper)
-        t = rearrange(time, 'b -> b 1 1')
+        t = time.unsqueeze(-1).unsqueeze(-1)
         φ = (1 - t) * x0 + t * x1
         flow = x1 - x0
 

model/dataset.py

@@ -9,8 +9,6 @@ import torchaudio
 from datasets import load_dataset, load_from_disk
 from datasets import Dataset as Dataset_
 
-from einops import rearrange
-
 from model.modules import MelSpec
 
 
@@ -54,11 +52,11 @@ class HFDataset(Dataset):
             resampler = torchaudio.transforms.Resample(sample_rate, self.target_sample_rate)
             audio_tensor = resampler(audio_tensor)
         
-        audio_tensor = rearrange(audio_tensor, 't -> 1 t')
+        audio_tensor = audio_tensor.unsqueeze(0)  # 't -> 1 t')
         
         mel_spec = self.mel_spectrogram(audio_tensor)
         
-        mel_spec = rearrange(mel_spec, '1 d t -> d t')
+        mel_spec = mel_spec.squeeze(0)  # '1 d t -> d t'
         
         text = row['text']
         
@@ -114,7 +112,7 @@ class CustomDataset(Dataset):
                 audio = resampler(audio)
             
             mel_spec = self.mel_spectrogram(audio)
-            mel_spec = rearrange(mel_spec, '1 d t -> d t')
+            mel_spec = mel_spec.squeeze(0)  # '1 d t -> d t')
         
         return dict(
             mel_spec = mel_spec,

model/modules.py

@@ -16,7 +16,6 @@ from torch import nn
 import torch.nn.functional as F
 import torchaudio
 
-from einops import rearrange
 from x_transformers.x_transformers import apply_rotary_pos_emb
 
 
@@ -54,7 +53,7 @@ class MelSpec(nn.Module):
 
     def forward(self, inp):
         if len(inp.shape) == 3:
-            inp = rearrange(inp, 'b 1 nw -> b nw')
+            inp = inp.squeeze(1) # 'b 1 nw -> b nw'
 
         assert len(inp.shape) == 2
 
@@ -101,9 +100,9 @@ class ConvPositionEmbedding(nn.Module):
             mask = mask[..., None]
             x = x.masked_fill(~mask, 0.)
 
-        x = rearrange(x, 'b n d -> b d n')
+        x = x.permute(0, 2, 1)
         x = self.conv1d(x)
-        out = rearrange(x, 'b d n -> b n d')
+        out = x.permute(0, 2, 1)
 
         if mask is not None:
             out = out.masked_fill(~mask, 0.)
@@ -345,7 +344,7 @@ class AttnProcessor:
         # 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 = rearrange(attn_mask, 'b n -> b 1 1 n')
+            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
@@ -360,7 +359,7 @@ class AttnProcessor:
         x = attn.to_out[1](x)
 
         if mask is not None:
-            mask = rearrange(mask, 'b n -> b n 1')
+            mask = mask.unsqueeze(-1)
             x = x.masked_fill(~mask, 0.)
 
         return x
@@ -422,7 +421,7 @@ class JointAttnProcessor:
         # 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 = rearrange(attn_mask, 'b n -> b 1 1 n')
+            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
@@ -445,7 +444,7 @@ class JointAttnProcessor:
             c = attn.to_out_c(c)
 
         if mask is not None:
-            mask = rearrange(mask, 'b n -> b n 1')
+            mask = mask.unsqueeze(-1)
             x = x.masked_fill(~mask, 0.)
             # c = c.masked_fill(~mask, 0.)  # no mask for c (text)
 

model/trainer.py

@@ -10,8 +10,6 @@ from torch.optim import AdamW
 from torch.utils.data import DataLoader, Dataset, SequentialSampler
 from torch.optim.lr_scheduler import LinearLR, SequentialLR
 
-from einops import rearrange
-
 from accelerate import Accelerator
 from accelerate.utils import DistributedDataParallelKwargs
 
@@ -222,7 +220,7 @@ class Trainer:
             for batch in progress_bar:
                 with self.accelerator.accumulate(self.model):
                     text_inputs = batch['text']
-                    mel_spec = rearrange(batch['mel'], 'b d n -> b n d')
+                    mel_spec = batch['mel'].permute(0, 2, 1)
                     mel_lengths = batch["mel_lengths"]
 
                     # TODO. add duration predictor training

model/utils.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import os
-import re
 import math
 import random
 import string
@@ -17,9 +16,6 @@ import torch.nn.functional as F
 from torch.nn.utils.rnn import pad_sequence
 import torchaudio
 
-import einx
-from einops import rearrange, reduce
-
 import jieba
 from pypinyin import lazy_pinyin, Style
 
@@ -57,7 +53,7 @@ def lens_to_mask(
         length = t.amax()
 
     seq = torch.arange(length, device = t.device)
-    return einx.less('n, b -> b n', seq, t)
+    return seq[None, :] < t[:, None]
 
 def mask_from_start_end_indices(
     seq_len: int['b'],
@@ -66,7 +62,9 @@ def mask_from_start_end_indices(
 ):
     max_seq_len = seq_len.max().item()  
     seq = torch.arange(max_seq_len, device = start.device).long()
-    return einx.greater_equal('n, b -> b n', seq, start) & einx.less('n, b -> b n', seq, end)
+    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'],
@@ -89,11 +87,11 @@ def maybe_masked_mean(
     if not exists(mask):
         return t.mean(dim = 1)
 
-    t = einx.where('b n, b n d, -> b n d', mask, t, 0.)
-    num = reduce(t, 'b n d -> b d', 'sum')
-    den = reduce(mask.float(), 'b n -> b', 'sum')
+    t = torch.where(mask[:, :, None], t, torch.tensor(0., device=t.device))
+    num = t.sum(dim=1)
+    den = mask.float().sum(dim=1)
 
-    return einx.divide('b d, b -> b d', num, den.clamp(min = 1.))
+    return num / den.clamp(min=1.)
 
 
 # simple utf-8 tokenizer, since paper went character based
@@ -239,7 +237,7 @@ def padded_mel_batch(ref_mels):
         padded_ref_mel = F.pad(mel, (0, max_mel_length - mel.shape[-1]), value = 0)
         padded_ref_mels.append(padded_ref_mel)
     padded_ref_mels = torch.stack(padded_ref_mels)
-    padded_ref_mels = rearrange(padded_ref_mels, 'b d n -> b n d')
+    padded_ref_mels = padded_ref_mels.permute(0, 2, 1)
     return padded_ref_mels
 
 
@@ -302,7 +300,7 @@ def get_inference_prompt(
 
         # to mel spectrogram
         ref_mel = mel_spectrogram(ref_audio)
-        ref_mel = rearrange(ref_mel, '1 d n -> d n')
+        ref_mel = ref_mel.squeeze(0)
 
         # deal with batch
         assert infer_batch_size > 0, "infer_batch_size should be greater than 0."

requirements.txt

@@ -3,8 +3,6 @@ bitsandbytes>0.37.0
 cached_path
 click
 datasets
-einops>=0.8.0
-einx>=0.3.0
 ema_pytorch>=0.5.2
 gradio
 jieba

scripts/eval_infer_batch.py

@@ -9,7 +9,6 @@ import argparse
 import torch
 import torchaudio
 from accelerate import Accelerator
-from einops import rearrange
 from vocos import Vocos
 
 from model import CFM, UNetT, DiT
@@ -187,7 +186,7 @@ with accelerator.split_between_processes(prompts_all) as prompts:
         # Final result
         for i, gen in enumerate(generated):
             gen = gen[ref_mel_lens[i]:total_mel_lens[i], :].unsqueeze(0)
-            gen_mel_spec = rearrange(gen, '1 n d -> 1 d n')
+            gen_mel_spec = gen.permute(0, 2, 1)
             generated_wave = vocos.decode(gen_mel_spec.cpu())
             if ref_rms_list[i] < target_rms:
                 generated_wave = generated_wave * ref_rms_list[i] / target_rms

speech_edit.py

@@ -3,7 +3,6 @@ import os
 import torch
 import torch.nn.functional as F
 import torchaudio
-from einops import rearrange
 from vocos import Vocos
 
 from model import CFM, UNetT, DiT, MMDiT
@@ -174,7 +173,7 @@ print(f"Generated mel: {generated.shape}")
 # Final result
 generated = generated.to(torch.float32)
 generated = generated[:, ref_audio_len:, :]
-generated_mel_spec = rearrange(generated, '1 n d -> 1 d n')
+generated_mel_spec = generated.permute(0, 2, 1)
 generated_wave = vocos.decode(generated_mel_spec.cpu())
 if rms < target_rms:
     generated_wave = generated_wave * rms / target_rms

train.py

@@ -56,7 +56,7 @@ def main():
             hop_length = hop_length,
         )
     
-    e2tts = CFM(
+    model = CFM(
         transformer = model_cls(
             **model_cfg,
             text_num_embeds = vocab_size, 
@@ -67,7 +67,7 @@ def main():
     )
 
     trainer = Trainer(
-        e2tts,
+        model,
         epochs, 
         learning_rate,
         num_warmup_updates = num_warmup_updates,