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,