| 1 | import argparse |
| 2 | import codecs |
| 3 | import os |
| 4 | import re |
| 5 | from datetime import datetime |
| 6 | from importlib.resources import files |
| 7 | from pathlib import Path |
| 8 | |
| 9 | import numpy as np |
| 10 | import soundfile as sf |
| 11 | import tomli |
| 12 | from cached_path import cached_path |
| 13 | from hydra.utils import get_class |
| 14 | from omegaconf import OmegaConf |
| 15 | from unidecode import unidecode |
| 16 | |
| 17 | from f5_tts.infer.utils_infer import ( |
| 18 | cfg_strength, |
| 19 | cross_fade_duration, |
| 20 | device, |
| 21 | fix_duration, |
| 22 | infer_process, |
| 23 | load_model, |
| 24 | load_vocoder, |
| 25 | mel_spec_type, |
| 26 | nfe_step, |
| 27 | preprocess_ref_audio_text, |
| 28 | remove_silence_for_generated_wav, |
| 29 | speed, |
| 30 | sway_sampling_coef, |
| 31 | target_rms, |
| 32 | ) |
| 33 | |
| 34 | |
| 35 | parser = argparse.ArgumentParser( |
| 36 | prog="python3 infer-cli.py", |
| 37 | description="Commandline interface for E2/F5 TTS with Advanced Batch Processing.", |
| 38 | epilog="Specify options above to override one or more settings from config.", |
| 39 | ) |
| 40 | parser.add_argument( |
| 41 | "-c", |
| 42 | "--config", |
| 43 | type=str, |
| 44 | default=os.path.join(files("f5_tts").joinpath("infer/examples/basic"), "basic.toml"), |
| 45 | help="The configuration file, default see infer/examples/basic/basic.toml", |
| 46 | ) |
| 47 | |
| 48 | |
| 49 | # Note. Not to provide default value here in order to read default from config file |
| 50 | |
| 51 | parser.add_argument( |
| 52 | "-m", |
| 53 | "--model", |
| 54 | type=str, |
| 55 | help="The model name: F5TTS_v1_Base | F5TTS_Base | E2TTS_Base | etc.", |
| 56 | ) |
| 57 | parser.add_argument( |
| 58 | "-mc", |
| 59 | "--model_cfg", |
| 60 | type=str, |
| 61 | help="The path to F5-TTS model config file .yaml", |
| 62 | ) |
| 63 | parser.add_argument( |
| 64 | "-p", |
| 65 | "--ckpt_file", |
| 66 | type=str, |
| 67 | help="The path to model checkpoint .pt, leave blank to use default", |
| 68 | ) |
| 69 | parser.add_argument( |
| 70 | "-v", |
| 71 | "--vocab_file", |
| 72 | type=str, |
| 73 | help="The path to vocab file .txt, leave blank to use default", |
| 74 | ) |
| 75 | parser.add_argument( |
| 76 | "-r", |
| 77 | "--ref_audio", |
| 78 | type=str, |
| 79 | help="The reference audio file.", |
| 80 | ) |
| 81 | parser.add_argument( |
| 82 | "-s", |
| 83 | "--ref_text", |
| 84 | type=str, |
| 85 | help="The transcript/subtitle for the reference audio", |
| 86 | ) |
| 87 | parser.add_argument( |
| 88 | "-t", |
| 89 | "--gen_text", |
| 90 | type=str, |
| 91 | help="The text to make model synthesize a speech", |
| 92 | ) |
| 93 | parser.add_argument( |
| 94 | "-f", |
| 95 | "--gen_file", |
| 96 | type=str, |
| 97 | help="The file with text to generate, will ignore --gen_text", |
| 98 | ) |
| 99 | parser.add_argument( |
| 100 | "-o", |
| 101 | "--output_dir", |
| 102 | type=str, |
| 103 | help="The path to output folder", |
| 104 | ) |
| 105 | parser.add_argument( |
| 106 | "-w", |
| 107 | "--output_file", |
| 108 | type=str, |
| 109 | help="The name of output file", |
| 110 | ) |
| 111 | parser.add_argument( |
| 112 | "--save_chunk", |
| 113 | action="store_true", |
| 114 | help="To save each audio chunks during inference", |
| 115 | ) |
| 116 | parser.add_argument( |
| 117 | "--no_legacy_text", |
| 118 | action="store_false", |
| 119 | help="Not to use lossy ASCII transliterations of unicode text in saved file names.", |
| 120 | ) |
| 121 | parser.add_argument( |
| 122 | "--remove_silence", |
| 123 | action="store_true", |
| 124 | help="To remove long silence found in ouput", |
| 125 | ) |
| 126 | parser.add_argument( |
| 127 | "--load_vocoder_from_local", |
| 128 | action="store_true", |
| 129 | help="To load vocoder from local dir, default to ../checkpoints/vocos-mel-24khz", |
| 130 | ) |
| 131 | parser.add_argument( |
| 132 | "--vocoder_name", |
| 133 | type=str, |
| 134 | choices=["vocos", "bigvgan"], |
| 135 | help=f"Used vocoder name: vocos | bigvgan, default {mel_spec_type}", |
| 136 | ) |
| 137 | parser.add_argument( |
| 138 | "--target_rms", |
| 139 | type=float, |
| 140 | help=f"Target output speech loudness normalization value, default {target_rms}", |
| 141 | ) |
| 142 | parser.add_argument( |
| 143 | "--cross_fade_duration", |
| 144 | type=float, |
| 145 | help=f"Duration of cross-fade between audio segments in seconds, default {cross_fade_duration}", |
| 146 | ) |
| 147 | parser.add_argument( |
| 148 | "--nfe_step", |
| 149 | type=int, |
| 150 | help=f"The number of function evaluation (denoising steps), default {nfe_step}", |
| 151 | ) |
| 152 | parser.add_argument( |
| 153 | "--cfg_strength", |
| 154 | type=float, |
| 155 | help=f"Classifier-free guidance strength, default {cfg_strength}", |
| 156 | ) |
| 157 | parser.add_argument( |
| 158 | "--sway_sampling_coef", |
| 159 | type=float, |
| 160 | help=f"Sway Sampling coefficient, default {sway_sampling_coef}", |
| 161 | ) |
| 162 | parser.add_argument( |
| 163 | "--speed", |
| 164 | type=float, |
| 165 | help=f"The speed of the generated audio, default {speed}", |
| 166 | ) |
| 167 | parser.add_argument( |
| 168 | "--fix_duration", |
| 169 | type=float, |
| 170 | help=f"Fix the total duration (ref and gen audios) in seconds, default {fix_duration}", |
| 171 | ) |
| 172 | parser.add_argument( |
| 173 | "--device", |
| 174 | type=str, |
| 175 | help="Specify the device to run on", |
| 176 | ) |
| 177 | args = parser.parse_args() |
| 178 | |
| 179 | |
| 180 | # config file |
| 181 | |
| 182 | config = tomli.load(open(args.config, "rb")) |
| 183 | |
| 184 | |
| 185 | # command-line interface parameters |
| 186 | |
| 187 | model = args.model or config.get("model", "F5TTS_v1_Base") |
| 188 | ckpt_file = args.ckpt_file or config.get("ckpt_file", "") |
| 189 | vocab_file = args.vocab_file or config.get("vocab_file", "") |
| 190 | |
| 191 | ref_audio = args.ref_audio or config.get("ref_audio", "infer/examples/basic/basic_ref_en.wav") |
| 192 | ref_text = ( |
| 193 | args.ref_text |
| 194 | if args.ref_text is not None |
| 195 | else config.get("ref_text", "Some call me nature, others call me mother nature.") |
| 196 | ) |
| 197 | gen_text = args.gen_text or config.get("gen_text", "Here we generate something just for test.") |
| 198 | gen_file = args.gen_file or config.get("gen_file", "") |
| 199 | |
| 200 | output_dir = args.output_dir or config.get("output_dir", "tests") |
| 201 | output_file = args.output_file or config.get( |
| 202 | "output_file", f"infer_cli_{datetime.now().strftime(r'%Y%m%d_%H%M%S')}.wav" |
| 203 | ) |
| 204 | |
| 205 | save_chunk = args.save_chunk or config.get("save_chunk", False) |
| 206 | use_legacy_text = args.no_legacy_text or config.get("no_legacy_text", False) # no_legacy_text is a store_false arg |
| 207 | if save_chunk and use_legacy_text: |
| 208 | print( |
| 209 | "\nWarning to --save_chunk: lossy ASCII transliterations of unicode text for legacy (.wav) file names, --no_legacy_text to disable.\n" |
| 210 | ) |
| 211 | |
| 212 | remove_silence = args.remove_silence or config.get("remove_silence", False) |
| 213 | load_vocoder_from_local = args.load_vocoder_from_local or config.get("load_vocoder_from_local", False) |
| 214 | |
| 215 | vocoder_name = args.vocoder_name or config.get("vocoder_name", mel_spec_type) |
| 216 | target_rms = args.target_rms or config.get("target_rms", target_rms) |
| 217 | cross_fade_duration = args.cross_fade_duration or config.get("cross_fade_duration", cross_fade_duration) |
| 218 | nfe_step = args.nfe_step or config.get("nfe_step", nfe_step) |
| 219 | cfg_strength = args.cfg_strength or config.get("cfg_strength", cfg_strength) |
| 220 | sway_sampling_coef = args.sway_sampling_coef or config.get("sway_sampling_coef", sway_sampling_coef) |
| 221 | speed = args.speed or config.get("speed", speed) |
| 222 | fix_duration = args.fix_duration or config.get("fix_duration", fix_duration) |
| 223 | device = args.device or config.get("device", device) |
| 224 | |
| 225 | |
| 226 | # patches for pip pkg user |
| 227 | if "infer/examples/" in ref_audio: |
| 228 | ref_audio = str(files("f5_tts").joinpath(f"{ref_audio}")) |
| 229 | if "infer/examples/" in gen_file: |
| 230 | gen_file = str(files("f5_tts").joinpath(f"{gen_file}")) |
| 231 | if "voices" in config: |
| 232 | for voice in config["voices"]: |
| 233 | voice_ref_audio = config["voices"][voice]["ref_audio"] |
| 234 | if "infer/examples/" in voice_ref_audio: |
| 235 | config["voices"][voice]["ref_audio"] = str(files("f5_tts").joinpath(f"{voice_ref_audio}")) |
| 236 | |
| 237 | |
| 238 | # ignore gen_text if gen_file provided |
| 239 | |
| 240 | if gen_file: |
| 241 | gen_text = codecs.open(gen_file, "r", "utf-8").read() |
| 242 | |
| 243 | |
| 244 | # output path |
| 245 | |
| 246 | wave_path = Path(output_dir) / output_file |
| 247 | # spectrogram_path = Path(output_dir) / "infer_cli_out.png" |
| 248 | if save_chunk: |
| 249 | output_chunk_dir = os.path.join(output_dir, f"{Path(output_file).stem}_chunks") |
| 250 | if not os.path.exists(output_chunk_dir): |
| 251 | os.makedirs(output_chunk_dir) |
| 252 | |
| 253 | |
| 254 | # load vocoder |
| 255 | |
| 256 | if vocoder_name == "vocos": |
| 257 | vocoder_local_path = "../checkpoints/vocos-mel-24khz" |
| 258 | elif vocoder_name == "bigvgan": |
| 259 | vocoder_local_path = "../checkpoints/bigvgan_v2_24khz_100band_256x" |
| 260 | |
| 261 | vocoder = load_vocoder( |
| 262 | vocoder_name=vocoder_name, is_local=load_vocoder_from_local, local_path=vocoder_local_path, device=device |
| 263 | ) |
| 264 | |
| 265 | |
| 266 | # load TTS model |
| 267 | |
| 268 | model_cfg = OmegaConf.load( |
| 269 | args.model_cfg or config.get("model_cfg", str(files("f5_tts").joinpath(f"configs/{model}.yaml"))) |
| 270 | ) |
| 271 | model_cls = get_class(f"f5_tts.model.{model_cfg.model.backbone}") |
| 272 | model_arc = model_cfg.model.arch |
| 273 | |
| 274 | repo_name, ckpt_step, ckpt_type = "F5-TTS", 1250000, "safetensors" |
| 275 | |
| 276 | if model != "F5TTS_Base": |
| 277 | assert vocoder_name == model_cfg.model.mel_spec.mel_spec_type |
| 278 | |
| 279 | # override for previous models |
| 280 | if model == "F5TTS_Base": |
| 281 | if vocoder_name == "vocos": |
| 282 | ckpt_step = 1200000 |
| 283 | elif vocoder_name == "bigvgan": |
| 284 | model = "F5TTS_Base_bigvgan" |
| 285 | ckpt_type = "pt" |
| 286 | elif model == "E2TTS_Base": |
| 287 | repo_name = "E2-TTS" |
| 288 | ckpt_step = 1200000 |
| 289 | |
| 290 | if not ckpt_file: |
| 291 | ckpt_file = str(cached_path(f"hf://SWivid/{repo_name}/{model}/model_{ckpt_step}.{ckpt_type}")) |
| 292 | elif ckpt_file.startswith("hf://"): |
| 293 | ckpt_file = str(cached_path(ckpt_file)) |
| 294 | |
| 295 | if vocab_file.startswith("hf://"): |
| 296 | vocab_file = str(cached_path(vocab_file)) |
| 297 | |
| 298 | print(f"Using {model}...") |
| 299 | ema_model = load_model( |
| 300 | model_cls, model_arc, ckpt_file, mel_spec_type=vocoder_name, vocab_file=vocab_file, device=device |
| 301 | ) |
| 302 | |
| 303 | |
| 304 | # inference process |
| 305 | |
| 306 | |
| 307 | def main(): |
| 308 | main_voice = {"ref_audio": ref_audio, "ref_text": ref_text} |
| 309 | if "voices" not in config: |
| 310 | voices = {"main": main_voice} |
| 311 | else: |
| 312 | voices = config["voices"] |
| 313 | voices["main"] = main_voice |
| 314 | for voice in voices: |
| 315 | print("Voice:", voice) |
| 316 | print("ref_audio ", voices[voice]["ref_audio"]) |
| 317 | voices[voice]["ref_audio"], voices[voice]["ref_text"] = preprocess_ref_audio_text( |
| 318 | voices[voice]["ref_audio"], voices[voice]["ref_text"] |
| 319 | ) |
| 320 | print("ref_audio_", voices[voice]["ref_audio"], "\n\n") |
| 321 | |
| 322 | generated_audio_segments = [] |
| 323 | reg1 = r"(?=\[\w+\])" |
| 324 | chunks = re.split(reg1, gen_text) |
| 325 | reg2 = r"\[(\w+)\]" |
| 326 | for text in chunks: |
| 327 | if not text.strip(): |
| 328 | continue |
| 329 | match = re.match(reg2, text) |
| 330 | if match: |
| 331 | voice = match[1] |
| 332 | else: |
| 333 | print("No voice tag found, using main.") |
| 334 | voice = "main" |
| 335 | if voice not in voices: |
| 336 | print(f"Voice {voice} not found, using main.") |
| 337 | voice = "main" |
| 338 | text = re.sub(reg2, "", text) |
| 339 | ref_audio_ = voices[voice]["ref_audio"] |
| 340 | ref_text_ = voices[voice]["ref_text"] |
| 341 | local_speed = voices[voice].get("speed", speed) |
| 342 | gen_text_ = text.strip() |
| 343 | print(f"Voice: {voice}") |
| 344 | audio_segment, final_sample_rate, spectrogram = infer_process( |
| 345 | ref_audio_, |
| 346 | ref_text_, |
| 347 | gen_text_, |
| 348 | ema_model, |
| 349 | vocoder, |
| 350 | mel_spec_type=vocoder_name, |
| 351 | target_rms=target_rms, |
| 352 | cross_fade_duration=cross_fade_duration, |
| 353 | nfe_step=nfe_step, |
| 354 | cfg_strength=cfg_strength, |
| 355 | sway_sampling_coef=sway_sampling_coef, |
| 356 | speed=local_speed, |
| 357 | fix_duration=fix_duration, |
| 358 | device=device, |
| 359 | ) |
| 360 | generated_audio_segments.append(audio_segment) |
| 361 | |
| 362 | if save_chunk: |
| 363 | if len(gen_text_) > 200: |
| 364 | gen_text_ = gen_text_[:200] + " ... " |
| 365 | if use_legacy_text: |
| 366 | gen_text_ = unidecode(gen_text_) |
| 367 | sf.write( |
| 368 | os.path.join(output_chunk_dir, f"{len(generated_audio_segments) - 1}_{gen_text_}.wav"), |
| 369 | audio_segment, |
| 370 | final_sample_rate, |
| 371 | ) |
| 372 | |
| 373 | if generated_audio_segments: |
| 374 | final_wave = np.concatenate(generated_audio_segments) |
| 375 | |
| 376 | if not os.path.exists(output_dir): |
| 377 | os.makedirs(output_dir) |
| 378 | |
| 379 | with open(wave_path, "wb") as f: |
| 380 | sf.write(f.name, final_wave, final_sample_rate) |
| 381 | # Remove silence |
| 382 | if remove_silence: |
| 383 | remove_silence_for_generated_wav(f.name) |
| 384 | print(f.name) |
| 385 | |
| 386 | |
| 387 | if __name__ == "__main__": |
| 388 | main() |
| 389 |