|
import json |
|
from pathlib import Path |
|
from typing import Union |
|
|
|
import numpy as np |
|
|
|
from style_bert_vits2.constants import DEFAULT_STYLE |
|
from style_bert_vits2.logging import logger |
|
|
|
|
|
def save_neutral_vector( |
|
wav_dir: Union[Path, str], |
|
output_dir: Union[Path, str], |
|
config_path: Union[Path, str], |
|
config_output_path: Union[Path, str], |
|
): |
|
wav_dir = Path(wav_dir) |
|
output_dir = Path(output_dir) |
|
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) |
|
mean = np.mean(x, axis=0) |
|
only_mean = np.stack([mean]) |
|
np.save(output_dir / "style_vectors.npy", only_mean) |
|
logger.info(f"Saved mean style vector to {output_dir}") |
|
|
|
with open(config_path, 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(config_output_path, "w", encoding="utf-8") as f: |
|
json.dump(json_dict, f, indent=2, ensure_ascii=False) |
|
logger.info(f"Saved style config to {config_output_path}") |
|
|
|
|
|
def save_styles_by_dirs( |
|
wav_dir: Union[Path, str], |
|
output_dir: Union[Path, str], |
|
config_path: Union[Path, str], |
|
config_output_path: Union[Path, str], |
|
): |
|
wav_dir = Path(wav_dir) |
|
output_dir = Path(output_dir) |
|
output_dir.mkdir(parents=True, exist_ok=True) |
|
config_path = Path(config_path) |
|
config_output_path = Path(config_output_path) |
|
|
|
subdirs = [d for d in wav_dir.iterdir() if d.is_dir()] |
|
subdirs.sort() |
|
if len(subdirs) in (0, 1): |
|
logger.info( |
|
f"At least 2 subdirectories are required for generating style vectors with respect to them, found {len(subdirs)}." |
|
) |
|
logger.info("Generating only neutral style vector instead.") |
|
save_neutral_vector(wav_dir, output_dir, config_path, config_output_path) |
|
return |
|
|
|
|
|
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) |
|
mean = np.mean(x, axis=0) |
|
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) |
|
mean = np.mean(x, axis=0) |
|
style_vectors.append(mean) |
|
names.append(style_dir.name) |
|
|
|
|
|
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'}") |
|
|
|
|
|
style2id = {name: i for i, name in enumerate(names)} |
|
with open(config_path, encoding="utf-8") as f: |
|
json_dict = json.load(f) |
|
json_dict["data"]["num_styles"] = len(names) |
|
json_dict["data"]["style2id"] = style2id |
|
with open(config_output_path, "w", encoding="utf-8") as f: |
|
json.dump(json_dict, f, indent=2, ensure_ascii=False) |
|
logger.info(f"Saved style config to {config_output_path}") |
|
|