Add args for Dataset and improve

.gitignore

@@ -17,3 +17,5 @@ venv/
 
 /scripts/test/
 *.zip
+*.csv
+*.bak

common/subprocess_utils.py

@@ -7,7 +7,7 @@ from .stdout_wrapper import SAFE_STDOUT
 python = sys.executable
 
 
-def run_script_with_log(cmd: list[str]) -> tuple[bool, str]:
+def run_script_with_log(cmd: list[str], ignore_warning=False) -> tuple[bool, str]:
     logger.info(f"Running: {' '.join(cmd)}")
     result = subprocess.run(
         [python] + cmd,
@@ -19,7 +19,7 @@ def run_script_with_log(cmd: list[str]) -> tuple[bool, str]:
         logger.error(f"Error: {' '.join(cmd)}")
         print(result.stderr)
         return False, result.stderr
-    elif result.stderr:
+    elif result.stderr and not ignore_warning:
         logger.warning(f"Warning: {' '.join(cmd)}")
         print(result.stderr)
         return True, result.stderr

common/tts_model.py

@@ -189,9 +189,9 @@ class ModelHolder:
         self, model_name: str, model_path: str
     ) -> tuple[gr.Dropdown, gr.Button]:
         if model_name not in self.model_files_dict:
-            raise Exception(f"モデル名{model_name}は存在しません")
+            raise ValueError(f"Model `{model_name}` is not found")
         if model_path not in self.model_files_dict[model_name]:
