Fix only-Neutral-style saving logic

This commit is contained in:
litagin02
2024-06-01 15:51:27 +09:00
parent 385181465f
commit 376ca08db8

View File

@@ -8,20 +8,14 @@ from style_bert_vits2.constants import DEFAULT_STYLE
from style_bert_vits2.logging import logger from style_bert_vits2.logging import logger
def set_style_config(json_path: Path, output_path: Path): def save_neutral_vector(
with open(json_path, encoding="utf-8") as f: wav_dir: Union[Path, str],
json_dict = json.load(f) output_dir: Union[Path, str],
json_dict["data"]["num_styles"] = 1 config_path: Union[Path, str],
json_dict["data"]["style2id"] = {DEFAULT_STYLE: 0} config_output_path: Union[Path, str],
with open(output_path, "w", encoding="utf-8") as f: ):
json.dump(json_dict, f, indent=2, ensure_ascii=False)
logger.info(f"Saving 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) wav_dir = Path(wav_dir)
output_path = Path(output_path) output_dir = Path(output_dir)
json_path = output_path / "config.json"
embs = [] embs = []
for file in wav_dir.rglob("*.npy"): for file in wav_dir.rglob("*.npy"):
xvec = np.load(file) xvec = np.load(file)
@@ -30,16 +24,16 @@ def save_neutral_vector(wav_dir: Union[Path, str], output_path: Union[Path, str]
x = np.concatenate(embs, axis=0) # (N, 256) x = np.concatenate(embs, axis=0) # (N, 256)
mean = np.mean(x, axis=0) # (256,) mean = np.mean(x, axis=0) # (256,)
only_mean = np.stack([mean]) # (1, 256) only_mean = np.stack([mean]) # (1, 256)
np.save(output_path, only_mean) np.save(output_dir / "style_vectors.npy", only_mean)
logger.info(f"Saved mean style vector to {output_path}") logger.info(f"Saved mean style vector to {output_dir}")
with open(json_path, encoding="utf-8") as f: with open(config_path, encoding="utf-8") as f:
json_dict = json.load(f) json_dict = json.load(f)
json_dict["data"]["num_styles"] = 1 json_dict["data"]["num_styles"] = 1
json_dict["data"]["style2id"] = {DEFAULT_STYLE: 0} json_dict["data"]["style2id"] = {DEFAULT_STYLE: 0}
with open(json_path, "w", encoding="utf-8") as f: with open(config_output_path, "w", encoding="utf-8") as f:
json.dump(json_dict, f, indent=2, ensure_ascii=False) json.dump(json_dict, f, indent=2, ensure_ascii=False)
logger.info(f"Saved style config to {json_path}") logger.info(f"Saved style config to {config_output_path}")
def save_styles_by_dirs( def save_styles_by_dirs(
@@ -61,8 +55,8 @@ def save_styles_by_dirs(
f"At least 2 subdirectories are required for generating style vectors with respect to them, found {len(subdirs)}." 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.") logger.info("Generating only neutral style vector instead.")
set_style_config(config_path, config_output_path) save_neutral_vector(wav_dir, output_dir, config_path, config_output_path)
save_neutral_vector(wav_dir, output_dir) return
# First get mean of all for Neutral # First get mean of all for Neutral
embs = [] embs = []