98 lines
3.5 KiB
Python
98 lines
3.5 KiB
Python
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 set_style_config(json_path: Path, output_path: Path):
|
|
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}
|
|
with open(output_path, "w", encoding="utf-8") as f:
|
|
json.dump(json_dict, f, indent=2, ensure_ascii=False)
|
|
logger.info(f"Save style config (only {DEFAULT_STYLE}) to {output_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)
|
|
embs.append(np.expand_dims(xvec, axis=0))
|
|
|
|
x = np.concatenate(embs, axis=0) # (N, 256)
|
|
mean = np.mean(x, axis=0) # (256,)
|
|
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, 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.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)
|
|
|
|
# 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, 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}")
|