-            raise Exception(f"pthファイル{model_path}は存在しません")
+            raise ValueError(f"Model file `{model_path}` is not found")
         self.current_model = Model(
             model_path=model_path,
             config_path=os.path.join(self.root_dir, model_name, "config.json"),

preprocess_text.py

@@ -74,6 +74,7 @@ def preprocess(
                             print(
                                 f"An error occurred while generating the training set and validation set! Details:\n{e}"
                             )
+                            raise e
 
     transcription_path = cleaned_path
     spk_utt_map = defaultdict(list)

resample.py

@@ -13,6 +13,7 @@ from common.stdout_wrapper import SAFE_STDOUT
 
 DEFAULT_BLOCK_SIZE: float = 0.400  # seconds
 
+
 class BlockSizeException(Exception):
     pass
 
@@ -37,7 +38,9 @@ def process(item):
             try:
                 wav = normalize_audio(wav, sr)
             except BlockSizeException:
-                logger.info(f"Skip normalize due to less than {DEFAULT_BLOCK_SIZE} second audio: {wav_path}")
+                logger.info(
+                    f"Skip normalize due to less than {DEFAULT_BLOCK_SIZE} second audio: {wav_path}"
+                )
         if args.trim:
             wav, _ = librosa.effects.trim(wav, top_db=30)
         soundfile.write(os.path.join(args.out_dir, spkdir, wav_name), wav, sr)

slice.py

@@ -1,11 +1,13 @@
 import argparse
 import os
 import shutil
+from pathlib import Path
 
 import soundfile as sf
 import torch
 from tqdm import tqdm
 
+from common.log import logger
 from common.stdout_wrapper import SAFE_STDOUT
 from resample import normalize_audio
 
@@ -18,40 +20,53 @@ vad_model, utils = torch.hub.load(
 (get_speech_timestamps, _, read_audio, *_) = utils
 
 
-def get_stamps(audio_file, min_silence_dur_ms=700, min_sec=2):
+def get_stamps(
+    audio_file, min_silence_dur_ms: int = 700, min_sec: float = 2, max_sec: float = 12
+):
     """
-    min_silence_dur_ms:
+    min_silence_dur_ms: int (ミリ秒):
         このミリ秒数以上を無音だと判断する。
         逆に、この秒数以下の無音区間では区切られない。
         小さくすると、音声がぶつ切りに小さくなりすぎ、
         大きくすると音声一つ一つが長くなりすぎる。
         データセットによってたぶん要調整。
-    min_sec:
+    min_sec: float (秒):
-        この秒数より小さい発話は無視する。TTSのためには2秒未満は切り捨てたほうがいいかも。
+        この秒数より小さい発話は無視する。
+    max_sec: float (秒):
+        この秒数より大きい発話は無視する。
     """
 
     sampling_rate = 16000  # 16kHzか8kHzのみ対応
 
+    min_ms = int(min_sec * 1000)
+
     wav = read_audio(audio_file, sampling_rate=sampling_rate)
     speech_timestamps = get_speech_timestamps(
         wav,
         vad_model,
         sampling_rate=sampling_rate,
         min_silence_duration_ms=min_silence_dur_ms,
-        min_speech_duration_ms=min_sec * 1000,
+        min_speech_duration_ms=min_ms,
+        max_speech_duration_s=max_sec,
     )
 
     return speech_timestamps
 
 
 def split_wav(
-    audio_file, target_dir="raw", max_sec=12, min_silence_dur_ms=700, min_sec=2, normalize=False
+    audio_file,
+    target_dir="raw",
+    min_sec=2,
+    max_sec=12,
+    min_silence_dur_ms=700,
+    normalize=False,
 ):
     margin = 200  # ミリ秒単位で、音声の前後に余裕を持たせる
-    upper_bound_ms = max_sec * 1000  # これ以上の長さの音声は無視する
-
     speech_timestamps = get_stamps(
-        audio_file, min_silence_dur_ms=min_silence_dur_ms, min_sec=min_sec
+        audio_file,
+        min_silence_dur_ms=min_silence_dur_ms,
+        min_sec=min_sec,
+        max_sec=max_sec,
     )
 
     data, sr = sf.read(audio_file)
@@ -67,8 +82,6 @@ def split_wav(
     for i, ts in enumerate(speech_timestamps):
         start_ms = max(ts["start"] / 16 - margin, 0)
         end_ms = min(ts["end"] / 16 + margin, total_ms)
-        if end_ms - start_ms > upper_bound_ms:
-            continue
 
         start_sample = int(start_ms / 1000 * sr)
         end_sample = int(end_ms / 1000 * sr)
@@ -85,12 +98,36 @@ def split_wav(
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
-    parser.add_argument("--max_sec", "-M", type=int, default=12)
+    parser.add_argument(
-    parser.add_argument("--min_sec", "-m", type=int, default=2)
+        "--min_sec", "-m", type=float, default=2, help="Minimum seconds of a slice"
-    parser.add_argument("--min_silence_dur_ms", "-s", type=int, default=700)
+    )
-    parser.add_argument("--input_dir", "-i", type=str, default="inputs")
+    parser.add_argument(
-    parser.add_argument("--output_dir", "-t", type=str, default="raw")
+        "--max_sec", "-M", type=float, default=12, help="Maximum seconds of a slice"
-    parser.add_argument("--normalize", action="store_true")
+    )
+    parser.add_argument(
+        "--input_dir",
+        "-i",
+        type=str,
+        default="inputs",
+        help="Directory of input wav files",
+    )
+    parser.add_argument(
+        "--output_dir",
+        "-t",
+        type=str,
+        default="raw",
+        help="Directory of output wav files",
+    )
+    parser.add_argument(
+        "--normalize", action="store_true", help="Whether to normalize loudness"
+    )
+    parser.add_argument(
+        "--min_silence_dur_ms",
+        "-s",
+        type=int,
+        default=700,
+        help="Silence above this duration (ms) is considered as a split point.",
+    )
     args = parser.parse_args()
 
     input_dir = args.input_dir
@@ -100,25 +137,22 @@ if __name__ == "__main__":
     min_silence_dur_ms = args.min_silence_dur_ms
     normalize = args.normalize
 
-    wav_files = [
+    wav_files = Path(input_dir).glob("**/*.wav")
-        os.path.join(input_dir, f)
+    wav_files = list(wav_files)
-        for f in os.listdir(input_dir)
+    if os.path.exists(output_dir):
-        if f.lower().endswith(".wav")
+        logger.warning(f"Output directory {output_dir} already exists, deleting...")
-    ]
-    if os.path.exists(output_dir):  # ディレクトリを削除
-        print(f"{output_dir}フォルダが存在するので、削除します。")
         shutil.rmtree(output_dir)
 
     total_sec = 0
     for wav_file in tqdm(wav_files, file=SAFE_STDOUT):
         time_sec = split_wav(
-            wav_file,
+            audio_file=str(wav_file),
-            output_dir,
+            target_dir=output_dir,
-            max_sec=max_sec,
             min_sec=min_sec,
+            max_sec=max_sec,
             min_silence_dur_ms=min_silence_dur_ms,
-            normalize=normalize
+            normalize=normalize,
         )
         total_sec += time_sec
 
-    print(f"Done! Total time: {total_sec / 60:.2f} min.")
+    logger.info(f"Slice done! Total time: {total_sec / 60:.2f} min.")

transcribe.py

@@ -1,16 +1,17 @@
 import argparse
 import os
-import sys
 
 from faster_whisper import WhisperModel
 from tqdm import tqdm
 
+from common.constants import Languages
+from common.log import logger
 from common.stdout_wrapper import SAFE_STDOUT
 
 
-def transcribe(wav_path, initial_prompt=None):
+def transcribe(wav_path, initial_prompt=None, language="ja"):
     segments, _ = model.transcribe(
-        wav_path, beam_size=5, language="ja", initial_prompt=initial_prompt
+        wav_path, beam_size=5, language=language, initial_prompt=initial_prompt
     )
     texts = [segment.text for segment in segments]
     return "".join(texts)
@@ -23,8 +24,13 @@ if __name__ == "__main__":
     parser.add_argument(
         "--initial_prompt", type=str, default="こんにちは。元気、ですかー？私は……ちゃんと元気だよ！"
     )
-    parser.add_argument("--speaker_name", type=str, default=None, required=True)
+    parser.add_argument(
+        "--language", type=str, default="ja", choices=["ja", "en", "zh"]
+    )
+    parser.add_argument("--speaker_name", type=str, required=True)
     parser.add_argument("--model", type=str, default="large-v3")
+    parser.add_argument("--device", type=str, default="cuda")
+    parser.add_argument("--compute_type", type=str, default="bfloat16")
 
     args = parser.parse_args()
 
@@ -33,26 +39,40 @@ if __name__ == "__main__":
     input_dir = args.input_dir
     output_file = args.output_file
     initial_prompt = args.initial_prompt
+    language = args.language
+    device = args.device
+    compute_type = args.compute_type
 
+    logger.info(
+        f"Loading Whisper model ({args.model}) with compute_type={compute_type}"
+    )
     try:
-        model = WhisperModel("large-v3", device="cuda", compute_type="bfloat16")
+        model = WhisperModel(args.model, device=device, compute_type=compute_type)
-    except Exception as e:
+    except ValueError as e:
-        # Maybe bfloat16 is not supported (e.g. in colab)
+        logger.warning(f"Failed to load model: {e}")
-        # I don't know actually bf16 is better than fp16 or not...
+        model = WhisperModel(args.model, device=device)
-        model = WhisperModel("large-v3", device="cuda", compute_type="float16")
 
     wav_files = [
         os.path.join(input_dir, f) for f in os.listdir(input_dir) if f.endswith(".wav")
     ]
     if os.path.exists(output_file):
-        print(f"{output_file}が存在するので、バックアップを{output_file}.bakに作成します。")
+        logger.warning(f"{output_file} exists, backing up to {output_file}.bak")
         if os.path.exists(output_file + ".bak"):
-            print(f"{output_file}.bakも存在するので、削除します。")
+            logger.warning(f"{output_file}.bak exists, deleting...")
             os.remove(output_file + ".bak")
         os.rename(output_file, output_file + ".bak")
 
+    if language == "ja":
+        language = Languages.JP
+    elif language == "en":
+        language = Languages.EN
+    elif language == "zh":
+        language = Languages.ZH
+    else:
+        raise ValueError(f"{language} is not supported.")
     with open(output_file, "w", encoding="utf-8") as f:
         for wav_file in tqdm(wav_files, file=SAFE_STDOUT):
             file_name = os.path.basename(wav_file)
             text = transcribe(wav_file, initial_prompt=initial_prompt)
-            f.write(f"{file_name}|{speaker_name}|JP|{text}\n")
+            f.write(f"{file_name}|{speaker_name}|{language}|{text}\n")
+            f.flush()

webui_dataset.py

@@ -6,9 +6,8 @@ from common.log import logger
 from common.subprocess_utils import run_script_with_log, second_elem_of
 
 
-def do_slice(model_name, normalize):
+def do_slice(model_name: str, min_sec: float, max_sec: float, input_dir="inputs"):
     logger.info("Start slicing...")
-    input_dir = "inputs"
     output_dir = os.path.join("Data", model_name, "raw")
     cmd = [
         "slice.py",
@@ -16,16 +15,18 @@ def do_slice(model_name, normalize):
         input_dir,
         "--output_dir",
         output_dir,
+        "--min_sec",
+        str(min_sec),
+        "--max_sec",
+        str(max_sec),
     ]
-    if normalize:
+    success, message = run_script_with_log(cmd, ignore_warning=True)
-        cmd.append("--normalize")
-    success, message = run_script_with_log(cmd)
     if not success:
         return f"Error: {message}"
     return "音声のスライスが完了しました。"
 
 
-def do_transcribe(model_name):
+def do_transcribe(model_name, whisper_model, compute_type, language, initial_prompt):
     input_dir = os.path.join("Data", model_name, "raw")
     output_file = os.path.join("Data", model_name, "esd.list")
     result = run_script_with_log(
@@ -37,6 +38,14 @@ def do_transcribe(model_name):
             output_file,
             "--speaker_name",
             model_name,
+            "--model",
+            whisper_model,
+            "--compute_type",
+            compute_type,
+            "--language",
+            language,
+            "--initial_prompt",
+            initial_prompt,
         ]
     )
     return "音声の文字起こしが完了しました。"
@@ -45,8 +54,6 @@ def do_transcribe(model_name):
 initial_md = """
 # 簡易学習用データセット作成ツール
 
-**注意**：より精密で高品質なデータセットを作成したい・書き起こしをいろいろ修正したい場合は、[Aivis Dataset](https://github.com/litagin02/Aivis-Dataset)をおすすめします。書き起こし部分もかなり工夫されています。このツールはあくまでスライスして書き起こすという簡易的なことしかしていません。
-
 Style-Bert-VITS2の学習用データセットを作成するためのツールです。与えられた音声からちょうどいい長さの発話区間を切り取りスライスし、それぞれの音声に対して文字起こしを行います。
 
 ## 必要なもの
@@ -57,11 +64,13 @@ Style-Bert-VITS2の学習用データセットを作成するためのツール
 1. `inputs`フォルダ直下にwavファイルをすべて入れる
 2. `モデル名`を入力して、`音声のスライス`ボタンを押す
 3. 完了したら、`音声の文字起こし`ボタンを押す
+4. 出来上がった音声ファイルたちは`Data/{モデル名}/raw`に、書き起こしファイルは`Data/{モデル名}/esd.list`に保存されます。
 
-細かいパラメータ調整とかがしたい人は、`slice.py`と`transcribe.py`を眺めて直接実行してください。
+## 注意
 
-また、出来上がった音声ファイルたちは`Data/{モデル名}/raw`に、書き起こしファイルは`Data/{モデル名}/esd.list`に保存されます。
+- 長すぎる秒数（12-15秒くらいより長い？）のwavファイルは学習に用いられないようです。また短すぎてもあまりよくない可能性もあります。
-書き起こしの結果をどれだけ修正すればいいかはデータセットに依存しそうです。
+- 書き起こしの結果をどれだけ修正すればいいかはデータセットに依存しそうです。
+- 手動で書き起こしをいろいろ修正したり結果を細かく確認したい場合は、[Aivis Dataset](https://github.com/litagin02/Aivis-Dataset)もおすすめします。書き起こし部分もかなり工夫されています。ですがファイル数が多い場合などは、このツールで簡易的に切り出してデータセットを作るだけでも十分という気もしています。
 """
 
 with gr.Blocks(theme="NoCrypt/miku") as app:
@@ -70,20 +79,45 @@ with gr.Blocks(theme="NoCrypt/miku") as app:
     with gr.Accordion("音声のスライス"):
         with gr.Row():
             with gr.Column():
-                normalize = gr.Checkbox(label="スライスされた音声の音量を正規化する", value=True)
+                min_sec = gr.Slider(
+                    minimum=0, maximum=10, value=2, step=0.5, label="この秒数未満は切り捨てる"
+                )
+                max_sec = gr.Slider(
+                    minimum=0, maximum=15, value=12, step=0.5, label="この秒数以上は切り捨てる"
+                )
                 slice_button = gr.Button("スライスを実行")
             result1 = gr.Textbox(label="結果")
     with gr.Row():
+        with gr.Column():
+            whisper_model = gr.Radio(
+                ["tiny", "base", "small", "medium", "large", "large-v2", "large-v3"],
+                value="large-v3",
+            )
+            compute_type = gr.Dropdown(
+                [
+                    "int8",
+                    "int8_float32",
+                    "int8_float16",
+                    "int8_bfloat16",
+                    "int16",
+                    "float16",
+                    "bfloat16",
+                    "float32",
+                ],
+                value="bfloat16",
+            )
+            language = gr.Dropdown(["ja", "en", "zh"], value="ja")
+            initial_prompt = gr.Textbox(label="初期プロンプトを入力してください（省略可）")
         transcribe_button = gr.Button("音声の文字起こし")
         result2 = gr.Textbox(label="結果")
     slice_button.click(
         do_slice,
-        inputs=[model_name, normalize],
+        inputs=[model_name, min_sec, max_sec],
         outputs=[result1],
     )
     transcribe_button.click(
         do_transcribe,
-        inputs=[model_name],
+        inputs=[model_name, whisper_model, compute_type, language, initial_prompt],
         outputs=[result2],
     )