Upload 6 files

cog.py

@@ -0,0 +1,180 @@
+# Prediction interface for Cog ⚙️
+# https://cog.run/python
+
+from cog import BasePredictor, Input, Path
+
+import os
+import re
+import torch
+import torchaudio
+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 librosa
+
+device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
+
+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
+
+
+class Predictor(BasePredictor):
+    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
+    def setup(self) -> None:
+        """Load the model into memory to make running multiple predictions efficient"""
+        # self.model = torch.load("./weights.pth")
+        print("Loading Whisper model...")
+        self.pipe = pipeline(
+            "automatic-speech-recognition",
+            model="openai/whisper-large-v3-turbo",
+            torch_dtype=torch.float16,
+            device=device,
+        )
+        print("Loading F5-TTS model...")
+
+        F5TTS_model_cfg = dict(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4)
+        self.F5TTS_ema_model, self.F5TTS_base_model = self.load_model("F5TTS_Base", DiT, F5TTS_model_cfg, 1200000)
+
+
+    def predict(
+        self,
+        gen_text: str = Input(description="Text to generate"),
+        ref_audio_orig: Path = Input(description="Reference audio"),
+        remove_silence: bool = Input(description="Remove silences", default=True),
+    ) -> Path:
+        """Run a single prediction on the model"""
+        model_choice = "F5-TTS"
+        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
+        ema_model = self.F5TTS_ema_model
+        base_model = self.F5TTS_base_model
+
+        if not ref_text.strip():
+            gr.Info("No reference text provided, transcribing reference audio...")
+            ref_text = outputs = self.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 F5-TTS")
+        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... This may take a moment")
+            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=".wav", delete=False) as tmp_wav:
+            wav_path = tmp_wav.name
+            torchaudio.save(wav_path, torch.tensor(generated_wave), target_sample_rate)
+
+        return wav_path

packages.txt

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

requirements.txt

@@ -0,0 +1,23 @@
+accelerate>=0.33.0
+cached_path
+click
+datasets
+einops>=0.8.0
+einx>=0.3.0
+ema_pytorch>=0.5.2
+gradio
+jieba
+librosa
+matplotlib
+numpy<=1.26.4
+pydub
+pypinyin
+safetensors
+soundfile
+tomli
+torchdiffeq
+tqdm>=4.65.0
+transformers
+vocos
+wandb
+x_transformers>=1.31.14

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}")