Upload f5_tts/infer/infer_batch_parallel.py with huggingface_hub

f5_tts/infer/infer_batch_parallel.py

@@ -0,0 +1,171 @@
+import argparse
+import codecs
+import os
+import re
+from pathlib import Path
+
+import numpy as np
+import soundfile as sf
+import tomli
+from cached_path import cached_path
+import pandas as pd
+from tqdm import tqdm
+
+from f5_tts.infer.utils_infer import (
+    infer_process,
+    load_model,
+    load_vocoder,
+    preprocess_ref_audio_text,
+    remove_silence_for_generated_wav,
+)
+from f5_tts.model import DiT, UNetT
+
+
+def run_batch_inference(prompt_paths, prompt_texts, texts, languages, categories, model_obj, vocoder, mel_spec_type, remove_silence, speed, output_dir):
+    count = 0
+    for ref_audio in prompt_paths:
+        if not isinstance(ref_audio, str) or not os.path.isfile(ref_audio):
+            print(f"Invalid ref_audio: {ref_audio}")
+            count += 1
+            print(count)
+            # raise ValueError(f"Invalid ref_audio: {ref_audio}")
+
+    for idx, (ref_audio, ref_text, text_gen, language, category) in tqdm(enumerate(zip(prompt_paths, prompt_texts, texts, languages, categories))):
+        voices = {"main": {"ref_audio": ref_audio, "ref_text": ref_text}}
+        for voice in voices:
+            voices[voice]["ref_audio"], voices[voice]["ref_text"] = preprocess_ref_audio_text(
+                voices[voice]["ref_audio"], voices[voice]["ref_text"]
+            )
+            print("Voice:", voice)
+            print("Ref_audio:", voices[voice]["ref_audio"])
+            print("Ref_text:", voices[voice]["ref_text"])
+
+        generated_audio_segments = []
+        reg1 = r"(?=\[\w+\])"
+        chunks = re.split(reg1, text_gen)
+        reg2 = r"\[(\w+)\]"
+        for text in chunks:
+            if not text.strip():
+                continue
+            match = re.match(reg2, text)
+            if match:
+                voice = match[1]
+            else:
+                print("No voice tag found, using main.")
+                voice = "main"
+            if voice not in voices:
+                print(f"Voice {voice} not found, using main.")
+                voice = "main"
+            text = re.sub(reg2, "", text)
+            gen_text = text.strip()
+            ref_audio = voices[voice]["ref_audio"]
+            ref_text = voices[voice]["ref_text"]
+            print(f"Voice: {voice}")
+            audio, final_sample_rate, spectragram = infer_process(
+                ref_audio, ref_text, gen_text, model_obj, vocoder, mel_spec_type=mel_spec_type, speed=speed
+            )
+            generated_audio_segments.append(audio)
+
+        if generated_audio_segments:
+            final_wave = np.concatenate(generated_audio_segments)
+            filename = f"{language.upper()}_{category.upper()}_{idx}.wav"
+            outfile_dir = os.path.join(output_dir, language)
+            os.makedirs(outfile_dir, exist_ok=True)
+            wave_path = Path(outfile_dir) / filename
+            with open(wave_path, "wb") as f:
+                sf.write(f.name, final_wave, final_sample_rate)
+                if remove_silence:
+                    remove_silence_for_generated_wav(f.name)
+                print(f"Generated audio saved to: {f.name}")
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        prog="python3 infer-cli.py",
+        description="Commandline interface for E2/F5 TTS with Advanced Batch Processing.",
+        epilog="Specify options above to override one or more settings from config.",
+    )
+
+    parser.add_argument(
+        "-m",
+        "--model",
+        help="F5-TTS | E2-TTS",
+    )
+    parser.add_argument(
+        "-p",
+        "--ckpt_file",
+        help="The Checkpoint .pt",
+    )
+    parser.add_argument(
+        "-v",
+        "--vocab_file",
+        help="The vocab .txt",
+    )
+
+    parser.add_argument(
+        "-f",
+        "--generate_csv",
+        type=str,
+    )
+    parser.add_argument(
+        "-o",
+        "--output_dir",
+        type=str,
+        help="Path to output folder..",
+    )
+    parser.add_argument(
+        "--remove_silence",
+        help="Remove silence.",
+    )
+    parser.add_argument("--vocoder_name", type=str, default="vocos", choices=["vocos", "bigvgan"], help="vocoder name")
+    parser.add_argument(
+        "--load_vocoder_from_local",
+        action="store_true",
+        help="load vocoder from local. Default: ../checkpoints/charactr/vocos-mel-24khz",
+    )
+    parser.add_argument(
+        "--speed",
+        type=float,
+        default=1.0,
+        help="Adjust the speed of the audio generation (default: 1.0)",
+    )
+    args = parser.parse_args()
+
+    # Read texts and prompts to generate
+    filepath = args.generate_csv
+    df = pd.read_csv(filepath)
+    prompt_paths = df['prompt_path'].tolist()
+    prompt_texts = df['prompt_text'].tolist()
+    texts = df['text'].tolist()
+    languages = df['language'].tolist()
+    categories = df['category'].tolist()
+
+    # Model config
+    model = args.model
+    ckpt_file = args.ckpt_file
+    vocab_file = args.vocab_file
+    remove_silence = args.remove_silence
+    speed = args.speed
+    vocoder_name = args.vocoder_name
+    mel_spec_type = args.vocoder_name
+    if vocoder_name == "vocos":
+        vocoder_local_path = "../checkpoints/vocos-mel-24khz"
+    elif vocoder_name == "bigvgan":
+        vocoder_local_path = "../checkpoints/bigvgan_v2_24khz_100band_256x"
+    vocoder = load_vocoder(vocoder_name=mel_spec_type, is_local=args.load_vocoder_from_local, local_path=vocoder_local_path)
+
+    # load models
+    model_cls = DiT
+    model_cfg = dict(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4)
+    print(f"Using {model}...")
+    ema_model = load_model(model_cls, model_cfg, ckpt_file, mel_spec_type=mel_spec_type, vocab_file=vocab_file)
+
+    # Batch inference
+    output_dir = args.output_dir
+    if not os.path.exists(output_dir):
+        os.makedirs(output_dir)
+    run_batch_inference(prompt_paths, prompt_texts, texts, languages, categories, ema_model, vocoder, mel_spec_type, remove_silence, speed, output_dir)
+
+
+if __name__ == "__main__":
+    main()