Refactor: typing and pathlib

slice.py

@@ -12,6 +12,8 @@ from style_bert_vits2.logging import logger
 from style_bert_vits2.utils.stdout_wrapper import SAFE_STDOUT
 
 
+# TODO: 並列処理による高速化
+
 vad_model, utils = torch.hub.load(
     repo_or_dir="snakers4/silero-vad",
     model="silero_vad",
@@ -23,7 +25,10 @@ vad_model, utils = torch.hub.load(
 
 
 def get_stamps(
-    audio_file, min_silence_dur_ms: int = 700, min_sec: float = 2, max_sec: float = 12
+    audio_file: Path,
+    min_silence_dur_ms: int = 700,
+    min_sec: float = 2,
+    max_sec: float = 12,
 ):
     """
     min_silence_dur_ms: int (ミリ秒):
@@ -42,7 +47,7 @@ def get_stamps(
 
     min_ms = int(min_sec * 1000)
 
-    wav = read_audio(audio_file, sampling_rate=sampling_rate)
+    wav = read_audio(str(audio_file), sampling_rate=sampling_rate)
     speech_timestamps = get_speech_timestamps(
         wav,
         vad_model,
@@ -56,13 +61,13 @@ def get_stamps(
 
 
 def split_wav(
-    audio_file,
+    audio_file: Path,
-    target_dir="raw",
+    target_dir: Path,
-    min_sec=2,
+    min_sec: float = 2,
-    max_sec=12,
+    max_sec: float = 12,
-    min_silence_dur_ms=700,
+    min_silence_dur_ms: int = 700,
-):
+) -> tuple[float, int]:
-    margin = 200  # ミリ秒単位で、音声の前後に余裕を持たせる
+    margin: int = 200  # ミリ秒単位で、音声の前後に余裕を持たせる
     speech_timestamps = get_stamps(
         audio_file,
         min_silence_dur_ms=min_silence_dur_ms,
@@ -74,10 +79,10 @@ def split_wav(
 
     total_ms = len(data) / sr * 1000
 
-    file_name = os.path.basename(audio_file).split(".")[0]
+    file_name = audio_file.stem
-    os.makedirs(target_dir, exist_ok=True)
+    target_dir.mkdir(parents=True, exist_ok=True)
 
-    total_time_ms = 0
+    total_time_ms: float = 0
     count = 0
 
     # タイムスタンプに従って分割し、ファイルに保存
@@ -89,7 +94,7 @@ def split_wav(
         end_sample = int(end_ms / 1000 * sr)
         segment = data[start_sample:end_sample]
 
-        sf.write(os.path.join(target_dir, f"{file_name}-{i}.wav"), segment, sr)
+        sf.write(str(target_dir / f"{file_name}-{i}.wav"), segment, sr)
         total_time_ms += end_ms - start_ms
         count += 1
 
@@ -126,20 +131,21 @@ if __name__ == "__main__":
     )
     args = parser.parse_args()
 
-    with open(os.path.join("configs", "paths.yml"), "r", encoding="utf-8") as f:
+    with open(Path("configs/paths.yml"), "r", encoding="utf-8") as f:
         path_config: dict[str, str] = yaml.safe_load(f.read())
         dataset_root = path_config["dataset_root"]
 
-    input_dir = args.input_dir
+    model_name = str(args.model_name)
-    output_dir = os.path.join(dataset_root, args.model_name, "raw")
+    input_dir = Path(args.input_dir)
-    min_sec = args.min_sec
+    output_dir = Path(dataset_root) / model_name / "raw"
-    max_sec = args.max_sec
+    min_sec: float = args.min_sec
-    min_silence_dur_ms = args.min_silence_dur_ms
+    max_sec: float = args.max_sec
+    min_silence_dur_ms: int = args.min_silence_dur_ms
 
     wav_files = Path(input_dir).glob("**/*.wav")
     wav_files = list(wav_files)
     logger.info(f"Found {len(wav_files)} wav files.")
-    if os.path.exists(output_dir):
+    if output_dir.exists():
         logger.warning(f"Output directory {output_dir} already exists, deleting...")
         shutil.rmtree(output_dir)
 
@@ -147,7 +153,7 @@ if __name__ == "__main__":
     total_count = 0
     for wav_file in tqdm(wav_files, file=SAFE_STDOUT):
         time_sec, count = split_wav(
-            audio_file=str(wav_file),
+            audio_file=wav_file,
             target_dir=output_dir,
             min_sec=min_sec,
             max_sec=max_sec,

style_gen.py

@@ -1,9 +1,11 @@
 import argparse
 import warnings
 from concurrent.futures import ThreadPoolExecutor
+from typing import Any
 
 import numpy as np
 import torch
+from numpy.typing import NDArray
 from tqdm import tqdm
 
 from config import config
@@ -11,11 +13,9 @@ from style_bert_vits2.logging import logger
 from style_bert_vits2.models.hyper_parameters import HyperParameters
 from style_bert_vits2.utils.stdout_wrapper import SAFE_STDOUT
 
-
 warnings.filterwarnings("ignore", category=UserWarning)
 from pyannote.audio import Inference, Model
 
-
 model = Model.from_pretrained("pyannote/wespeaker-voxceleb-resnet34-LM")
 inference = Inference(model, window="whole")
 device = torch.device(config.style_gen_config.device)
@@ -29,11 +29,11 @@ class NaNValueError(ValueError):
 
 
 # 推論時にインポートするために短いが関数を書く
-def get_style_vector(wav_path):
+def get_style_vector(wav_path: str) -> NDArray[Any]:
-    return inference(wav_path)
+    return inference(wav_path)  # type: ignore
 
 
-def save_style_vector(wav_path):
+def save_style_vector(wav_path: str):
     try:
         style_vec = get_style_vector(wav_path)
     except Exception as e:
@@ -48,20 +48,15 @@ def save_style_vector(wav_path):
     np.save(f"{wav_path}.npy", style_vec)  # `test.wav` -> `test.wav.npy`
 
 
-def process_line(line):
+def process_line(line: str):
-    wavname = line.split("|")[0]
+    wav_path = line.split("|")[0]
     try:
-        save_style_vector(wavname)
+        save_style_vector(wav_path)
         return line, None
     except NaNValueError:
         return line, "nan_error"
 
 
-def save_average_style_vector(style_vectors, filename="style_vectors.npy"):
-    average_vector = np.mean(style_vectors, axis=0)
-    np.save(filename, average_vector)
-
-
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument(
@@ -71,14 +66,14 @@ if __name__ == "__main__":
         "--num_processes", type=int, default=config.style_gen_config.num_processes
     )
     args, _ = parser.parse_known_args()
-    config_path = args.config
+    config_path: str = args.config
-    num_processes = args.num_processes
+    num_processes: int = args.num_processes
 
     hps = HyperParameters.load_from_json(config_path)
 
     device = config.style_gen_config.device
 
-    training_lines = []
+    training_lines: list[str] = []
     with open(hps.data.training_files, encoding="utf-8") as f:
         training_lines.extend(f.readlines())
     with ThreadPoolExecutor(max_workers=num_processes) as executor:
@@ -99,7 +94,7 @@ if __name__ == "__main__":
             f"Found NaN value in {len(nan_training_lines)} files: {nan_files}, so they will be deleted from training data."
         )
 
-    val_lines = []
+    val_lines: list[str] = []
     with open(hps.data.validation_files, encoding="utf-8") as f:
         val_lines.extend(f.readlines())
 

transcribe.py

@@ -2,6 +2,7 @@ import argparse
 import os
 import sys
 from pathlib import Path
+from typing import Optional
 
 import yaml
 from faster_whisper import WhisperModel
@@ -12,7 +13,9 @@ from style_bert_vits2.logging import logger
 from style_bert_vits2.utils.stdout_wrapper import SAFE_STDOUT
 
 
-def transcribe(wav_path: Path, initial_prompt=None, language="ja"):
+def transcribe(
+    wav_path: Path, initial_prompt: Optional[str] = None, language: str = "ja"
+):
     segments, _ = model.transcribe(
         str(wav_path), beam_size=5, language=language, initial_prompt=initial_prompt
     )
@@ -45,10 +48,10 @@ if __name__ == "__main__":
 
     input_dir = dataset_root / model_name / "raw"
     output_file = dataset_root / model_name / "esd.list"
-    initial_prompt = args.initial_prompt
+    initial_prompt: str = args.initial_prompt
-    language = args.language
+    language: str = args.language
-    device = args.device
+    device: str = args.device
-    compute_type = args.compute_type
+    compute_type: str = args.compute_type
 
     output_file.parent.mkdir(parents=True, exist_ok=True)