Spaces:
Running
on
Zero
Running
on
Zero
File size: 6,689 Bytes
06cc563 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 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 |