Super-squash branch 'main' using huggingface_hub

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Yushen CHEN
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,13 @@
+---
+title: E2/F5 TTS
+emoji: 🗣️
+colorFrom: green
+colorTo: green
+sdk: gradio
+app_file: app.py
+pinned: true
+short_description: 'E2-TTS & F5-TTS: Zero-Shot Voice Cloning (Unofficial Demo)'
+sdk_version: 5.0.1
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,207 @@
+import os
+import re
+import torch
+import torchaudio
+import gradio as gr
+import numpy as np
+import tempfile
+from einops import rearrange
+from ema_pytorch import EMA
+from vocos import Vocos
+from pydub import AudioSegment
+from model import CFM, UNetT, DiT, MMDiT
+from cached_path import cached_path
+from model.utils import (
+    get_tokenizer, 
+    convert_char_to_pinyin, 
+    save_spectrogram,
+)
+from transformers import pipeline
+import spaces
+import librosa
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+pipe = pipeline(
+    "automatic-speech-recognition",
+    model="openai/whisper-large-v3-turbo",
+    torch_dtype=torch.float16,
+    device=device,
+)
+
+# --------------------- Settings -------------------- #
+
+target_sample_rate = 24000
+n_mel_channels = 100
+hop_length = 256
+target_rms = 0.1
+nfe_step = 32  # 16, 32
+cfg_strength = 2.0
+ode_method = 'euler'
+sway_sampling_coef = -1.0
+speed = 1.0
+# fix_duration = 27  # None or float (duration in seconds)
+fix_duration = None
+
+def load_model(exp_name, model_cls, model_cfg, ckpt_step):
+    checkpoint = torch.load(str(cached_path(f"hf://SWivid/F5-TTS/{exp_name}/model_{ckpt_step}.pt")), map_location=device)
+    vocab_char_map, vocab_size = get_tokenizer("Emilia_ZH_EN", "pinyin")
+    model = CFM(
+        transformer=model_cls(
+            **model_cfg,
+            text_num_embeds=vocab_size,
+            mel_dim=n_mel_channels
+        ),
+        mel_spec_kwargs=dict(
+            target_sample_rate=target_sample_rate,
+            n_mel_channels=n_mel_channels,
+            hop_length=hop_length,
+        ),
+        odeint_kwargs=dict(
+            method=ode_method,
+        ),
+        vocab_char_map=vocab_char_map,
+    ).to(device)
+
+    ema_model = EMA(model, include_online_model=False).to(device)
+    ema_model.load_state_dict(checkpoint['ema_model_state_dict'])
+    ema_model.copy_params_from_ema_to_model()
+
+    return ema_model, model
+
+# load models
+F5TTS_model_cfg = dict(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4)
+E2TTS_model_cfg = dict(dim=1024, depth=24, heads=16, ff_mult=4)
+
+F5TTS_ema_model, F5TTS_base_model = load_model("F5TTS_Base", DiT, F5TTS_model_cfg, 1200000)
+E2TTS_ema_model, E2TTS_base_model = load_model("E2TTS_Base", UNetT, E2TTS_model_cfg, 1200000)
+
+@spaces.GPU
+def infer(ref_audio_orig, ref_text, gen_text, exp_name, remove_silence):
+    print(gen_text)
+    if len(gen_text) > 200:
+        raise gr.Error("Please keep your text under 200 chars.")
+    gr.Info("Converting audio...")
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as f:
+        aseg = AudioSegment.from_file(ref_audio_orig)
+        audio_duration = len(aseg)
+        if audio_duration > 15000:
+            gr.Warning("Audio is over 15s, clipping to only first 15s.")
+            aseg = aseg[:15000]
+        aseg.export(f.name, format="wav")
+        ref_audio = f.name
+    if exp_name == "F5-TTS":
+        ema_model = F5TTS_ema_model
+        base_model = F5TTS_base_model
+    elif exp_name == "E2-TTS":
+        ema_model = E2TTS_ema_model
+        base_model = E2TTS_base_model
+    
+    if not ref_text.strip():
+        gr.Info("No reference text provided, transcribing reference audio...")
+        ref_text = outputs = pipe(
+            ref_audio,
+            chunk_length_s=30,
+            batch_size=128,
+            generate_kwargs={"task": "transcribe"},
+            return_timestamps=False,
+        )['text'].strip()
+        gr.Info("Finished transcription")
+    else:
+        gr.Info("Using custom reference text...")
+    audio, sr = torchaudio.load(ref_audio)
+
+    rms = torch.sqrt(torch.mean(torch.square(audio)))
+    if rms < target_rms:
+        audio = audio * target_rms / rms
+    if sr != target_sample_rate:
+        resampler = torchaudio.transforms.Resample(sr, target_sample_rate)
+        audio = resampler(audio)
+    audio = audio.to(device)
+
+    # Prepare the text
+    text_list = [ref_text + gen_text]
+    final_text_list = convert_char_to_pinyin(text_list)
+
+    # Calculate duration
+    ref_audio_len = audio.shape[-1] // hop_length
+    # if fix_duration is not None:
+    #     duration = int(fix_duration * target_sample_rate / hop_length)
+    # else:
+    zh_pause_punc = r"。，、；：？！"
+    ref_text_len = len(ref_text) + len(re.findall(zh_pause_punc, ref_text))
+    gen_text_len = len(gen_text) + len(re.findall(zh_pause_punc, gen_text))
+    duration = ref_audio_len + int(ref_audio_len / ref_text_len * gen_text_len / speed)
+
+    # inference
+    gr.Info(f"Generating audio using {exp_name}")
+    with torch.inference_mode():
+        generated, _ = base_model.sample(
+            cond=audio,
+            text=final_text_list,
+            duration=duration,
+            steps=nfe_step,
+            cfg_strength=cfg_strength,
+            sway_sampling_coef=sway_sampling_coef,
+        )
+
+    generated = generated[:, ref_audio_len:, :]
+    generated_mel_spec = rearrange(generated, '1 n d -> 1 d n')
+    gr.Info("Running vocoder")
+    vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")
+    generated_wave = vocos.decode(generated_mel_spec.cpu())
+    if rms < target_rms:
+        generated_wave = generated_wave * rms / target_rms
+
+    # wav -> numpy
+    generated_wave = generated_wave.squeeze().cpu().numpy()
+
+    if remove_silence:
+        gr.Info("Removing audio silences")
+        non_silent_intervals = librosa.effects.split(generated_wave, top_db=30)
+        non_silent_wave = np.array([])
+        for interval in non_silent_intervals:
+            start, end = interval
+            non_silent_wave = np.concatenate([non_silent_wave, generated_wave[start:end]])
+        generated_wave = non_silent_wave
+
+
+    # spectogram
+    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_spectrogram:
+        spectrogram_path = tmp_spectrogram.name
+        save_spectrogram(generated_mel_spec[0].cpu().numpy(), spectrogram_path)
+
+    return (target_sample_rate, generated_wave), spectrogram_path
+
+with gr.Blocks() as app:
+    gr.Markdown("""
+# E2/F5 TTS
+
+This is an unofficial E2/F5 TTS demo. This demo supports the following TTS models:
+
+* [E2-TTS](https://arxiv.org/abs/2406.18009) (Embarrassingly Easy Fully Non-Autoregressive Zero-Shot TTS)
+* [F5-TTS](https://arxiv.org/abs/2410.06885) (A Fairytaler that Fakes Fluent and Faithful Speech with Flow Matching)
+
+This demo is based on the [F5-TTS](https://github.com/SWivid/F5-TTS) codebase, which is based on an [unofficial E2-TTS implementation](https://github.com/lucidrains/e2-tts-pytorch).
+
+The checkpoints support English and Chinese.
+
+**NOTE: Reference text will be automatically transcribed with Whisper if not provided. For best results, keep your reference clips short (<15s). Ensure the audio is fully uploaded before generating.**
+""")
+
+    ref_audio_input = gr.Audio(label="Reference Audio", type="filepath")
+    gen_text_input = gr.Textbox(label="Text to Generate (max 200 chars.)", lines=4)
+    model_choice = gr.Radio(choices=["F5-TTS", "E2-TTS"], label="Choose TTS Model", value="F5-TTS")
+    generate_btn = gr.Button("Synthesize", variant="primary")
+    with gr.Accordion("Advanced Settings", open=False):
+        ref_text_input = gr.Textbox(label="Reference Text", info="Leave blank to automatically transcribe the reference audio. If you enter text it will override automatic transcription.", lines=2)
+        remove_silence = gr.Checkbox(label="[EXPERIMENTAL] Remove Silences", info="The model tends to leave silences, we can manually remove silences if needed. This may produce strange results and is not guarenteed to work.")
+    
+    audio_output = gr.Audio(label="Synthesized Audio")
+    spectrogram_output = gr.Image(label="Spectrogram")
+
+    generate_btn.click(infer, inputs=[ref_audio_input, ref_text_input, gen_text_input, model_choice, remove_silence], outputs=[audio_output, spectrogram_output])
+    gr.Markdown("Unofficial demo by [mrfakename](https://x.com/realmrfakename)")
+
+
+app.queue().launch()

ckpts/README.md

@@ -0,0 +1,10 @@
+
+Pretrained model ckpts. https://huggingface.co/SWivid/F5-TTS
+
+```
+ckpts/
+    E2TTS_Base/
+        model_1200000.pt
+    F5TTS_Base/
+        model_1200000.pt
+```

data/Emilia_ZH_EN_pinyin/vocab.txt

