Feat: make style vectors w.r.t. subdirs structure

default_style.py

@@ -9,7 +9,7 @@ from style_bert_vits2.logging import logger
 
 
 def set_style_config(json_path: Path, output_path: Path):
-    with open(json_path, "r", encoding="utf-8") as f:
+    with open(json_path, encoding="utf-8") as f:
         json_dict = json.load(f)
     json_dict["data"]["num_styles"] = 1
     json_dict["data"]["style2id"] = {DEFAULT_STYLE: 0}
@@ -21,6 +21,7 @@ def set_style_config(json_path: Path, output_path: Path):
 def save_neutral_vector(wav_dir: Union[Path, str], output_path: Union[Path, str]):
     wav_dir = Path(wav_dir)
     output_path = Path(output_path)
+    json_path = output_path / "config.json"
     embs = []
     for file in wav_dir.rglob("*.npy"):
         xvec = np.load(file)
@@ -31,3 +32,63 @@ def save_neutral_vector(wav_dir: Union[Path, str], output_path: Union[Path, str]
     only_mean = np.stack([mean])  # (1, 256)
     np.save(output_path, only_mean)
     logger.info(f"Saved mean style vector to {output_path}")
+
+    with open(json_path, "r", encoding="utf-8") as f:
+        json_dict = json.load(f)
+    json_dict["data"]["num_styles"] = 1
+    json_dict["data"]["style2id"] = {DEFAULT_STYLE: 0}
+    with open(json_path, "w", encoding="utf-8") as f:
+        json.dump(json_dict, f, indent=2, ensure_ascii=False)
+    logger.info(f"Saved style config to {json_path}")
+
+
+def save_styles_by_dirs(wav_dir: Union[Path, str], output_dir: Union[Path, str]):
+    wav_dir = Path(wav_dir)
+    output_dir = Path(output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    json_path = output_dir / "config.json"
+
+    subdirs = [d for d in wav_dir.iterdir() if d.is_dir()]
+    subdirs.sort()
+    if len(subdirs) in (0, 1):
+        logger.warning("No style directories found. Saving only neutral style.")
+        save_neutral_vector(wav_dir, output_dir)
+
+    # First get mean of all for Neutral
+    embs = []
+    for file in wav_dir.rglob("*.npy"):
+        xvec = np.load(file)
+        embs.append(np.expand_dims(xvec, axis=0))
+    x = np.concatenate(embs, axis=0)  # (N, 256)
+    mean = np.mean(x, axis=0)  # (256,)
+    style_vectors = [mean]
+
+    names = [DEFAULT_STYLE]
+    for style_dir in subdirs:
+        npy_files = list(style_dir.rglob("*.npy"))
+        if not npy_files:
+            continue
+        embs = []
+        for file in npy_files:
+            xvec = np.load(file)
+            embs.append(np.expand_dims(xvec, axis=0))
+
+        x = np.concatenate(embs, axis=0)  # (N, 256)
+        mean = np.mean(x, axis=0)  # (256,)
+        style_vectors.append(mean)
+        names.append(style_dir.name)
+
+    # Stack them to make (num_styles, 256)
+    style_vectors_npy = np.stack(style_vectors, axis=0)
+    np.save(output_dir / "style_vectors.npy", style_vectors_npy)
+    logger.info(f"Saved style vectors to {output_dir / 'style_vectors.npy'}")
+
+    # Save style2id config to json
+    style2id = {name: i for i, name in enumerate(names)}
+    with open(json_path, "r", encoding="utf-8") as f:
+        json_dict = json.load(f)
+    json_dict["data"]["num_styles"] = len(names)
+    json_dict["data"]["style2id"] = style2id
+    with open(json_path, "w", encoding="utf-8") as f:
+        json.dump(json_dict, f, indent=2, ensure_ascii=False)
+    logger.info(f"Saved style config to {json_path}")