mirror of
https://github.com/storytold/storyteller-ml.git
synced 2026-10-16 00:13:07 +00:00
Paul Sajna
f86b44d8fc
* 48KHz softvc * add some notes * almost forgot the inference script * fix readme link * requirements.txt
76 lines
2.4 KiB
Python
76 lines
2.4 KiB
Python
import argparse
|
|
|
|
from pathlib import Path
|
|
import numpy as np
|
|
from tqdm import tqdm
|
|
|
|
import torch
|
|
|
|
from acoustic.model import AcousticModel
|
|
from torch.nn.modules.utils import consume_prefix_in_state_dict_if_present
|
|
|
|
|
|
def repeat_expand_2d(content: torch.Tensor, target_len: int) -> torch.Tensor:
|
|
# content : [h, t]
|
|
src_len = content.shape[-1]
|
|
if target_len < src_len:
|
|
return content[:, :target_len]
|
|
else:
|
|
return torch.nn.functional.interpolate(
|
|
content.unsqueeze(0), size=target_len, mode="nearest"
|
|
).squeeze(0)
|
|
|
|
|
|
def generate(args):
|
|
print("Loading acoustic model checkpoint")
|
|
#acoustic = torch.hub.load("bshall/acoustic-model:main", f"hubert_{args.model}").cuda()
|
|
acoustic = AcousticModel().cuda().eval()
|
|
checkpoint = torch.load("exp/rue/model-2000.pt", map_location="cuda")["acoustic-model"]
|
|
consume_prefix_in_state_dict_if_present(checkpoint, "module.")
|
|
acoustic.load_state_dict(checkpoint)
|
|
|
|
print(f"Generating from {args.in_dir} -> {args.out_dir}")
|
|
for path in tqdm(list(args.in_dir.rglob("*.npy"))):
|
|
units = np.load(path)
|
|
units_dtype = torch.long if args.model == "discrete" else torch.float
|
|
units = torch.tensor(units, dtype=units_dtype).cuda()
|
|
units = units.T
|
|
print(units.shape)
|
|
target_len = units.shape[1]*3
|
|
units = units.squeeze(0)
|
|
units = repeat_expand_2d(units, target_len).T.unsqueeze(0)
|
|
|
|
with torch.inference_mode():
|
|
mel_ = acoustic.generate(units)
|
|
mel_ = mel_.transpose(1, 2)
|
|
|
|
out_path = args.out_dir / path.relative_to(args.in_dir)
|
|
out_path.parent.mkdir(exist_ok=True, parents=True)
|
|
np.save(out_path.with_suffix(".npy"), mel_.squeeze().cpu().numpy())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(
|
|
description="Generate spectrograms from input speech units (discrete or soft)."
|
|
)
|
|
parser.add_argument(
|
|
"model",
|
|
help="available models (HuBERT-Soft or HuBERT-Discrete)",
|
|
choices=["soft", "discrete"],
|
|
)
|
|
parser.add_argument(
|
|
"in_dir",
|
|
metavar="in-dir",
|
|
help="path to the dataset directory.",
|
|
type=Path,
|
|
)
|
|
parser.add_argument(
|
|
"out_dir",
|
|
metavar="out-dir",
|
|
help="path to the output directory.",
|
|
type=Path,
|
|
)
|
|
args = parser.parse_args()
|
|
generate(args)
|
|
|