@@ -0,0 +1,2545 @@
+ 
+!
+"
+#
+$
+%
+&
+'
+(
+)
+*
++
+,
+-
+.
+/
+0
+1
+2
+3
+4
+5
+6
+7
+8
+9
+:
+;
+=
+>
+?
+@
+A
+B
+C
+D
+E
+F
+G
+H
+I
+J
+K
+L
+M
+N
+O
+P
+Q
+R
+S
+T
+U
+V
+W
+X
+Y
+Z
+[
+\
+]
+_
+a
+a1
+ai1
+ai2
+ai3
+ai4
+an1
+an3
+an4
+ang1
+ang2
+ang4
+ao1
+ao2
+ao3
+ao4
+b
+ba
+ba1
+ba2
+ba3
+ba4
+bai1
+bai2
+bai3
+bai4
+ban1
+ban2
+ban3
+ban4
+bang1
+bang2
+bang3
+bang4
+bao1
+bao2
+bao3
+bao4
+bei
+bei1
+bei2
+bei3
+bei4
+ben1
+ben2
+ben3
+ben4
+beng
+beng1
+beng2
+beng3
+beng4
+bi1
+bi2
+bi3
+bi4
+bian1
+bian2
+bian3
+bian4
+biao1
+biao2
+biao3
+bie1
+bie2
+bie3
+bie4
+bin1
+bin4
+bing1
+bing2
+bing3
+bing4
+bo
+bo1
+bo2
+bo3
+bo4
+bu2
+bu3
+bu4
+c
+ca1
+cai1
+cai2
+cai3
+cai4
+can1
+can2
+can3
+can4
+cang1
+cang2
+cao1
+cao2
+cao3
+ce4
+cen1
+cen2
+ceng1
+ceng2
+ceng4
+cha1
+cha2
+cha3
+cha4
+chai1
+chai2
+chan1
+chan2
+chan3
+chan4
+chang1
+chang2
+chang3
+chang4
+chao1
+chao2
+chao3
+che1
+che2
+che3
+che4
+chen1
+chen2
+chen3
+chen4
+cheng1
+cheng2
+cheng3
+cheng4
+chi1
+chi2
+chi3
+chi4
+chong1
+chong2
+chong3
+chong4
+chou1
+chou2
+chou3
+chou4
+chu1
+chu2
+chu3
+chu4
+chua1
+chuai1
+chuai2
+chuai3
+chuai4
+chuan1
+chuan2
+chuan3
+chuan4
+chuang1
+chuang2
+chuang3
+chuang4
+chui1
+chui2
+chun1
+chun2
+chun3
+chuo1
+chuo4
+ci1
+ci2
+ci3
+ci4
+cong1
+cong2
+cou4
+cu1
+cu4
+cuan1
+cuan2
+cuan4
+cui1
+cui3
+cui4
+cun1
+cun2
+cun4
+cuo1
+cuo2
+cuo4
+d
+da
+da1
+da2
+da3
+da4
+dai1
+dai2
+dai3
+dai4
+dan1
+dan2
+dan3
+dan4
+dang1
+dang2
+dang3
+dang4
+dao1
+dao2
+dao3
+dao4
+de
+de1
+de2
+dei3
+den4
+deng1
+deng2
+deng3
+deng4
+di1
+di2
+di3
+di4
+dia3
+dian1
+dian2
+dian3
+dian4
+diao1
+diao3
+diao4
+die1
+die2
+die4
+ding1
+ding2
+ding3
+ding4
+diu1
+dong1
+dong3
+dong4
+dou1
+dou2
+dou3
+dou4
+du1
+du2
+du3
+du4
+duan1
+duan2
+duan3
+duan4
+dui1
+dui4
+dun1
+dun3
+dun4
+duo1
+duo2
+duo3
+duo4
+e
+e1
+e2
+e3
+e4
+ei2
+en1
+en4
+er
+er2
+er3
+er4
+f
+fa1
+fa2
+fa3
+fa4
+fan1
+fan2
+fan3
+fan4
+fang1
+fang2
+fang3
+fang4
+fei1
+fei2
+fei3
+fei4
+fen1
+fen2
+fen3
+fen4
+feng1
+feng2
+feng3
+feng4
+fo2
+fou2
+fou3
+fu1
+fu2
+fu3
+fu4
+g
+ga1
+ga2
+ga3
+ga4
+gai1
+gai2
+gai3
+gai4
+gan1
+gan2
+gan3
+gan4
+gang1
+gang2
+gang3
+gang4
+gao1
+gao2
+gao3
+gao4
+ge1
+ge2
+ge3
+ge4
+gei2
+gei3
+gen1
+gen2
+gen3
+gen4
+geng1
+geng3
+geng4
+gong1
+gong3
+gong4
+gou1
+gou2
+gou3
+gou4
+gu
+gu1
+gu2
+gu3
+gu4
+gua1
+gua2
+gua3
+gua4
+guai1
+guai2
+guai3
+guai4
+guan1
+guan2
+guan3
+guan4
+guang1
+guang2
+guang3
+guang4
+gui1
+gui2
+gui3
+gui4
+gun3
+gun4
+guo1
+guo2
+guo3
+guo4
+h
+ha1
+ha2
+ha3
+hai1
+hai2
+hai3
+hai4
+han1
+han2
+han3
+han4
+hang1
+hang2
+hang4
+hao1
+hao2
+hao3
+hao4
+he1
+he2
+he4
+hei1
+hen2
+hen3
+hen4
+heng1
+heng2
+heng4
+hong1
+hong2
+hong3
+hong4
+hou1
+hou2
+hou3
+hou4
+hu1
+hu2
+hu3
+hu4
+hua1
+hua2
+hua4
+huai2
+huai4
+huan1
+huan2
+huan3
+huan4
+huang1
+huang2
+huang3
+huang4
+hui1
+hui2
+hui3
+hui4
+hun1
+hun2
+hun4
+huo
+huo1
+huo2
+huo3
+huo4
+i
+j
+ji1
+ji2
+ji3
+ji4
+jia
+jia1
+jia2
+jia3
+jia4
+jian1
+jian2
+jian3
+jian4
+jiang1
+jiang2
+jiang3
+jiang4
+jiao1
+jiao2
+jiao3
+jiao4
+jie1
+jie2
+jie3
+jie4
+jin1
+jin2
+jin3
+jin4
+jing1
+jing2
+jing3
+jing4
+jiong3
+jiu1
+jiu2
+jiu3
+jiu4
+ju1
+ju2
+ju3
+ju4
+juan1
+juan2
+juan3
+juan4
+jue1
+jue2
+jue4
+jun1
+jun4
+k
+ka1
+ka2
+ka3
+kai1
+kai2
+kai3
+kai4
+kan1
+kan2
+kan3
+kan4
+kang1
+kang2
+kang4
+kao1
+kao2
+kao3
+kao4
+ke1
+ke2
+ke3
+ke4
+ken3
+keng1
+kong1
+kong3
+kong4
+kou1
+kou2
+kou3
+kou4
+ku1
+ku2
+ku3
+ku4
+kua1
+kua3
+kua4
+kuai3
+kuai4
+kuan1
+kuan2
+kuan3
+kuang1
+kuang2
+kuang4
+kui1
+kui2
+kui3
+kui4
+kun1
+kun3
+kun4
+kuo4
+l
+la
+la1
+la2
+la3
+la4
+lai2
+lai4
+lan2
+lan3
+lan4
+lang1
+lang2
+lang3
+lang4
+lao1
+lao2
+lao3
+lao4
+le
+le1
+le4
+lei
+lei1
+lei2
+lei3
+lei4
+leng1
+leng2
+leng3
+leng4
+li
+li1
+li2
+li3
+li4
+lia3
+lian2
+lian3
+lian4
+liang2
+liang3
+liang4
+liao1
+liao2
+liao3
+liao4
+lie1
+lie2
+lie3
+lie4
+lin1
+lin2
+lin3
+lin4
+ling2
+ling3
+ling4
+liu1
+liu2
+liu3
+liu4
+long1
+long2
+long3
+long4
+lou1
+lou2
+lou3
+lou4
+lu1
+lu2
+lu3
+lu4
+luan2
+luan3
+luan4
+lun1
+lun2
+lun4
+luo1
+luo2
+luo3
+luo4
+lv2
+lv3
+lv4
+lve3
+lve4
+m
+ma
+ma1
+ma2
+ma3
+ma4
+mai2
+mai3
+mai4
+man1
+man2
+man3
+man4
+mang2
+mang3
+mao1
+mao2
+mao3
+mao4
+me
+mei2
+mei3
+mei4
+men
+men1
+men2
+men4
+meng
+meng1
+meng2
+meng3
+meng4
+mi1
+mi2
+mi3
+mi4
+mian2
+mian3
+mian4
+miao1
+miao2
+miao3
+miao4
+mie1
+mie4
+min2
+min3
+ming2
+ming3
+ming4
+miu4
+mo1
+mo2
+mo3
+mo4
+mou1
+mou2
+mou3
+mu2
+mu3
+mu4
+n
+n2
+na1
+na2
+na3
+na4
+nai2
+nai3
+nai4
+nan1
+nan2
+nan3
+nan4
+nang1
+nang2
+nang3
+nao1
+nao2
+nao3
+nao4
+ne
+ne2
+ne4
+nei3
+nei4
+nen4
+neng2
+ni1
+ni2
+ni3
+ni4
+nian1
+nian2
+nian3
+nian4
+niang2
+niang4
+niao2
+niao3
+niao4
+nie1
+nie4
+nin2
+ning2
+ning3
+ning4
+niu1
+niu2
+niu3
+niu4
+nong2
+nong4
+nou4
+nu2
+nu3
+nu4
+nuan3
+nuo2
+nuo4
+nv2
+nv3
+nve4
+o
+o1
+o2
+ou1
+ou2
+ou3
+ou4
+p
+pa1
+pa2
+pa4
+pai1
+pai2
+pai3
+pai4
+pan1
+pan2
+pan4
+pang1
+pang2
+pang4
+pao1
+pao2
+pao3
+pao4
+pei1
+pei2
+pei4
+pen1
+pen2
+pen4
+peng1
+peng2
+peng3
+peng4
+pi1
+pi2
+pi3
+pi4
+pian1
+pian2
+pian4
+piao1
+piao2
+piao3
+piao4
+pie1
+pie2
+pie3
+pin1
+pin2
+pin3
+pin4
+ping1
+ping2
+po1
+po2
+po3
+po4
+pou1
+pu1
+pu2
+pu3
+pu4
+q
+qi1
+qi2
+qi3
+qi4
+qia1
+qia3
+qia4
+qian1
+qian2
+qian3
+qian4
+qiang1
+qiang2
+qiang3
+qiang4
+qiao1
+qiao2
+qiao3
+qiao4
+qie1
+qie2
+qie3
+qie4
+qin1
+qin2
+qin3
+qin4
+qing1
+qing2
+qing3
+qing4
+qiong1
+qiong2
+qiu1
+qiu2
+qiu3
+qu1
+qu2
+qu3
+qu4
+quan1
+quan2
+quan3
+quan4
+que1
+que2
+que4
+qun2
+r
+ran2
+ran3
+rang1
+rang2
+rang3
+rang4
+rao2
+rao3
+rao4
+re2
+re3
+re4
+ren2
+ren3
+ren4
+reng1
+reng2
+ri4
+rong1
+rong2
+rong3
+rou2
+rou4
+ru2
+ru3
+ru4
+ruan2
+ruan3
+rui3
+rui4
+run4
+ruo4
+s
+sa1
+sa2
+sa3
+sa4
+sai1
+sai4
+san1
+san2
+san3
+san4
+sang1
+sang3
+sang4
+sao1
+sao2
+sao3
+sao4
+se4
+sen1
+seng1
+sha1
+sha2
+sha3
+sha4
+shai1
+shai2
+shai3
+shai4
+shan1
+shan3
+shan4
+shang
+shang1
+shang3
+shang4
+shao1
+shao2
+shao3
+shao4
+she1
+she2
+she3
+she4
+shei2
+shen1
+shen2
+shen3
+shen4
+sheng1
+sheng2
+sheng3
+sheng4
+shi
+shi1
+shi2
+shi3
+shi4
+shou1
+shou2
+shou3
+shou4
+shu1
+shu2
+shu3
+shu4
+shua1
+shua2
+shua3
+shua4
+shuai1
+shuai3
+shuai4
+shuan1
+shuan4
+shuang1
+shuang3
+shui2
+shui3
+shui4
+shun3
+shun4
+shuo1
+shuo4
+si1
+si2
+si3
+si4
+song1
+song3
+song4
+sou1
+sou3
+sou4
+su1
+su2
+su4
+suan1
+suan4
+sui1
+sui2
+sui3
+sui4
+sun1
+sun3
+suo
+suo1
+suo2
+suo3
+t
+ta1
+ta2
+ta3
+ta4
+tai1
+tai2
+tai4
+tan1
+tan2
+tan3
+tan4
+tang1
+tang2
+tang3
+tang4
+tao1
+tao2
+tao3
+tao4
+te4
+teng2
+ti1
+ti2
+ti3
+ti4
+tian1
+tian2
+tian3
+tiao1
+tiao2
+tiao3
+tiao4
+tie1
+tie2
+tie3
+tie4
+ting1
+ting2
+ting3
+tong1
+tong2
+tong3
+tong4
+tou
+tou1
+tou2
+tou4
+tu1
+tu2
+tu3
+tu4
+tuan1
+tuan2
+tui1
+tui2
+tui3
+tui4
+tun1
+tun2
+tun4
+tuo1
+tuo2
+tuo3
+tuo4
+u
+v
+w
+wa
+wa1
+wa2
+wa3
+wa4
+wai1
+wai3
+wai4
+wan1
+wan2
+wan3
+wan4
+wang1
+wang2
+wang3
+wang4
+wei1
+wei2
+wei3
+wei4
+wen1
+wen2
+wen3
+wen4
+weng1
+weng4
+wo1
+wo2
+wo3
+wo4
+wu1
+wu2
+wu3
+wu4
+x
+xi1
+xi2
+xi3
+xi4
+xia1
+xia2
+xia4
+xian1
+xian2
+xian3
+xian4
+xiang1
+xiang2
+xiang3
+xiang4
+xiao1
+xiao2
+xiao3
+xiao4
+xie1
+xie2
+xie3
+xie4
+xin1
+xin2
+xin4
+xing1
+xing2
+xing3
+xing4
+xiong1
+xiong2
+xiu1
+xiu3
+xiu4
+xu
+xu1
+xu2
+xu3
+xu4
+xuan1
+xuan2
+xuan3
+xuan4
+xue1
+xue2
+xue3
+xue4
+xun1
+xun2
+xun4
+y
+ya
+ya1
+ya2
+ya3
+ya4
+yan1
+yan2
+yan3
+yan4
+yang1
+yang2
+yang3
+yang4
+yao1
+yao2
+yao3
+yao4
+ye1
+ye2
+ye3
+ye4
+yi
+yi1
+yi2
+yi3
+yi4
+yin1
+yin2
+yin3
+yin4
+ying1
+ying2
+ying3
+ying4
+yo1
+yong1
+yong2
+yong3
+yong4
+you1
+you2
+you3
+you4
+yu1
+yu2
+yu3
+yu4
+yuan1
+yuan2
+yuan3
+yuan4
+yue1
+yue4
+yun1
+yun2
+yun3
+yun4
+z
+za1
+za2
+za3
+zai1
+zai3
+zai4
+zan1
+zan2
+zan3
+zan4
+zang1
+zang4
+zao1
+zao2
+zao3
+zao4
+ze2
+ze4
+zei2
+zen3
+zeng1
+zeng4
+zha1
+zha2
+zha3
+zha4
+zhai1
+zhai2
+zhai3
+zhai4
+zhan1
+zhan2
+zhan3
+zhan4
+zhang1
+zhang2
+zhang3
+zhang4
+zhao1
+zhao2
+zhao3
+zhao4
+zhe
+zhe1
+zhe2
+zhe3
+zhe4
+zhen1
+zhen2
+zhen3
+zhen4
+zheng1
+zheng2
+zheng3
+zheng4
+zhi1
+zhi2
+zhi3
+zhi4
+zhong1
+zhong2
+zhong3
+zhong4
+zhou1
+zhou2
+zhou3
+zhou4
+zhu1
+zhu2
+zhu3
+zhu4
+zhua1
+zhua2
+zhua3
+zhuai1
+zhuai3
+zhuai4
+zhuan1
+zhuan2
+zhuan3
+zhuan4
+zhuang1
+zhuang4
+zhui1
+zhui4
+zhun1
+zhun2
+zhun3
+zhuo1
+zhuo2
+zi
+zi1
+zi2
+zi3
+zi4
+zong1
+zong2
+zong3
+zong4
+zou1
+zou2
+zou3
+zou4
+zu1
+zu2
+zu3
+zuan1
+zuan3
+zuan4
+zui2
+zui3
+zui4
+zun1
+zuo
+zuo1
+zuo2
+zuo3
+zuo4
+{
+~
+¡
+¢
+£
+¥
+§
+¨
+©
+«
+®
+¯
+°
+±
+²
+³
+´
+µ
+·
+¹
+º
+»
+¼
+½
+¾
+¿
+À
+Á
+Â
+Ã
+Ä
+Å
+Æ
+Ç
+È
+É
+Ê
+Í
+Î
+Ñ
+Ó
+Ö
+×
+Ø
+Ú
+Ü
+Ý
+Þ
+ß
+à
+á
+â
+ã
+ä
+å
+æ
+ç
+è
+é
+ê
+ë
+ì
+í
+î
+ï
+ð
+ñ
+ò
+ó
+ô
+õ
+ö
+ø
+ù
+ú
+û
+ü
+ý
+Ā
+ā
+ă
+ą
+ć
+Č
+č
+Đ
+đ
+ē
+ė
+ę
+ě
+ĝ
+ğ
+ħ
+ī
+į
+İ
+ı
+Ł
+ł
+ń
+ņ
+ň
+ŋ
+Ō
+ō
+ő
+œ
+ř
+Ś
+ś
+Ş
+ş
+Š
+š
+Ť
+ť
+ũ
+ū
+ź
+Ż
+ż
+Ž
+ž
+ơ
+ư
+ǎ
+ǐ
+ǒ
+ǔ
+ǚ
+ș
+ț
+ɑ
+ɔ
+ɕ
+ə
+ɛ
+ɜ
+ɡ
+ɣ
+ɪ
+ɫ
+ɴ
+ɹ
+ɾ
+ʃ
+ʊ
+ʌ
+ʒ
+ʔ
+ʰ
+ʷ
+ʻ
+ʾ
+ʿ
+ˈ
+ː
+˙
+˜
+ˢ
+́
+̅
+Α
+Β
+Δ
+Ε
+Θ
+Κ
+Λ
+Μ
+Ξ
+Π
+Σ
+Τ
+Φ
+Χ
+Ψ
+Ω
+ά
+έ
+ή
+ί
+α
+β
+γ
+δ
+ε
+ζ
+η
+θ
+ι
+κ
+λ
+μ
+ν
+ξ
+ο
+π
+ρ
+ς
+σ
+τ
+υ
+φ
+χ
+ψ
+ω
+ϊ
+ό
+ύ
+ώ
+ϕ
+ϵ
+Ё
+А
+Б
+В
+Г
+Д
+Е
+Ж
+З
+И
+Й
+К
+Л
+М
+Н
+О
+П
+Р
+С
+Т
+У
+Ф
+Х
+Ц
+Ч
+Ш
+Щ
+Ы
+Ь
+Э
+Ю
+Я
+а
+б
+в
+г
+д
+е
+ж
+з
+и
+й
+к
+л
+м
+н
+о
+п
+р
+с
+т
+у
+ф
+х
+ц
+ч
+ш
+щ
+ъ
+ы
+ь
+э
+ю
+я
+ё
+і
+ְ
+ִ
+ֵ
+ֶ
+ַ
+ָ
+ֹ
+ּ
+־
+ׁ
+א
+ב
+ג
+ד
+ה
+ו
+ז
+ח
+ט
+י
+כ
+ל
+ם
+מ
+ן
+נ
+ס
+ע
+פ
+ק
+ר
+ש
+ת
+أ
+ب
+ة
+ت
+ج
+ح
+د
+ر
+ز
+س
+ص
+ط
+ع
+ق
+ك
+ل
+م
+ن
+ه
+و
+ي
+َ
+ُ
+ِ
+ْ
+ก
+ข
+ง
+จ
+ต
+ท
+น
+ป
+ย
+ร
+ว
+ส
+ห
+อ
+ฮ
+ั
+า
+ี
+ึ
+โ
+ใ
+ไ
+่
+้
+์
+ḍ
+Ḥ
+ḥ
+ṁ
+ṃ
+ṅ
+ṇ
+Ṛ
+ṛ
+Ṣ
+ṣ
+Ṭ
+ṭ
+ạ
+ả
+Ấ
+ấ
+ầ
+ậ
+ắ
+ằ
+ẻ
+ẽ
+ế
+ề
+ể
+ễ
+ệ
+ị
+ọ
+ỏ
+ố
+ồ
+ộ
+ớ
+ờ
+ở
+ụ
+ủ
+ứ
+ữ
+ἀ
+ἁ
+Ἀ
+ἐ
+ἔ
+ἰ
+ἱ
+ὀ
+ὁ
+ὐ
+ὲ
+ὸ
+���
+᾽
+ῆ
+ῇ
+ῶ
+‎
+‑
+‒
+–
+—
+―
+‖
+†
+‡
+•
+…
+‧
+‬
+′
+″
+⁄
+⁡
+⁰
+⁴
+⁵
+⁶
+⁷
+⁸
+⁹
+₁
+₂
+₃
+€
+₱
+₹
+₽
+℃
+ℏ
+ℓ
+№
+ℝ
+™
+⅓
+⅔
+⅛
+→
+∂
+∈
+∑
+−
+∗
+√
+∞
+∫
+≈
+≠
+≡
+≤
+≥
+⋅
+⋯
+█
+♪
+⟨
+⟩
+、
+。
+《
+》
+「
+」
+【
+】
+あ
+う
+え
+お
+か
+が
+き
+ぎ
+く
+ぐ
+け
+げ
+こ
+ご
+さ
+し
+じ
+す
+ず
+せ
+ぜ
+そ
+ぞ
+た
+だ
+ち
+っ
+つ
+で
+と
+ど
+な
+に
+ね
+の
+は
+ば
+ひ
+ぶ
+へ
+べ
+ま
+み
+む
+め
+も
+ゃ
+や
+ゆ
+ょ
+よ
+ら
+り
+る
+れ
+ろ
+わ
+を
+ん
+ァ
+ア
+ィ
+イ
+ウ
+ェ
+エ
+オ
+カ
+ガ
+キ
+ク
+ケ
+ゲ
+コ
+ゴ
+サ
+ザ
+シ
+ジ
+ス
+ズ
+セ
+ゾ
+タ
+ダ
+チ
+ッ
+ツ
+テ
+デ
+ト
+ド
+ナ
+ニ
+ネ
+ノ
+バ
+パ
+ビ
+ピ
+フ
+プ
+ヘ
+ベ
+ペ
+ホ
+ボ
+ポ
+マ
+ミ
+ム
+メ
+モ
+ャ
+ヤ
+ュ
+ユ
+ョ
+ヨ
+ラ
+リ
+ル
+レ
+ロ
+ワ
+ン
+・
+ー
+ㄋ
+ㄍ
+ㄎ
+ㄏ
+ㄓ
+ㄕ
+ㄚ
+ㄜ
+ㄟ
+ㄤ
+ㄥ
+ㄧ
+ㄱ
+ㄴ
+ㄷ
+ㄹ
+ㅁ
+ㅂ
+ㅅ
+ㅈ
+ㅍ
+ㅎ
+ㅏ
+ㅓ
+ㅗ
+ㅜ
+ㅡ
+ㅣ
+㗎
+가
+각
+간
+갈
+감
+갑
+갓
+갔
+강
+같
+개
+거
+건
+걸
+겁
+것
+겉
+게
+겠
+겨
+결
+겼
+경
+계
+고
+곤
+골
+곱
+공
+과
+관
+광
+교
+구
+국
+굴
+귀
+귄
+그
+근
+글
+금
+기
+긴
+길
+까
+깍
+깔
+깜
+깨
+께
+꼬
+꼭
+꽃
+꾸
+꿔
+끔
+끗
+끝
+끼
+나
+난
+날
+남
+납
+내
+냐
+냥
+너
+넘
+넣
+네
+녁
+년
+녕
+노
+녹
+놀
+누
+눈
+느
+는
+늘
+니
+님
+닙
+다
+닥
+단
+달
+닭
+당
+대
+더
+덕
+던
+덥
+데
+도
+독
+동
+돼
+됐
+되
+된
+될
+두
+둑
+둥
+드
+들
+등
+디
+따
+딱
+딸
+땅
+때
+떤
+떨
+떻
+또
+똑
+뚱
+뛰
+뜻
+띠
+라
+락
+란
+람
+랍
+랑
+래
+랜
+러
+런
+럼
+렇
+레
+려
+력
+렵
+렸
+로
+록
+롬
+루
+르
+른
+를
+름
+릉
+리
+릴
+림
+마
+막
+만
+많
+말
+맑
+맙
+맛
+매
+머
+먹
+멍
+메
+면
+명
+몇
+모
+목
+몸
+못
+무
+문
+물
+뭐
+뭘
+미
+민
+밌
+밑
+바
+박
+밖
+반
+받
+발
+밤
+밥
+방
+배
+백
+밸
+뱀
+버
+번
+벌
+벚
+베
+벼
+벽
+별
+병
+보
+복
+본
+볼
+봐
+봤
+부
+분
+불
+비
+빔
+빛
+빠
+빨
+뼈
+뽀
+뿅
+쁘
+사
+산
+살
+삼
+샀
+상
+새
+색
+생
+서
+선
+설
+섭
+섰
+성
+세
+셔
+션
+셨
+소
+속
+손
+송
+수
+숙
+순
+술
+숫
+숭
+숲
+쉬
+쉽
+스
+슨
+습
+슷
+시
+식
+신
+실
+싫
+심
+십
+싶
+싸
+써
+쓰
+쓴
+씌
+씨
+씩
+씬
+아
+악
+안
+않
+알
+야
+약
+얀
+양
+얘
+어
+언
+얼
+엄
+업
+없
+었
+엉
+에
+여
+역
+연
+염
+엽
+영
+옆
+예
+옛
+오
+온
+올
+옷
+옹
+와
+왔
+왜
+요
+욕
+용
+우
+운
+울
+웃
+워
+원
+월
+웠
+위
+윙
+유
+육
+윤
+으
+은
+을
+음
+응
+의
+이
+익
+인
+일
+읽
+임
+입
+있
+자
+작
+잔
+잖
+잘
+잡
+잤
+장
+재
+저
+전
+점
+정
+제
+져
+졌
+조
+족
+좀
+종
+좋
+죠
+주
+준
+줄
+중
+줘
+즈
+즐
+즘
+지
+진
+집
+짜
+짝
+쩌
+쪼
+쪽
+쫌
+쭈
+쯔
+찌
+찍
+차
+착
+찾
+책
+처
+천
+철
+체
+쳐
+쳤
+초
+촌
+추
+출
+춤
+춥
+춰
+치
+친
+칠
+침
+칩
+칼
+커
+켓
+코
+콩
+쿠
+퀴
+크
+큰
+큽
+키
+킨
+타
+태
+터
+턴
+털
+테
+토
+통
+투
+트
+특
+튼
+틀
+티
+팀
+파
+팔
+패
+페
+펜
+펭
+평
+포
+폭
+표
+품
+풍
+프
+플
+피
+필
+하
+학
+한
+할
+함
+합
+항
+해
+햇
+했
+행
+허
+험
+형
+혜
+호
+혼
+홀
+화
+회
+획
+후
+휴
+흐
+흔
+희
+히
+힘
+ﷺ
+ﷻ
+！
+，
+？
+�
+𠮶

model/__init__.py

@@ -0,0 +1,7 @@
+from model.cfm import CFM
+
+from model.backbones.unett import UNetT
+from model.backbones.dit import DiT
+from model.backbones.mmdit import MMDiT
+
+from model.trainer import Trainer

model/backbones/README.md

@@ -0,0 +1,20 @@
+## 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

model/backbones/dit.py

@@ -0,0 +1,158 @@
+"""
+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 einops import repeat
+
+from x_transformers.x_transformers import RotaryEmbedding
+
+from 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):
+        batch, text_len = text.shape[0], text.shape[1]
+        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
+        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):
+        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
+        cond: float['b n d'],  # masked cond audio
+        text: int['b nt'],     # text
+        time: float['b'] | float[''],  # time step
+        drop_audio_cond,  # cfg for cond audio
+        drop_text,        # cfg for text
+        mask: bool['b n'] | None = None,
+    ):
+        batch, seq_len = x.shape[0], x.shape[1]
+        if time.ndim == 0:
+            time = repeat(time, ' -> b', b = 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

model/backbones/mmdit.py

@@ -0,0 +1,136 @@
+"""
+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 einops import repeat
+
+from x_transformers.x_transformers import RotaryEmbedding
+
+from 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']:
+        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):
+        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
+        cond: float['b n d'],  # masked cond audio
+        text: int['b nt'],     # text
+        time: float['b'] | float[''],  # time step
+        drop_audio_cond,  # cfg for cond audio
+        drop_text,        # cfg for text
+        mask: bool['b n'] | None = None,
+    ):
+        batch = x.shape[0]
+        if time.ndim == 0:
+            time = repeat(time, ' -> b', b = 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

model/backbones/unett.py

@@ -0,0 +1,201 @@
+"""
+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 einops import repeat, pack, unpack
+
+from x_transformers import RMSNorm
+from x_transformers.x_transformers import RotaryEmbedding
+
+from 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):
+        batch, text_len = text.shape[0], text.shape[1]
+        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
+        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):
+        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
+        cond: float['b n d'],  # masked cond audio
+        text: int['b nt'],     # text
+        time: float['b'] | float[''],  # time step
+        drop_audio_cond,  # cfg for cond audio
+        drop_text,        # cfg for text
+        mask: bool['b n'] | None = None,
+    ):
+        batch, seq_len = x.shape[0], x.shape[1]
+        if time.ndim == 0:
+            time = repeat(time, ' -> b', b = 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, ps = pack((t, x), 'b * d')
+        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 = unpack(self.norm_out(x), ps, 'b * d')
+
+        return self.proj_out(x)

model/cfm.py

@@ -0,0 +1,273 @@
+"""
+ein notation:
+b - batch
+n - sequence
+nt - text sequence
+nw - raw wave length
+d - dimension
+"""
+
+from __future__ import annotations
+from typing import Callable
+from random import random
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+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, 
+    lens_to_mask, mask_from_frac_lengths,
+) 
+
+
+class CFM(nn.Module):
+    def __init__(
+        self,
+        transformer: nn.Module,
+        sigma = 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.),
+        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'],
+        text: int['b nt'] | list[str],
+        duration: int | int['b'],
+        *,
+        lens: int['b'] | None = None,
+        steps = 32,
+        cfg_strength = 1., 
+        sway_sampling_coef = None,
+        seed: int | None = None,
+        max_duration = 4096, 
+        vocoder: Callable[[float['b d n']], float['b nw']] | None = None,
+        no_ref_audio = False,
+        duplicate_test = False,
+        t_inter = 0.1,
+    ):
+        self.eval()
+
+        # raw wave
+
+        if cond.ndim == 2:
+            cond = self.mel_spec(cond)
+            cond = rearrange(cond, 'b d n -> b n d')
+            assert cond.shape[-1] == self.num_channels
+
+        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 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.)
+            
+        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')
+        step_cond = torch.where(cond_mask, cond, torch.zeros_like(cond))  # allow direct control (cut cond audio) with lens passed in
+
+        mask = lens_to_mask(duration)
+
+        # 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))
+        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)
+        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 = rearrange(out, 'b n d -> b d n')
+            out = vocoder(out)
+
+        return out, trajectory
+
+    def forward(
+        self,
+        inp: float['b n d'] | float['b nw'], # mel or raw wave
+        text: int['b nt'] | list[str],
+        *,
+        lens: int['b'] | None = None,
+        noise_scheduler: str | None = None,
+    ):
+        # handle raw wave
+        if inp.ndim == 2:
+            inp = self.mel_spec(inp)
+            inp = rearrange(inp, 'b d n -> b n d')
+            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 = rearrange(time, 'b -> b 1 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

model/dataset.py

@@ -0,0 +1,242 @@
+import json
+import random
+from tqdm import tqdm
+
+import torch
+import torch.nn.functional as F
+from torch.utils.data import Dataset, Sampler
+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
+
+
+class HFDataset(Dataset):
+    def __init__(
+        self,
+        hf_dataset: Dataset,
+        target_sample_rate = 24_000,
+        n_mel_channels = 100,
+        hop_length = 256,
+    ):
+        self.data = hf_dataset
+        self.target_sample_rate = target_sample_rate
+        self.hop_length = hop_length
+        self.mel_spectrogram = MelSpec(target_sample_rate=target_sample_rate, n_mel_channels=n_mel_channels, hop_length=hop_length)
+        
+    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 = rearrange(audio_tensor, 't -> 1 t')
+        
+        mel_spec = self.mel_spectrogram(audio_tensor)
+        
+        mel_spec = rearrange(mel_spec, '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,
+        preprocessed_mel = False,
+    ):
+        self.data = custom_dataset
+        self.durations = durations
+        self.target_sample_rate = target_sample_rate
+        self.hop_length = hop_length
+        self.preprocessed_mel = preprocessed_mel
+        if not preprocessed_mel:
+            self.mel_spectrogram = MelSpec(target_sample_rate=target_sample_rate, hop_length=hop_length, n_mel_channels=n_mel_channels)
+
+    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 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 = rearrange(mel_spec, '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=f"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,
+        dataset_type: str = "CustomDataset", 
+        audio_type: str = "raw", 
+        mel_spec_kwargs: dict = dict()
+        ) -> CustomDataset | HFDataset:
+    
+    print("Loading dataset ...")
+
+    if dataset_type == "CustomDataset":
+        if audio_type == "raw":
+            try:
+                train_dataset = load_from_disk(f"data/{dataset_name}_{tokenizer}/raw")
+            except:
+                train_dataset = Dataset_.from_file(f"data/{dataset_name}_{tokenizer}/raw.arrow")
+            preprocessed_mel = False
+        elif audio_type == "mel":
+            train_dataset = Dataset_.from_file(f"data/{dataset_name}_{tokenizer}/mel.arrow")
+            preprocessed_mel = True
+        with open(f"data/{dataset_name}_{tokenizer}/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="./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,
+    )

model/ecapa_tdnn.py

@@ -0,0 +1,268 @@
+# just for speaker similarity evaluation, third-party code
+
+# From https://github.com/microsoft/UniSpeech/blob/main/downstreams/speaker_verification/models/
+# part of the code is borrowed from https://github.com/lawlict/ECAPA-TDNN
+
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+''' Res2Conv1d + BatchNorm1d + ReLU
+'''
+
+class Res2Conv1dReluBn(nn.Module):
+    '''
+    in_channels == out_channels == channels
+    '''
+
+    def __init__(self, channels, kernel_size=1, stride=1, padding=0, dilation=1, bias=True, scale=4):
+        super().__init__()
+        assert channels % scale == 0, "{} % {} != 0".format(channels, scale)
+        self.scale = scale
+        self.width = channels // scale
+        self.nums = scale if scale == 1 else scale - 1
+
+        self.convs = []
+        self.bns = []
+        for i in range(self.nums):
+            self.convs.append(nn.Conv1d(self.width, self.width, kernel_size, stride, padding, dilation, bias=bias))
+            self.bns.append(nn.BatchNorm1d(self.width))
+        self.convs = nn.ModuleList(self.convs)
+        self.bns = nn.ModuleList(self.bns)
+
+    def forward(self, x):
+        out = []
+        spx = torch.split(x, self.width, 1)
+        for i in range(self.nums):
+            if i == 0:
+                sp = spx[i]
+            else:
+                sp = sp + spx[i]
+            # Order: conv -> relu -> bn
+            sp = self.convs[i](sp)
+            sp = self.bns[i](F.relu(sp))
+            out.append(sp)
+        if self.scale != 1:
+            out.append(spx[self.nums])
+        out = torch.cat(out, dim=1)
+
+        return out
+
+
+''' Conv1d + BatchNorm1d + ReLU
+'''
+
+class Conv1dReluBn(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size=1, stride=1, padding=0, dilation=1, bias=True):
+        super().__init__()
+        self.conv = nn.Conv1d(in_channels, out_channels, kernel_size, stride, padding, dilation, bias=bias)
+        self.bn = nn.BatchNorm1d(out_channels)
+
+    def forward(self, x):
+        return self.bn(F.relu(self.conv(x)))
+
+
+''' The SE connection of 1D case.
+'''
+
+class SE_Connect(nn.Module):
+    def __init__(self, channels, se_bottleneck_dim=128):
+        super().__init__()
+        self.linear1 = nn.Linear(channels, se_bottleneck_dim)
+        self.linear2 = nn.Linear(se_bottleneck_dim, channels)
+
+    def forward(self, x):
+        out = x.mean(dim=2)
+        out = F.relu(self.linear1(out))
+        out = torch.sigmoid(self.linear2(out))
+        out = x * out.unsqueeze(2)
+
+        return out
+
+
+''' SE-Res2Block of the ECAPA-TDNN architecture.
+'''
+
+# def SE_Res2Block(channels, kernel_size, stride, padding, dilation, scale):
+#     return nn.Sequential(
+#         Conv1dReluBn(channels, 512, kernel_size=1, stride=1, padding=0),
+#         Res2Conv1dReluBn(512, kernel_size, stride, padding, dilation, scale=scale),
+#         Conv1dReluBn(512, channels, kernel_size=1, stride=1, padding=0),
+#         SE_Connect(channels)
+#     )
+
+class SE_Res2Block(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, dilation, scale, se_bottleneck_dim):
+        super().__init__()
+        self.Conv1dReluBn1 = Conv1dReluBn(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
+        self.Res2Conv1dReluBn = Res2Conv1dReluBn(out_channels, kernel_size, stride, padding, dilation, scale=scale)
+        self.Conv1dReluBn2 = Conv1dReluBn(out_channels, out_channels, kernel_size=1, stride=1, padding=0)
+        self.SE_Connect = SE_Connect(out_channels, se_bottleneck_dim)
+
+        self.shortcut = None
+        if in_channels != out_channels:
+            self.shortcut = nn.Conv1d(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                kernel_size=1,
+            )
+
+    def forward(self, x):
+        residual = x
+        if self.shortcut:
+            residual = self.shortcut(x)
+
+        x = self.Conv1dReluBn1(x)
+        x = self.Res2Conv1dReluBn(x)
+        x = self.Conv1dReluBn2(x)
+        x = self.SE_Connect(x)
+
+        return x + residual
+
+
+''' Attentive weighted mean and standard deviation pooling.
+'''
+
+class AttentiveStatsPool(nn.Module):
+    def __init__(self, in_dim, attention_channels=128, global_context_att=False):
+        super().__init__()
+        self.global_context_att = global_context_att
+
+        # Use Conv1d with stride == 1 rather than Linear, then we don't need to transpose inputs.
+        if global_context_att:
+            self.linear1 = nn.Conv1d(in_dim * 3, attention_channels, kernel_size=1)  # equals W and b in the paper
+        else:
+            self.linear1 = nn.Conv1d(in_dim, attention_channels, kernel_size=1)  # equals W and b in the paper
+        self.linear2 = nn.Conv1d(attention_channels, in_dim, kernel_size=1)  # equals V and k in the paper
+
+    def forward(self, x):
+
+        if self.global_context_att:
+            context_mean = torch.mean(x, dim=-1, keepdim=True).expand_as(x)
+            context_std = torch.sqrt(torch.var(x, dim=-1, keepdim=True) + 1e-10).expand_as(x)
+            x_in = torch.cat((x, context_mean, context_std), dim=1)
+        else:
+            x_in = x
+
+        # DON'T use ReLU here! In experiments, I find ReLU hard to converge.
+        alpha = torch.tanh(self.linear1(x_in))
+        # alpha = F.relu(self.linear1(x_in))
+        alpha = torch.softmax(self.linear2(alpha), dim=2)
+        mean = torch.sum(alpha * x, dim=2)
+        residuals = torch.sum(alpha * (x ** 2), dim=2) - mean ** 2
+        std = torch.sqrt(residuals.clamp(min=1e-9))
+        return torch.cat([mean, std], dim=1)
+
+
+class ECAPA_TDNN(nn.Module):
+    def __init__(self, feat_dim=80, channels=512, emb_dim=192, global_context_att=False,
+                 feat_type='wavlm_large', sr=16000, feature_selection="hidden_states", update_extract=False, config_path=None):
+        super().__init__()
+
+        self.feat_type = feat_type
+        self.feature_selection = feature_selection
+        self.update_extract = update_extract
+        self.sr = sr
+        
+        torch.hub._validate_not_a_forked_repo=lambda a,b,c: True
+        try:
+            local_s3prl_path = os.path.expanduser("~/.cache/torch/hub/s3prl_s3prl_main")
+            self.feature_extract = torch.hub.load(local_s3prl_path, feat_type, source='local', config_path=config_path)
+        except:
+            self.feature_extract = torch.hub.load('s3prl/s3prl', feat_type)
+
+        if len(self.feature_extract.model.encoder.layers) == 24 and hasattr(self.feature_extract.model.encoder.layers[23].self_attn, "fp32_attention"):
+            self.feature_extract.model.encoder.layers[23].self_attn.fp32_attention = False
+        if len(self.feature_extract.model.encoder.layers) == 24 and hasattr(self.feature_extract.model.encoder.layers[11].self_attn, "fp32_attention"):
+            self.feature_extract.model.encoder.layers[11].self_attn.fp32_attention = False
+
+        self.feat_num = self.get_feat_num()
+        self.feature_weight = nn.Parameter(torch.zeros(self.feat_num))
+
+        if feat_type != 'fbank' and feat_type != 'mfcc':
+            freeze_list = ['final_proj', 'label_embs_concat', 'mask_emb', 'project_q', 'quantizer']
+            for name, param in self.feature_extract.named_parameters():
+                for freeze_val in freeze_list:
+                    if freeze_val in name:
+                        param.requires_grad = False
+                        break
+
+        if not self.update_extract:
+            for param in self.feature_extract.parameters():
+                param.requires_grad = False
+
+        self.instance_norm = nn.InstanceNorm1d(feat_dim)
+        # self.channels = [channels] * 4 + [channels * 3]
+        self.channels = [channels] * 4 + [1536]
+
+        self.layer1 = Conv1dReluBn(feat_dim, self.channels[0], kernel_size=5, padding=2)
+        self.layer2 = SE_Res2Block(self.channels[0], self.channels[1], kernel_size=3, stride=1, padding=2, dilation=2, scale=8, se_bottleneck_dim=128)
+        self.layer3 = SE_Res2Block(self.channels[1], self.channels[2], kernel_size=3, stride=1, padding=3, dilation=3, scale=8, se_bottleneck_dim=128)
+        self.layer4 = SE_Res2Block(self.channels[2], self.channels[3], kernel_size=3, stride=1, padding=4, dilation=4, scale=8, se_bottleneck_dim=128)
+
+        # self.conv = nn.Conv1d(self.channels[-1], self.channels[-1], kernel_size=1)
+        cat_channels = channels * 3
+        self.conv = nn.Conv1d(cat_channels, self.channels[-1], kernel_size=1)
+        self.pooling = AttentiveStatsPool(self.channels[-1], attention_channels=128, global_context_att=global_context_att)
+        self.bn = nn.BatchNorm1d(self.channels[-1] * 2)
+        self.linear = nn.Linear(self.channels[-1] * 2, emb_dim)
+
+
+    def get_feat_num(self):
+        self.feature_extract.eval()
+        wav = [torch.randn(self.sr).to(next(self.feature_extract.parameters()).device)]
+        with torch.no_grad():
+            features = self.feature_extract(wav)
+        select_feature = features[self.feature_selection]
+        if isinstance(select_feature, (list, tuple)):
+            return len(select_feature)
+        else:
+            return 1
+
+    def get_feat(self, x):
+        if self.update_extract:
+            x = self.feature_extract([sample for sample in x])
+        else:
+            with torch.no_grad():
+                if self.feat_type == 'fbank' or self.feat_type == 'mfcc':
+                    x = self.feature_extract(x) + 1e-6  # B x feat_dim x time_len
+                else:
+                    x = self.feature_extract([sample for sample in x])
+
+        if self.feat_type == 'fbank':
+            x = x.log()
+
+        if self.feat_type != "fbank" and self.feat_type != "mfcc":
+            x = x[self.feature_selection]
+            if isinstance(x, (list, tuple)):
+                x = torch.stack(x, dim=0)
+            else:
+                x = x.unsqueeze(0)
+            norm_weights = F.softmax(self.feature_weight, dim=-1).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
+            x = (norm_weights * x).sum(dim=0)
+            x = torch.transpose(x, 1, 2) + 1e-6
+
+        x = self.instance_norm(x)
+        return x
+
+    def forward(self, x):
+        x = self.get_feat(x)
+
+        out1 = self.layer1(x)
+        out2 = self.layer2(out1)
+        out3 = self.layer3(out2)
+        out4 = self.layer4(out3)
+
+        out = torch.cat([out2, out3, out4], dim=1)
+        out = F.relu(self.conv(out))
+        out = self.bn(self.pooling(out))
+        out = self.linear(out)
+
+        return out
+
+
+def ECAPA_TDNN_SMALL(feat_dim, emb_dim=256, feat_type='wavlm_large', sr=16000, feature_selection="hidden_states", update_extract=False, config_path=None):
+    return ECAPA_TDNN(feat_dim=feat_dim, channels=512, emb_dim=emb_dim,
+                      feat_type=feat_type, sr=sr, feature_selection=feature_selection, update_extract=update_extract, config_path=config_path)

model/modules.py

@@ -0,0 +1,575 @@
+"""
+ein notation:
+b - batch
+n - sequence
+nt - text sequence
+nw - raw wave length
+d - dimension
+"""
+
+from __future__ import annotations
+from typing import Optional
+import math
+
+import torch
+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
+
+
+# raw wav to mel spec
+
+class MelSpec(nn.Module):
+    def __init__(
+        self,
+        filter_length = 1024,
+        hop_length = 256,
+        win_length = 1024,
+        n_mel_channels = 100,
+        target_sample_rate = 24_000,
+        normalize = False,
+        power = 1,
+        norm = None,
+        center = True,
+    ):
+        super().__init__()
+        self.n_mel_channels = n_mel_channels
+
+        self.mel_stft = torchaudio.transforms.MelSpectrogram(
+            sample_rate = target_sample_rate,
+            n_fft = filter_length,
+            win_length = win_length,
+            hop_length = hop_length,
+            n_mels = n_mel_channels,
+            power = power,
+            center = center,
+            normalized = normalize,
+            norm = norm,
+        )
+
+        self.register_buffer('dummy', torch.tensor(0), persistent = False)
+
+    def forward(self, inp):
+        if len(inp.shape) == 3:
+            inp = rearrange(inp, 'b 1 nw -> b nw')
+
+        assert len(inp.shape) == 2
+
+        if self.dummy.device != inp.device:
+            self.to(inp.device)
+
+        mel = self.mel_stft(inp)
+        mel = mel.clamp(min = 1e-5).log()
+        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):
+        if mask is not None:
+            mask = mask[..., None]
+            x = x.masked_fill(~mask, 0.)
+
+        x = rearrange(x, 'b n d -> b d n')
+        x = self.conv1d(x)
+        out = rearrange(x, 'b d n -> b n d')
+
+        if mask is not None:
+            out = out.masked_fill(~mask, 0.)
+
+        return out
+
+
+# rotary positional embedding related
+
+def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0, theta_rescale_factor=1.):
+    # 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.):
+    # 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., 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
+        c: float['b n d'] = None,  # context c
+        mask: bool['b n'] | None = None,
+        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
+        mask: bool['b n'] | None = None,
+        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.) if xpos_scale is not None else (1., 1.)
+
+            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 = rearrange(attn_mask, '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 = rearrange(mask, 'b n -> b n 1')
+            x = x.masked_fill(~mask, 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
+        c: float['b nt d'] = None,  # context c, here text
+        mask: bool['b n'] | None = None,
+        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.) if xpos_scale is not None else (1., 1.)
+            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.) if xpos_scale is not None else (1., 1.)
+            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 = rearrange(attn_mask, '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 = rearrange(mask, 'b n -> b n 1')
+            x = x.masked_fill(~mask, 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']):
+        time_hidden = self.time_embed(timestep)
+        time = self.time_mlp(time_hidden)  # b d
+        return time

model/trainer.py

@@ -0,0 +1,245 @@
+from __future__ import annotations
+
+import os
+import gc
+from tqdm import tqdm
+import wandb
+
+import torch
+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
+
+from ema_pytorch import EMA
+
+from model import CFM
+from model.utils import exists, default
+from model.dataset import DynamicBatchSampler, collate_fn
+
+
+# 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,
+        wandb_project = "test_e2-tts",
+        wandb_run_name = "test_run",
+        wandb_resume_id: str = None,
+        last_per_steps = None,
+        accelerate_kwargs: dict = dict(),
+        ema_kwargs: dict = dict()
+    ):
+        
+        ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters = True)
+
+        self.accelerator = Accelerator(
+            log_with = "wandb",
+            kwargs_handlers = [ddp_kwargs],
+            gradient_accumulation_steps = grad_accumulation_steps,
+            **accelerate_kwargs
+        )
+        
+        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}
+            )
+
+        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.noise_scheduler = noise_scheduler
+
+        self.duration_predictor = duration_predictor
+
+        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 == True:
+                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(os.listdir(self.checkpoint_path), 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}", map_location="cpu")
+        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.is_main:
+            self.ema_model.load_state_dict(checkpoint['ema_model_state_dict'])
+
+        if self.scheduler:
+            self.scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
+        
+        step = checkpoint['step']
+        del checkpoint; gc.collect()
+        return step
+
+    def train(self, train_dataset: Dataset, num_workers=16, resumable_with_seed: int = None):
+        
+        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, 
+                                          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,
+                                          batch_sampler=batch_sampler)
+        else:
+            raise ValueError(f"batch_size_type must be either 'sample' or 'frame', but recieved {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 = rearrange(batch['mel'], 'b d n -> b n d')
+                    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)
+                
+                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 global_step % self.last_per_steps == 0:
+                    self.save_checkpoint(global_step, last=True)
+        
+        self.accelerator.end_training()

model/utils.py

@@ -0,0 +1,545 @@
+from __future__ import annotations
+
+import os
+import re
+import math
+import random
+import string
+from tqdm import tqdm
+from collections import defaultdict
+
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pylab as plt
+
+import torch
+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
+import zhconv
+from zhon.hanzi import punctuation
+from jiwer import compute_measures
+
+from funasr import AutoModel
+from faster_whisper import WhisperModel
+
+from model.ecapa_tdnn import ECAPA_TDNN_SMALL
+from model.modules import MelSpec
+
+
+# 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']:
+
+    if not exists(length):
+        length = t.amax()
+
+    seq = torch.arange(length, device = t.device)
+    return einx.less('n, b -> b n', seq, t)
+
+def mask_from_start_end_indices(
+    seq_len: int['b'],
+    start: int['b'],
+    end: int['b']
+):
+    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)
+
+def mask_from_frac_lengths(
+    seq_len: int['b'],
+    frac_lengths: float['b']
+):
+    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']:
+
+    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')
+
+    return einx.divide('b d, b -> b d', num, den.clamp(min = 1.))
+
+
+# simple utf-8 tokenizer, since paper went character based
+def list_str_to_tensor(
+    text: list[str],
+    padding_value = -1
+) -> int['b nt']:
+    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']:
+    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
+    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"]:
+        with open (f"data/{dataset_name}_{tokenizer}/vocab.txt", "r") 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
+
+    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
+    for text in text_list:
+        char_list = []
+        text = text.translate(god_knows_why_en_testset_contains_zh_quote)
+        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
+
+
+# save spectrogram
+def save_spectrogram(spectrogram, path):
+    plt.figure(figsize=(12, 4))
+    plt.imshow(spectrogram, origin='lower', aspect='auto')
+    plt.colorbar()
+    plt.savefig(path)
+    plt.close()
+
+
+# seedtts testset metainfo: utt, prompt_text, prompt_wav, gt_text, gt_wav
+def get_seedtts_testset_metainfo(metalst):
+    f = open(metalst); lines = f.readlines(); f.close()
+    metainfo = []
+    for line in lines:
+        if len(line.strip().split('|')) == 5:
+            utt, prompt_text, prompt_wav, gt_text, gt_wav = line.strip().split('|')
+        elif len(line.strip().split('|')) == 4:
+            utt, prompt_text, prompt_wav, gt_text = line.strip().split('|')
+            gt_wav = os.path.join(os.path.dirname(metalst), "wavs", utt + ".wav")
+        if not os.path.isabs(prompt_wav):
+            prompt_wav = os.path.join(os.path.dirname(metalst), prompt_wav)
+        metainfo.append((utt, prompt_text, prompt_wav, gt_text, gt_wav))
+    return metainfo
+
+
+# librispeech test-clean metainfo: gen_utt, ref_txt, ref_wav, gen_txt, gen_wav
+def get_librispeech_test_clean_metainfo(metalst, librispeech_test_clean_path):
+    f = open(metalst); lines = f.readlines(); f.close()
+    metainfo = []
+    for line in lines:
+        ref_utt, ref_dur, ref_txt, gen_utt, gen_dur, gen_txt = line.strip().split('\t')
+
+        # ref_txt = ref_txt[0] + ref_txt[1:].lower() + '.'  # if use librispeech test-clean (no-pc)
+        ref_spk_id, ref_chaptr_id, _ =  ref_utt.split('-')
+        ref_wav = os.path.join(librispeech_test_clean_path, ref_spk_id, ref_chaptr_id, ref_utt + '.flac')
+
+        # gen_txt = gen_txt[0] + gen_txt[1:].lower() + '.'  # if use librispeech test-clean (no-pc)
+        gen_spk_id, gen_chaptr_id, _ =  gen_utt.split('-')
+        gen_wav = os.path.join(librispeech_test_clean_path, gen_spk_id, gen_chaptr_id, gen_utt + '.flac')
+
+        metainfo.append((gen_utt, ref_txt, ref_wav, " " + gen_txt, gen_wav))
+
+    return metainfo
+
+
+# padded to max length mel batch
+def padded_mel_batch(ref_mels):
+    max_mel_length = torch.LongTensor([mel.shape[-1] for mel in ref_mels]).amax()
+    padded_ref_mels = []
+    for mel in 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')
+    return padded_ref_mels
+
+
+# get prompts from metainfo containing: utt, prompt_text, prompt_wav, gt_text, gt_wav
+
+def get_inference_prompt(
+    metainfo, 
+    speed = 1., tokenizer = "pinyin", polyphone = True, 
+    target_sample_rate = 24000, n_mel_channels = 100, hop_length = 256, target_rms = 0.1,
+    use_truth_duration = False,
+    infer_batch_size = 1, num_buckets = 200, min_secs = 3, max_secs = 40,
+):
+    prompts_all = []
+
+    min_tokens = min_secs * target_sample_rate // hop_length
+    max_tokens = max_secs * target_sample_rate // hop_length
+
+    batch_accum = [0] * num_buckets
+    utts, ref_rms_list, ref_mels, ref_mel_lens, total_mel_lens, final_text_list = \
+        ([[] for _ in range(num_buckets)] for _ in range(6))
+
+    mel_spectrogram = MelSpec(target_sample_rate=target_sample_rate, n_mel_channels=n_mel_channels, hop_length=hop_length)
+
+    for utt, prompt_text, prompt_wav, gt_text, gt_wav in tqdm(metainfo, desc="Processing prompts..."):
+
+        # Audio
+        ref_audio, ref_sr = torchaudio.load(prompt_wav)
+        ref_rms = torch.sqrt(torch.mean(torch.square(ref_audio)))
+        if ref_rms < target_rms:
+            ref_audio = ref_audio * target_rms / ref_rms
+        assert ref_audio.shape[-1] > 5000, f"Empty prompt wav: {prompt_wav}, or torchaudio backend issue."
+        if ref_sr != target_sample_rate:
+            resampler = torchaudio.transforms.Resample(ref_sr, target_sample_rate)
+            ref_audio = resampler(ref_audio)
+
+        # Text
+        text = [prompt_text + gt_text]
+        if tokenizer == "pinyin":
+            text_list = convert_char_to_pinyin(text, polyphone = polyphone)
+        else:
+            text_list = text
+
+        # Duration, mel frame length
+        ref_mel_len = ref_audio.shape[-1] // hop_length
+        if use_truth_duration:
+            gt_audio, gt_sr = torchaudio.load(gt_wav)
+            if gt_sr != target_sample_rate:
+                resampler = torchaudio.transforms.Resample(gt_sr, target_sample_rate)
+                gt_audio = resampler(gt_audio)
+            total_mel_len = ref_mel_len + int(gt_audio.shape[-1] / hop_length / speed)
+
+            # # test vocoder resynthesis
+            # ref_audio = gt_audio
+        else:
+            zh_pause_punc = r"。，、；：？！"
+            ref_text_len = len(prompt_text) + len(re.findall(zh_pause_punc, prompt_text))
+            gen_text_len = len(gt_text) + len(re.findall(zh_pause_punc, gt_text))
+            total_mel_len = ref_mel_len + int(ref_mel_len / ref_text_len * gen_text_len / speed)
+
+        # to mel spectrogram
+        ref_mel = mel_spectrogram(ref_audio)
+        ref_mel = rearrange(ref_mel, '1 d n -> d n')
+
+        # deal with batch
+        assert infer_batch_size > 0, "infer_batch_size should be greater than 0."
+        assert min_tokens <= total_mel_len <= max_tokens, \
+            f"Audio {utt} has duration {total_mel_len*hop_length//target_sample_rate}s out of range [{min_secs}, {max_secs}]."
+        bucket_i = math.floor((total_mel_len - min_tokens) / (max_tokens - min_tokens + 1) * num_buckets)
+
+        utts[bucket_i].append(utt)
+        ref_rms_list[bucket_i].append(ref_rms)
+        ref_mels[bucket_i].append(ref_mel)
+        ref_mel_lens[bucket_i].append(ref_mel_len)
+        total_mel_lens[bucket_i].append(total_mel_len)
+        final_text_list[bucket_i].extend(text_list)
+
+        batch_accum[bucket_i] += total_mel_len
+
+        if batch_accum[bucket_i] >= infer_batch_size:
+            # print(f"\n{len(ref_mels[bucket_i][0][0])}\n{ref_mel_lens[bucket_i]}\n{total_mel_lens[bucket_i]}")
+            prompts_all.append((
+                utts[bucket_i], 
+                ref_rms_list[bucket_i], 
+                padded_mel_batch(ref_mels[bucket_i]), 
+                ref_mel_lens[bucket_i], 
+                total_mel_lens[bucket_i], 
+                final_text_list[bucket_i]
+            ))
+            batch_accum[bucket_i] = 0
+            utts[bucket_i], ref_rms_list[bucket_i], ref_mels[bucket_i], ref_mel_lens[bucket_i], total_mel_lens[bucket_i], final_text_list[bucket_i] = [], [], [], [], [], []
+
+    # add residual
+    for bucket_i, bucket_frames in enumerate(batch_accum):
+        if bucket_frames > 0:
+            prompts_all.append((
+                utts[bucket_i], 
+                ref_rms_list[bucket_i], 
+                padded_mel_batch(ref_mels[bucket_i]), 
+                ref_mel_lens[bucket_i], 
+                total_mel_lens[bucket_i], 
+                final_text_list[bucket_i]
+            ))
+    # not only leave easy work for last workers
+    random.seed(666)
+    random.shuffle(prompts_all)
+
+    return prompts_all
+
+
+# get wav_res_ref_text of seed-tts test metalst
+# https://github.com/BytedanceSpeech/seed-tts-eval
+
+def get_seed_tts_test(metalst, gen_wav_dir, gpus):
+    f = open(metalst)
+    lines = f.readlines()
+    f.close()
+
+    test_set_ = []
+    for line in tqdm(lines):
+        if len(line.strip().split('|')) == 5:
+            utt, prompt_text, prompt_wav, gt_text, gt_wav = line.strip().split('|')
+        elif len(line.strip().split('|')) == 4:
+            utt, prompt_text, prompt_wav, gt_text = line.strip().split('|')
+
+        if not os.path.exists(os.path.join(gen_wav_dir, utt + '.wav')):
+            continue
+        gen_wav = os.path.join(gen_wav_dir, utt + '.wav')
+        if not os.path.isabs(prompt_wav):
+            prompt_wav = os.path.join(os.path.dirname(metalst), prompt_wav)
+
+        test_set_.append((gen_wav, prompt_wav, gt_text))
+
+    num_jobs = len(gpus)
+    if num_jobs == 1:
+        return [(gpus[0], test_set_)]
+    
+    wav_per_job = len(test_set_) // num_jobs + 1
+    test_set = []
+    for i in range(num_jobs):
+        test_set.append((gpus[i], test_set_[i*wav_per_job:(i+1)*wav_per_job]))
+
+    return test_set
+
+
+# get librispeech test-clean cross sentence test
+
+def get_librispeech_test(metalst, gen_wav_dir, gpus, librispeech_test_clean_path, eval_ground_truth = False):
+    f = open(metalst)
+    lines = f.readlines()
+    f.close()
+
+    test_set_ = []
+    for line in tqdm(lines):
+        ref_utt, ref_dur, ref_txt, gen_utt, gen_dur, gen_txt = line.strip().split('\t')
+
+        if eval_ground_truth:
+            gen_spk_id, gen_chaptr_id, _ =  gen_utt.split('-')
+            gen_wav = os.path.join(librispeech_test_clean_path, gen_spk_id, gen_chaptr_id, gen_utt + '.flac')
+        else:
+            if not os.path.exists(os.path.join(gen_wav_dir, gen_utt + '.wav')):
+                raise FileNotFoundError(f"Generated wav not found: {gen_utt}")
+            gen_wav = os.path.join(gen_wav_dir, gen_utt + '.wav')
+
+        ref_spk_id, ref_chaptr_id, _ =  ref_utt.split('-')
+        ref_wav = os.path.join(librispeech_test_clean_path, ref_spk_id, ref_chaptr_id, ref_utt + '.flac')
+
+        test_set_.append((gen_wav, ref_wav, gen_txt))
+
+    num_jobs = len(gpus)
+    if num_jobs == 1:
+        return [(gpus[0], test_set_)]
+    
+    wav_per_job = len(test_set_) // num_jobs + 1
+    test_set = []
+    for i in range(num_jobs):
+        test_set.append((gpus[i], test_set_[i*wav_per_job:(i+1)*wav_per_job]))
+
+    return test_set
+
+
+# load asr model
+
+def load_asr_model(lang, ckpt_dir = ""):
+    if lang == "zh":
+        model = AutoModel(
+            model = os.path.join(ckpt_dir, "paraformer-zh"), 
+            # vad_model = os.path.join(ckpt_dir, "fsmn-vad"), 
+            # punc_model = os.path.join(ckpt_dir, "ct-punc"),
+            # spk_model = os.path.join(ckpt_dir, "cam++"), 
+            disable_update=True,
+            )  # following seed-tts setting
+    elif lang == "en":
+        model_size = "large-v3" if ckpt_dir == "" else ckpt_dir
+        model = WhisperModel(model_size, device="cuda", compute_type="float16")
+    return model
+
+
+# WER Evaluation, the way Seed-TTS does
+
+def run_asr_wer(args):
+    rank, lang, test_set, ckpt_dir = args
+
+    if lang == "zh":
+        torch.cuda.set_device(rank)
+    elif lang == "en":
+        os.environ["CUDA_VISIBLE_DEVICES"] = str(rank)
+    else:
+        raise NotImplementedError("lang support only 'zh' (funasr paraformer-zh), 'en' (faster-whisper-large-v3), for now.")
+
+    asr_model = load_asr_model(lang, ckpt_dir = ckpt_dir)
+
+    punctuation_all = punctuation + string.punctuation
+    wers = []
+
+    for gen_wav, prompt_wav, truth in tqdm(test_set):
+        if lang == "zh":
+            res = asr_model.generate(input=gen_wav, batch_size_s=300, disable_pbar=True)
+            hypo = res[0]["text"]
+            hypo = zhconv.convert(hypo, 'zh-cn')
+        elif lang == "en":
+            segments, _ = asr_model.transcribe(gen_wav, beam_size=5, language="en")
+            hypo = ''
+            for segment in segments:
+                hypo = hypo + ' ' + segment.text
+
+        # raw_truth = truth
+        # raw_hypo = hypo
+
+        for x in punctuation_all:
+            truth = truth.replace(x, '')
+            hypo = hypo.replace(x, '')
+
+        truth = truth.replace('  ', ' ')
+        hypo = hypo.replace('  ', ' ')
+
+        if lang == "zh":
+            truth = " ".join([x for x in truth])
+            hypo = " ".join([x for x in hypo])
+        elif lang == "en":
+            truth = truth.lower()
+            hypo = hypo.lower()
+
+        measures = compute_measures(truth, hypo)
+        wer = measures["wer"]
+
+        # ref_list = truth.split(" ")
+        # subs = measures["substitutions"] / len(ref_list)
+        # dele = measures["deletions"] / len(ref_list)
+        # inse = measures["insertions"] / len(ref_list)
+
+        wers.append(wer)
+
+    return wers
+
+
+# SIM Evaluation
+
+def run_sim(args):
+    rank, test_set, ckpt_dir = args
+    device = f"cuda:{rank}"
+
+    model = ECAPA_TDNN_SMALL(feat_dim=1024, feat_type='wavlm_large', config_path=None)
+    state_dict = torch.load(ckpt_dir, map_location=lambda storage, loc: storage)
+    model.load_state_dict(state_dict['model'], strict=False)
+
+    use_gpu=True if torch.cuda.is_available() else False
+    if use_gpu:
+        model = model.cuda(device)
+    model.eval()
+
+    sim_list = []
+    for wav1, wav2, truth in tqdm(test_set):
+
+        wav1, sr1 = torchaudio.load(wav1)
+        wav2, sr2 = torchaudio.load(wav2)
+
+        resample1 = torchaudio.transforms.Resample(orig_freq=sr1, new_freq=16000)
+        resample2 = torchaudio.transforms.Resample(orig_freq=sr2, new_freq=16000)
+        wav1 = resample1(wav1)
+        wav2 = resample2(wav2)
+
+        if use_gpu:
+            wav1 = wav1.cuda(device)
+            wav2 = wav2.cuda(device)
+        with torch.no_grad():
+            emb1 = model(wav1)
+            emb2 = model(wav2)
+        
+        sim = F.cosine_similarity(emb1, emb2)[0].item()
+        # print(f"VSim score between two audios: {sim:.4f} (-1.0, 1.0).")
+        sim_list.append(sim)
+    
+    return sim_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

packages.txt

@@ -0,0 +1 @@
+ffmpeg

requirements.txt

@@ -0,0 +1,24 @@
+accelerate>=0.33.0
+datasets
+einops>=0.8.0
+einx>=0.3.0
+ema_pytorch>=0.5.2
+faster_whisper
+funasr
+jieba
+jiwer
+librosa
+matplotlib
+pypinyin
+torch>=2.0
+torchaudio>=2.3.0
+torchdiffeq
+tqdm>=4.65.0
+transformers
+vocos
+wandb
+x_transformers>=1.31.14
+zhconv
+zhon
+cached_path
+pydub

scripts/count_max_epoch.py

@@ -0,0 +1,32 @@
+'''ADAPTIVE BATCH SIZE'''
+print('Adaptive batch size: using grouping batch sampler, frames_per_gpu fixed fed in')
+print('  -> least padding, gather wavs with accumulated frames in a batch\n')
+
+# data
+total_hours = 95282
+mel_hop_length = 256
+mel_sampling_rate = 24000
+
+# target
+wanted_max_updates = 1000000
+
+# train params
+gpus = 8
+frames_per_gpu = 38400  # 8 * 38400 = 307200
+grad_accum = 1
+
+# intermediate
+mini_batch_frames = frames_per_gpu * grad_accum * gpus
+mini_batch_hours = mini_batch_frames * mel_hop_length / mel_sampling_rate / 3600
+updates_per_epoch = total_hours / mini_batch_hours
+steps_per_epoch = updates_per_epoch * grad_accum
+
+# result
+epochs = wanted_max_updates / updates_per_epoch
+print(f"epochs should be set to: {epochs:.0f} ({epochs/grad_accum:.1f} x gd_acum {grad_accum})")
+print(f"progress_bar should show approx. 0/{updates_per_epoch:.0f} updates")
+print(f"                      or approx. 0/{steps_per_epoch:.0f} steps")
+
+# others
+print(f"total {total_hours:.0f} hours")
+print(f"mini-batch of {mini_batch_frames:.0f} frames, {mini_batch_hours:.2f} hours per mini-batch")

scripts/count_params_gflops.py

@@ -0,0 +1,35 @@
+import sys, os
+sys.path.append(os.getcwd())
+
+from model import M2_TTS, UNetT, DiT, MMDiT
+
+import torch
+import thop
+
+
+''' ~155M '''
+# transformer =     UNetT(dim = 768, depth = 20, heads = 12, ff_mult = 4)
+# transformer =     UNetT(dim = 768, depth = 20, heads = 12, ff_mult = 4, text_dim = 512, conv_layers = 4)
+# transformer =       DiT(dim = 768, depth = 18, heads = 12, ff_mult = 2)
+# transformer =       DiT(dim = 768, depth = 18, heads = 12, ff_mult = 2, text_dim = 512, conv_layers = 4)
+# transformer =       DiT(dim = 768, depth = 18, heads = 12, ff_mult = 2, text_dim = 512, conv_layers = 4, long_skip_connection = True)
+# transformer =     MMDiT(dim = 512, depth = 16, heads = 16, ff_mult = 2)
+
+''' ~335M '''
+# FLOPs: 622.1 G, Params: 333.2 M
+# transformer =     UNetT(dim = 1024, depth = 24, heads = 16, ff_mult = 4)
+# FLOPs: 363.4 G, Params: 335.8 M
+transformer =       DiT(dim = 1024, depth = 22, heads = 16, ff_mult = 2, text_dim = 512, conv_layers = 4)
+
+
+model = M2_TTS(transformer=transformer)
+target_sample_rate = 24000
+n_mel_channels = 100
+hop_length = 256
+duration = 20
+frame_length = int(duration * target_sample_rate / hop_length)
+text_length = 150
+
+flops, params = thop.profile(model, inputs=(torch.randn(1, frame_length, n_mel_channels), torch.zeros(1, text_length, dtype=torch.long)))
+print(f"FLOPs: {flops / 1e9} G")
+print(f"Params: {params / 1e6} M")

scripts/eval_librispeech_test_clean.py

@@ -0,0 +1,67 @@
+# Evaluate with Librispeech test-clean, ~3s prompt to generate 4-10s audio (the way of valle/voicebox evaluation)
+
+import sys, os
+sys.path.append(os.getcwd())
+
+import multiprocessing as mp
+import numpy as np
+
+from model.utils import (
+    get_librispeech_test,
+    run_asr_wer,
+    run_sim,
+)
+
+
+eval_task = "wer"  # sim | wer
+lang = "en"
+metalst = "data/librispeech_pc_test_clean_cross_sentence.lst"
+librispeech_test_clean_path = "<SOME_PATH>/LibriSpeech/test-clean"  # test-clean path
+gen_wav_dir = "PATH_TO_GENERATED"  # generated wavs
+
+gpus = [0,1,2,3,4,5,6,7]
+test_set = get_librispeech_test(metalst, gen_wav_dir, gpus, librispeech_test_clean_path)
+
+## In LibriSpeech, some speakers utilized varying voice characteristics for different characters in the book,
+## leading to a low similarity for the ground truth in some cases.
+# test_set = get_librispeech_test(metalst, gen_wav_dir, gpus, librispeech_test_clean_path, eval_ground_truth = True)  # eval ground truth
+
+local = False
+if local:  # use local custom checkpoint dir
+    asr_ckpt_dir = "../checkpoints/Systran/faster-whisper-large-v3"
+else:
+    asr_ckpt_dir = ""  # auto download to cache dir
+
+wavlm_ckpt_dir = "../checkpoints/UniSpeech/wavlm_large_finetune.pth"
+
+
+# --------------------------- WER ---------------------------
+
+if eval_task == "wer":
+    wers = []
+
+    with mp.Pool(processes=len(gpus)) as pool:
+        args = [(rank, lang, sub_test_set, asr_ckpt_dir) for (rank, sub_test_set) in test_set]
+        results = pool.map(run_asr_wer, args)
+        for wers_ in results:
+            wers.extend(wers_)
+
+    wer = round(np.mean(wers)*100, 3)
+    print(f"\nTotal {len(wers)} samples")
+    print(f"WER      : {wer}%")
+
+
+# --------------------------- SIM ---------------------------
+
+if eval_task == "sim":
+    sim_list = []
+
+    with mp.Pool(processes=len(gpus)) as pool:
+        args = [(rank, sub_test_set, wavlm_ckpt_dir) for (rank, sub_test_set) in test_set]
+        results = pool.map(run_sim, args)
+        for sim_ in results:
+            sim_list.extend(sim_)
+
+    sim = round(sum(sim_list)/len(sim_list), 3)
+    print(f"\nTotal {len(sim_list)} samples")
+    print(f"SIM      : {sim}")

scripts/eval_seedtts_testset.py

@@ -0,0 +1,69 @@
+# Evaluate with Seed-TTS testset
+
+import sys, os
+sys.path.append(os.getcwd())
+
+import multiprocessing as mp
+import numpy as np
+
+from model.utils import (
+    get_seed_tts_test,
+    run_asr_wer,
+    run_sim,
+)
+
+
+eval_task = "wer"  # sim | wer
+lang = "zh"        # zh | en
+metalst = f"data/seedtts_testset/{lang}/meta.lst"  # seed-tts testset
+# gen_wav_dir = f"data/seedtts_testset/{lang}/wavs"  # ground truth wavs
+gen_wav_dir = f"PATH_TO_GENERATED"  # generated wavs
+
+
+# NOTE. paraformer-zh result will be slightly different according to the number of gpus, cuz batchsize is different
+#       zh 1.254 seems a result of 4 workers wer_seed_tts 
+gpus = [0,1,2,3,4,5,6,7]
+test_set = get_seed_tts_test(metalst, gen_wav_dir, gpus)
+
+local = False
+if local:  # use local custom checkpoint dir
+    if lang == "zh":
+        asr_ckpt_dir = "../checkpoints/funasr" # paraformer-zh dir under funasr
+    elif lang == "en":
+        asr_ckpt_dir = "../checkpoints/Systran/faster-whisper-large-v3"
+else:
+    asr_ckpt_dir = ""  # auto download to cache dir
+
+wavlm_ckpt_dir = "../checkpoints/UniSpeech/wavlm_large_finetune.pth"
+
+
+# --------------------------- WER ---------------------------
+
+if eval_task == "wer":
+    wers = []
+
+    with mp.Pool(processes=len(gpus)) as pool:
+        args = [(rank, lang, sub_test_set, asr_ckpt_dir) for (rank, sub_test_set) in test_set]
+        results = pool.map(run_asr_wer, args)
+        for wers_ in results:
+            wers.extend(wers_)
+
+    wer = round(np.mean(wers)*100, 3)
+    print(f"\nTotal {len(wers)} samples")
+    print(f"WER      : {wer}%")
+
+
+# --------------------------- SIM ---------------------------
+
+if eval_task == "sim":
+    sim_list = []
+
+    with mp.Pool(processes=len(gpus)) as pool:
+        args = [(rank, sub_test_set, wavlm_ckpt_dir) for (rank, sub_test_set) in test_set]
+        results = pool.map(run_sim, args)
+        for sim_ in results:
+            sim_list.extend(sim_)
+
+    sim = round(sum(sim_list)/len(sim_list), 3)
+    print(f"\nTotal {len(sim_list)} samples")
+    print(f"SIM      : {sim}")

scripts/prepare_emilia.py

@@ -0,0 +1,143 @@
+# Emilia Dataset: https://huggingface.co/datasets/amphion/Emilia-Dataset/tree/fc71e07
+# if use updated new version, i.e. WebDataset, feel free to modify / draft your own script
+
+# generate audio text map for Emilia ZH & EN
+# evaluate for vocab size
+
+import sys, os
+sys.path.append(os.getcwd())
+
+from pathlib import Path
+import json
+from tqdm import tqdm
+from concurrent.futures import ProcessPoolExecutor
+
+from datasets import Dataset
+from datasets.arrow_writer import ArrowWriter
+
+from model.utils import (
+    repetition_found,
+    convert_char_to_pinyin,
+)
+
+
+out_zh = {"ZH_B00041_S06226", "ZH_B00042_S09204", "ZH_B00065_S09430", "ZH_B00065_S09431", "ZH_B00066_S09327", "ZH_B00066_S09328"}
+zh_filters = ["い", "て"]
+# seems synthesized audios, or heavily code-switched
+out_en = {
+    "EN_B00013_S00913", "EN_B00042_S00120", "EN_B00055_S04111", "EN_B00061_S00693", "EN_B00061_S01494", "EN_B00061_S03375",
+    
+    "EN_B00059_S00092", "EN_B00111_S04300", "EN_B00100_S03759", "EN_B00087_S03811", "EN_B00059_S00950", "EN_B00089_S00946", "EN_B00078_S05127", "EN_B00070_S04089", "EN_B00074_S09659", "EN_B00061_S06983", "EN_B00061_S07060", "EN_B00059_S08397", "EN_B00082_S06192", "EN_B00091_S01238", "EN_B00089_S07349", "EN_B00070_S04343", "EN_B00061_S02400", "EN_B00076_S01262", "EN_B00068_S06467", "EN_B00076_S02943", "EN_B00064_S05954", "EN_B00061_S05386", "EN_B00066_S06544", "EN_B00076_S06944", "EN_B00072_S08620", "EN_B00076_S07135", "EN_B00076_S09127", "EN_B00065_S00497", "EN_B00059_S06227", "EN_B00063_S02859", "EN_B00075_S01547", "EN_B00061_S08286", "EN_B00079_S02901", "EN_B00092_S03643", "EN_B00096_S08653", "EN_B00063_S04297", "EN_B00063_S04614", "EN_B00079_S04698", "EN_B00104_S01666", "EN_B00061_S09504", "EN_B00061_S09694", "EN_B00065_S05444", "EN_B00063_S06860", "EN_B00065_S05725", "EN_B00069_S07628", "EN_B00083_S03875", "EN_B00071_S07665", "EN_B00071_S07665", "EN_B00062_S04187", "EN_B00065_S09873", "EN_B00065_S09922", "EN_B00084_S02463", "EN_B00067_S05066", "EN_B00106_S08060", "EN_B00073_S06399", "EN_B00073_S09236", "EN_B00087_S00432", "EN_B00085_S05618", "EN_B00064_S01262", "EN_B00072_S01739", "EN_B00059_S03913", "EN_B00069_S04036", "EN_B00067_S05623", "EN_B00060_S05389", "EN_B00060_S07290", "EN_B00062_S08995",
+}
+en_filters = ["ا", "い", "て"]
+
+
+def deal_with_audio_dir(audio_dir):
+    audio_jsonl = audio_dir.with_suffix(".jsonl")
+    sub_result, durations = [], []
+    vocab_set = set()
+    bad_case_zh = 0
+    bad_case_en = 0
+    with open(audio_jsonl, "r") as f:
+        lines = f.readlines()
+        for line in tqdm(lines, desc=f"{audio_jsonl.stem}"):
+            obj = json.loads(line)
+            text = obj["text"]
+            if obj['language'] == "zh":
+                if obj["wav"].split("/")[1] in out_zh or any(f in text for f in zh_filters) or repetition_found(text):
+                    bad_case_zh += 1
+                    continue
+                else:
+                    text = text.translate(str.maketrans({',': '，', '!': '！', '?': '？'}))  # not "。" cuz much code-switched
+            if obj['language'] == "en":
+                if obj["wav"].split("/")[1] in out_en or any(f in text for f in en_filters) or repetition_found(text, length=4):
+                    bad_case_en += 1
+                    continue
+            if tokenizer == "pinyin":
+                text = convert_char_to_pinyin([text], polyphone = polyphone)[0]
+            duration = obj["duration"]
+            sub_result.append({"audio_path": str(audio_dir.parent / obj["wav"]), "text": text, "duration": duration})
+            durations.append(duration)
+            vocab_set.update(list(text))
+    return sub_result, durations, vocab_set, bad_case_zh, bad_case_en
+
+
+def main():
+    assert tokenizer in ["pinyin", "char"]
+    result = []
+    duration_list = []
+    text_vocab_set = set()
+    total_bad_case_zh = 0
+    total_bad_case_en = 0
+
+    # process raw data
+    executor = ProcessPoolExecutor(max_workers=max_workers)
+    futures = []
+    for lang in langs:
+        dataset_path = Path(os.path.join(dataset_dir, lang))
+        [
+            futures.append(executor.submit(deal_with_audio_dir, audio_dir))
+            for audio_dir in dataset_path.iterdir()
+            if audio_dir.is_dir()
+        ]
+    for futures in tqdm(futures, total=len(futures)):
+        sub_result, durations, vocab_set, bad_case_zh, bad_case_en = futures.result()
+        result.extend(sub_result)
+        duration_list.extend(durations)
+        text_vocab_set.update(vocab_set)
+        total_bad_case_zh += bad_case_zh
+        total_bad_case_en += bad_case_en
+    executor.shutdown()
+
+    # save preprocessed dataset to disk
+    if not os.path.exists(f"data/{dataset_name}"):
+        os.makedirs(f"data/{dataset_name}")
+    print(f"\nSaving to data/{dataset_name} ...")
+    # dataset = Dataset.from_dict({"audio_path": audio_path_list, "text": text_list, "duration": duration_list})  # oom
+    # dataset.save_to_disk(f"data/{dataset_name}/raw", max_shard_size="2GB")
+    with ArrowWriter(path=f"data/{dataset_name}/raw.arrow") as writer:
+        for line in tqdm(result, desc=f"Writing to raw.arrow ..."):
+            writer.write(line)
+
+    # dup a json separately saving duration in case for DynamicBatchSampler ease
+    with open(f"data/{dataset_name}/duration.json", 'w', encoding='utf-8') as f:
+        json.dump({"duration": duration_list}, f, ensure_ascii=False)
+
+    # vocab map, i.e. tokenizer
+    # add alphabets and symbols (optional, if plan to ft on de/fr etc.)
+    # if tokenizer == "pinyin":
+    #     text_vocab_set.update([chr(i) for i in range(32, 127)] + [chr(i) for i in range(192, 256)])
+    with open(f"data/{dataset_name}/vocab.txt", "w") as f:
+        for vocab in sorted(text_vocab_set):
+            f.write(vocab + "\n")
+
+    print(f"\nFor {dataset_name}, sample count: {len(result)}")
+    print(f"For {dataset_name}, vocab size is: {len(text_vocab_set)}")
+    print(f"For {dataset_name}, total {sum(duration_list)/3600:.2f} hours")
+    if "ZH" in langs: print(f"Bad zh transcription case: {total_bad_case_zh}")
+    if "EN" in langs: print(f"Bad en transcription case: {total_bad_case_en}\n")
+
+
+if __name__ == "__main__":
+
+    max_workers = 32
+
+    tokenizer = "pinyin"  # "pinyin" | "char"
+    polyphone = True
+
+    langs = ["ZH", "EN"]
+    dataset_dir = "<SOME_PATH>/Emilia_Dataset/raw"
+    dataset_name = f"Emilia_{'_'.join(langs)}_{tokenizer}"
+    print(f"\nPrepare for {dataset_name}\n")
+
+    main()
+
+    # Emilia               ZH & EN
+    # samples count       37837916   (after removal)
+    # pinyin vocab size       2543   (polyphone)
+    # total duration      95281.87   (hours)
+    # bad zh asr cnt        230435   (samples)
+    # bad eh asr cnt         37217   (samples)
+
+    # vocab size may be slightly different due to jieba tokenizer and pypinyin (e.g. way of polyphoneme)
+    # please be careful if using pretrained model, make sure the vocab.txt is same

scripts/prepare_wenetspeech4tts.py

@@ -0,0 +1,116 @@
+# generate audio text map for WenetSpeech4TTS
+# evaluate for vocab size
+
+import sys, os
+sys.path.append(os.getcwd())
+
+import json
+from tqdm import tqdm
+from concurrent.futures import ProcessPoolExecutor
+
+import torchaudio
+from datasets import Dataset
+
+from model.utils import convert_char_to_pinyin
+
+
+def deal_with_sub_path_files(dataset_path, sub_path):
+    print(f"Dealing with: {sub_path}")
+
+    text_dir = os.path.join(dataset_path, sub_path, "txts")
+    audio_dir = os.path.join(dataset_path, sub_path, "wavs")
+    text_files = os.listdir(text_dir)
+
+    audio_paths, texts, durations = [], [], []
+    for text_file in tqdm(text_files):
+        with open(os.path.join(text_dir, text_file), 'r', encoding='utf-8') as file:
+            first_line = file.readline().split("\t")
+        audio_nm = first_line[0]
+        audio_path = os.path.join(audio_dir, audio_nm + ".wav")
+        text = first_line[1].strip()
+
+        audio_paths.append(audio_path)
+
+        if tokenizer == "pinyin":
+            texts.extend(convert_char_to_pinyin([text], polyphone = polyphone))
+        elif tokenizer == "char":
+            texts.append(text)
+
+        audio, sample_rate = torchaudio.load(audio_path)
+        durations.append(audio.shape[-1] / sample_rate)
+
+    return audio_paths, texts, durations
+
+
+def main():
+    assert tokenizer in ["pinyin", "char"]
+
+    audio_path_list, text_list, duration_list = [], [], []
+    
+    executor = ProcessPoolExecutor(max_workers=max_workers)
+    futures = []
+    for dataset_path in dataset_paths:
+        sub_items = os.listdir(dataset_path)
+        sub_paths = [item for item in sub_items if os.path.isdir(os.path.join(dataset_path, item))]
+        for sub_path in sub_paths:
+            futures.append(executor.submit(deal_with_sub_path_files, dataset_path, sub_path))
+    for future in tqdm(futures, total=len(futures)):
+        audio_paths, texts, durations = future.result()
+        audio_path_list.extend(audio_paths)
+        text_list.extend(texts)
+        duration_list.extend(durations)
+    executor.shutdown()
+
+    if not os.path.exists("data"):
+        os.makedirs("data")
+
+    print(f"\nSaving to data/{dataset_name}_{tokenizer} ...")
+    dataset = Dataset.from_dict({"audio_path": audio_path_list, "text": text_list, "duration": duration_list})
+    dataset.save_to_disk(f"data/{dataset_name}_{tokenizer}/raw", max_shard_size="2GB")  # arrow format
+
+    with open(f"data/{dataset_name}_{tokenizer}/duration.json", 'w', encoding='utf-8') as f:
+        json.dump({"duration": duration_list}, f, ensure_ascii=False)  # dup a json separately saving duration in case for DynamicBatchSampler ease
+
+    print("\nEvaluating vocab size (all characters and symbols / all phonemes) ...")
+    text_vocab_set = set()
+    for text in tqdm(text_list):
+        text_vocab_set.update(list(text))
+
+    # add alphabets and symbols (optional, if plan to ft on de/fr etc.)
+    if tokenizer == "pinyin":
+        text_vocab_set.update([chr(i) for i in range(32, 127)] + [chr(i) for i in range(192, 256)])
+
+    with open(f"data/{dataset_name}_{tokenizer}/vocab.txt", "w") as f:
+        for vocab in sorted(text_vocab_set):
+            f.write(vocab + "\n")
+    print(f"\nFor {dataset_name}, sample count: {len(text_list)}")
+    print(f"For {dataset_name}, vocab size is: {len(text_vocab_set)}\n")
+    
+
+if __name__ == "__main__":
+
+    max_workers = 32
+
+    tokenizer = "pinyin"  # "pinyin" | "char"
+    polyphone = True
+    dataset_choice = 1  # 1: Premium, 2: Standard, 3: Basic
+
+    dataset_name = ["WenetSpeech4TTS_Premium", "WenetSpeech4TTS_Standard", "WenetSpeech4TTS_Basic"][dataset_choice-1]
+    dataset_paths = [
+        "<SOME_PATH>/WenetSpeech4TTS/Basic",
+        "<SOME_PATH>/WenetSpeech4TTS/Standard",
+        "<SOME_PATH>/WenetSpeech4TTS/Premium",
+        ][-dataset_choice:]
+    print(f"\nChoose Dataset: {dataset_name}\n")
+
+    main()
+
+    # Results (if adding alphabets with accents and symbols):
+    # WenetSpeech4TTS       Basic     Standard     Premium
+    # samples count       3932473      1941220      407494
+    # pinyin vocab size      1349         1348        1344   (no polyphone)
+    #                           -            -        1459   (polyphone) 
+    # char   vocab size      5264         5219        5042
+    
+    # vocab size may be slightly different due to jieba tokenizer and pypinyin (e.g. way of polyphoneme)
+    # please be careful if using pretrained model, make sure the vocab.txt is same

test_infer_batch.py

@@ -0,0 +1,202 @@
+import os
+import time
+import random
+from tqdm import tqdm
+import argparse
+
+import torch
+import torchaudio
+from accelerate import Accelerator
+from einops import rearrange
+from ema_pytorch import EMA
+from vocos import Vocos
+
+from model import CFM, UNetT, DiT
+from model.utils import (
+    get_tokenizer, 
+    get_seedtts_testset_metainfo, 
+    get_librispeech_test_clean_metainfo, 
+    get_inference_prompt,
+)
+
+accelerator = Accelerator()
+device = f"cuda:{accelerator.process_index}"
+
+
+# --------------------- Dataset Settings -------------------- #
+
+target_sample_rate = 24000
+n_mel_channels = 100
+hop_length = 256
+target_rms = 0.1
+
+tokenizer = "pinyin"
+
+
+# ---------------------- infer setting ---------------------- #
+
+parser = argparse.ArgumentParser(description="batch inference")
+
+parser.add_argument('-s', '--seed', default=None, type=int)
+parser.add_argument('-d', '--dataset', default="Emilia_ZH_EN")
+parser.add_argument('-n', '--expname', required=True)
+parser.add_argument('-c', '--ckptstep', default=1200000, type=int)
+
+parser.add_argument('-nfe', '--nfestep', default=32, type=int)
+parser.add_argument('-o', '--odemethod', default="euler")
+parser.add_argument('-ss', '--swaysampling', default=-1, type=float)
+
+parser.add_argument('-t', '--testset', required=True)
+
+args = parser.parse_args()
+
+
+seed = args.seed
+dataset_name = args.dataset
+exp_name = args.expname
+ckpt_step = args.ckptstep
+checkpoint = torch.load(f"ckpts/{exp_name}/model_{ckpt_step}.pt", map_location=device)
+
+nfe_step = args.nfestep
+ode_method = args.odemethod
+sway_sampling_coef = args.swaysampling
+
+testset = args.testset
+
+
+infer_batch_size = 1  # max frames. 1 for ddp single inference (recommended)
+cfg_strength = 2.
+speed = 1.
+use_truth_duration = False
+no_ref_audio = False
+
+
+if exp_name == "F5TTS_Base":
+    model_cls = DiT
+    model_cfg = dict(dim = 1024, depth = 22, heads = 16, ff_mult = 2, text_dim = 512, conv_layers = 4)
+
+elif exp_name == "E2TTS_Base":
+    model_cls = UNetT
+    model_cfg = dict(dim = 1024, depth = 24, heads = 16, ff_mult = 4)
+
+
+if testset == "ls_pc_test_clean":
+    metalst = "data/librispeech_pc_test_clean_cross_sentence.lst"
+    librispeech_test_clean_path = "<SOME_PATH>/LibriSpeech/test-clean"  # test-clean path
+    metainfo = get_librispeech_test_clean_metainfo(metalst, librispeech_test_clean_path)
+    
+elif testset == "seedtts_test_zh":
+    metalst = "data/seedtts_testset/zh/meta.lst"
+    metainfo = get_seedtts_testset_metainfo(metalst)
+
+elif testset == "seedtts_test_en":
+    metalst = "data/seedtts_testset/en/meta.lst"
+    metainfo = get_seedtts_testset_metainfo(metalst)
+
+
+# path to save genereted wavs
+if seed is None: seed = random.randint(-10000, 10000)
+output_dir = f"results/{exp_name}_{ckpt_step}/{testset}/" \
+    f"seed{seed}_{ode_method}_nfe{nfe_step}" \
+    f"{f'_ss{sway_sampling_coef}' if sway_sampling_coef else ''}" \
+    f"_cfg{cfg_strength}_speed{speed}" \
+    f"{'_gt-dur' if use_truth_duration else ''}" \
+    f"{'_no-ref-audio' if no_ref_audio else ''}"
+
+
+# -------------------------------------------------#
+
+use_ema = True
+
+prompts_all = get_inference_prompt(
+    metainfo, 
+    speed = speed, 
+    tokenizer = tokenizer, 
+    target_sample_rate = target_sample_rate, 
+    n_mel_channels = n_mel_channels,
+    hop_length = hop_length,
+    target_rms = target_rms,
+    use_truth_duration = use_truth_duration,
+    infer_batch_size = infer_batch_size,
+)
+
+# Vocoder model
+local = False
+if local:
+    vocos_local_path = "../checkpoints/charactr/vocos-mel-24khz"
+    vocos = Vocos.from_hparams(f"{vocos_local_path}/config.yaml")
+    state_dict = torch.load(f"{vocos_local_path}/pytorch_model.bin", map_location=device)
+    vocos.load_state_dict(state_dict)
+    vocos.eval()
+else:
+    vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")
+
+# Tokenizer
+vocab_char_map, vocab_size = get_tokenizer(dataset_name, tokenizer)
+
+# Model
+model = CFM(
+    transformer = model_cls(
+        **model_cfg,
+        text_num_embeds = vocab_size, 
+        mel_dim = n_mel_channels
+    ),
+    mel_spec_kwargs = dict(
+        target_sample_rate = target_sample_rate, 
+        n_mel_channels = n_mel_channels,
+        hop_length = hop_length,
+    ),
+    odeint_kwargs = dict(
+        method = ode_method,
+    ),
+    vocab_char_map = vocab_char_map,
+).to(device)
+
+if use_ema == True:
+    ema_model = EMA(model, include_online_model = False).to(device)
+    ema_model.load_state_dict(checkpoint['ema_model_state_dict'])
+    ema_model.copy_params_from_ema_to_model()
+else:
+    model.load_state_dict(checkpoint['model_state_dict'])
+
+if not os.path.exists(output_dir) and accelerator.is_main_process:
+    os.makedirs(output_dir)
+
+# start batch inference
+accelerator.wait_for_everyone()
+start = time.time()
+
+with accelerator.split_between_processes(prompts_all) as prompts:
+
+    for prompt in tqdm(prompts, disable=not accelerator.is_local_main_process):
+        utts, ref_rms_list, ref_mels, ref_mel_lens, total_mel_lens, final_text_list = prompt
+        ref_mels = ref_mels.to(device)
+        ref_mel_lens = torch.tensor(ref_mel_lens, dtype = torch.long).to(device)
+        total_mel_lens = torch.tensor(total_mel_lens, dtype = torch.long).to(device)
+        
+        # Inference
+        with torch.inference_mode():
+            generated, _ = model.sample(
+                cond = ref_mels,
+                text = final_text_list,
+                duration = total_mel_lens,
+                lens = ref_mel_lens,
+                steps = nfe_step,
+                cfg_strength = cfg_strength,
+                sway_sampling_coef = sway_sampling_coef,
+                no_ref_audio = no_ref_audio,
+                seed = seed,
+            )
+        # 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')
+            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
+            torchaudio.save(f"{output_dir}/{utts[i]}.wav", generated_wave, target_sample_rate)
+
+accelerator.wait_for_everyone()
+if accelerator.is_main_process:
+    timediff = time.time() - start
+    print(f"Done batch inference in {timediff / 60 :.2f} minutes.")

test_infer_batch.sh

@@ -0,0 +1,13 @@
+#!/bin/bash
+
+# e.g. F5-TTS, 16 NFE
+accelerate launch test_infer_batch.py -n "F5TTS_Base" -t "seedtts_test_zh" -nfe 16
+accelerate launch test_infer_batch.py -n "F5TTS_Base" -t "seedtts_test_en" -nfe 16
+accelerate launch test_infer_batch.py -n "F5TTS_Base" -t "ls_pc_test_clean" -nfe 16
+
+# e.g. Vanilla E2 TTS, 32 NFE
+accelerate launch test_infer_batch.py -n "E2TTS_Base" -t "seedtts_test_zh" -o "midpoint" -ss 0
+accelerate launch test_infer_batch.py -n "E2TTS_Base" -t "seedtts_test_en" -o "midpoint" -ss 0
+accelerate launch test_infer_batch.py -n "E2TTS_Base" -t "ls_pc_test_clean" -o "midpoint" -ss 0
+
+# etc.

test_infer_single.py

@@ -0,0 +1,162 @@
+import os
+import re
+
+import torch
+import torchaudio
+from einops import rearrange
+from ema_pytorch import EMA
+from vocos import Vocos
+
+from model import CFM, UNetT, DiT, MMDiT
+from model.utils import (
+    get_tokenizer, 
+    convert_char_to_pinyin, 
+    save_spectrogram,
+)
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+
+# --------------------- Dataset Settings -------------------- #
+
+target_sample_rate = 24000
+n_mel_channels = 100
+hop_length = 256
+target_rms = 0.1
+
+tokenizer = "pinyin"
+dataset_name = "Emilia_ZH_EN"
+
+
+# ---------------------- infer setting ---------------------- #
+
+seed = None  # int | None
+
+exp_name = "F5TTS_Base"  # F5TTS_Base | E2TTS_Base
+ckpt_step = 1200000
+
+nfe_step = 32  # 16, 32
+cfg_strength = 2.
+ode_method = 'euler'  # euler | midpoint
+sway_sampling_coef = -1.
+speed = 1.
+fix_duration = 27  # None (will linear estimate. if code-switched, consider fix) | float (total in seconds, include ref audio) 
+
+if exp_name == "F5TTS_Base":
+    model_cls = DiT
+    model_cfg = dict(dim = 1024, depth = 22, heads = 16, ff_mult = 2, text_dim = 512, conv_layers = 4)
+
+elif exp_name == "E2TTS_Base":
+    model_cls = UNetT
+    model_cfg = dict(dim = 1024, depth = 24, heads = 16, ff_mult = 4)
+
+checkpoint = torch.load(f"ckpts/{exp_name}/model_{ckpt_step}.pt", map_location=device)
+output_dir = "tests"
+
+ref_audio = "tests/ref_audio/test_en_1_ref_short.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."
+
+# ref_audio = "tests/ref_audio/test_zh_1_ref_short.wav"
+# ref_text = "对，这就是我，万人敬仰的太乙真人。"
+# gen_text = "突然，身边一阵笑声。我看着他们，意气风发地挺直了胸膛，甩了甩那稍显肉感的双臂，轻笑道：\"我身上的肉，是为了掩饰我爆棚的魅力，否则，岂不吓坏了你们呢？\""
+
+
+# -------------------------------------------------#
+
+use_ema = True
+
+if not os.path.exists(output_dir):
+    os.makedirs(output_dir)
+
+# Vocoder model
+local = False
+if local:
+    vocos_local_path = "../checkpoints/charactr/vocos-mel-24khz"
+    vocos = Vocos.from_hparams(f"{vocos_local_path}/config.yaml")
+    state_dict = torch.load(f"{vocos_local_path}/pytorch_model.bin", map_location=device)
+    vocos.load_state_dict(state_dict)
+    vocos.eval()
+else:
+    vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")
+
+# Tokenizer
+vocab_char_map, vocab_size = get_tokenizer(dataset_name, tokenizer)
+
+# Model
+model = CFM(
+    transformer = model_cls(
+        **model_cfg,
+        text_num_embeds = vocab_size, 
+        mel_dim = n_mel_channels
+    ),
+    mel_spec_kwargs = dict(
+        target_sample_rate = target_sample_rate, 
+        n_mel_channels = n_mel_channels,
+        hop_length = hop_length,
+    ),
+    odeint_kwargs = dict(
+        method = ode_method,
+    ),
+    vocab_char_map = vocab_char_map,
+).to(device)
+
+if use_ema == True:
+    ema_model = EMA(model, include_online_model = False).to(device)
+    ema_model.load_state_dict(checkpoint['ema_model_state_dict'])
+    ema_model.copy_params_from_ema_to_model()
+else:
+    model.load_state_dict(checkpoint['model_state_dict'])
+
+# Audio
+audio, sr = torchaudio.load(ref_audio)
+rms = torch.sqrt(torch.mean(torch.square(audio)))
+if rms < target_rms:
+    audio = audio * target_rms / rms
+if sr != target_sample_rate:
+    resampler = torchaudio.transforms.Resample(sr, target_sample_rate)
+    audio = resampler(audio)
+audio = audio.to(device)
+
+# Text
+text_list = [ref_text + gen_text]
+if tokenizer == "pinyin":
+    final_text_list = convert_char_to_pinyin(text_list)
+else:
+    final_text_list = [text_list]
+print(f"text  : {text_list}")
+print(f"pinyin: {final_text_list}")
+
+# Duration
+ref_audio_len = audio.shape[-1] // hop_length
+if fix_duration is not None:
+    duration = int(fix_duration * target_sample_rate / hop_length)
+else:  # simple linear scale calcul
+    zh_pause_punc = r"。，、；：？！"
+    ref_text_len = len(ref_text) + len(re.findall(zh_pause_punc, ref_text))
+    gen_text_len = len(gen_text) + len(re.findall(zh_pause_punc, gen_text))
+    duration = ref_audio_len + int(ref_audio_len / ref_text_len * gen_text_len / speed)
+
+# Inference
+with torch.inference_mode():
+    generated, trajectory = model.sample(
+        cond = audio,
+        text = final_text_list,
+        duration = duration,
+        steps = nfe_step,
+        cfg_strength = cfg_strength,
+        sway_sampling_coef = sway_sampling_coef,
+        seed = seed,
+    )
+print(f"Generated mel: {generated.shape}")
+
+# Final result
+generated = generated[:, ref_audio_len:, :]
+generated_mel_spec = rearrange(generated, '1 n d -> 1 d n')
+generated_wave = vocos.decode(generated_mel_spec.cpu())
+if rms < target_rms:
+    generated_wave = generated_wave * rms / target_rms
+
+save_spectrogram(generated_mel_spec[0].cpu().numpy(), f"{output_dir}/test_single.png")
+torchaudio.save(f"{output_dir}/test_single.wav", generated_wave, target_sample_rate)
+print(f"Generated wav: {generated_wave.shape}")

test_train.py

@@ -0,0 +1,91 @@
+from model import CFM, UNetT, DiT, MMDiT, Trainer
+from model.utils import get_tokenizer
+from model.dataset import load_dataset
+
+
+# -------------------------- Dataset Settings --------------------------- #
+
+target_sample_rate = 24000
+n_mel_channels = 100
+hop_length = 256
+
+tokenizer = "pinyin"
+dataset_name = "Emilia_ZH_EN"
+
+
+# -------------------------- Training Settings -------------------------- #
+
+exp_name = "F5TTS_Base"  # F5TTS_Base | E2TTS_Base
+
+learning_rate = 7.5e-5
+
+batch_size_per_gpu = 38400  # 8 GPUs, 8 * 38400 = 307200
+batch_size_type = "frame"  # "frame" or "sample"
+max_samples = 64  # max sequences per batch if use frame-wise batch_size. we set 32 for small models, 64 for base models
+grad_accumulation_steps = 1  # note: updates = steps / grad_accumulation_steps
+max_grad_norm = 1.
+
+epochs = 11  # use linear decay, thus epochs control the slope
+num_warmup_updates = 20000  # warmup steps
+save_per_updates = 50000  # save checkpoint per steps
+last_per_steps = 5000  # save last checkpoint per steps
+
+# model params
+if exp_name == "F5TTS_Base":
+    wandb_resume_id = None
+    model_cls = DiT
+    model_cfg = dict(dim = 1024, depth = 22, heads = 16, ff_mult = 2, text_dim = 512, conv_layers = 4)
+elif exp_name == "E2TTS_Base":
+    wandb_resume_id = None
+    model_cls = UNetT
+    model_cfg = dict(dim = 1024, depth = 24, heads = 16, ff_mult = 4)
+
+
+# ----------------------------------------------------------------------- #
+
+def main():
+
+    vocab_char_map, vocab_size = get_tokenizer(dataset_name, tokenizer)
+
+    mel_spec_kwargs = dict(
+            target_sample_rate = target_sample_rate, 
+            n_mel_channels = n_mel_channels,
+            hop_length = hop_length,
+        )
+    
+    e2tts = CFM(
+        transformer = model_cls(
+            **model_cfg,
+            text_num_embeds = vocab_size, 
+            mel_dim = n_mel_channels
+        ),
+        mel_spec_kwargs = mel_spec_kwargs,
+        vocab_char_map = vocab_char_map,
+    )
+
+    trainer = Trainer(
+        e2tts,
+        epochs, 
+        learning_rate,
+        num_warmup_updates = num_warmup_updates,
+        save_per_updates = save_per_updates, 
+        checkpoint_path = f'ckpts/{exp_name}',
+        batch_size = batch_size_per_gpu, 
+        batch_size_type = batch_size_type,
+        max_samples = max_samples,
+        grad_accumulation_steps = grad_accumulation_steps,
+        max_grad_norm = max_grad_norm,
+        wandb_project = "CFM-TTS",
+        wandb_run_name = exp_name,
+        wandb_resume_id = wandb_resume_id,
+        last_per_steps = last_per_steps,
+    )
+
+    train_dataset = load_dataset(dataset_name, tokenizer, mel_spec_kwargs=mel_spec_kwargs)
+    trainer.train(train_dataset, 
+                  resumable_with_seed = 666 # seed for shuffling dataset
+                  )
+
+
+if __name__ == '__main__':
+    main()