File size: 4,561 Bytes
380b209 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 | """Run the f0+intensity -> mel conformer.
The real module, loaded from the trained checkpoint. This is the path the
published numbers come from: `generate()` sizes its output from the real
input length and takes a padding mask, so a phrase is predicted at its own
duration.
The alternative was the ONNX export, whose input width is baked in at 2000
frames -- every phrase padded to ten seconds, with the decoder attending
over the padding. Measured against the training pipeline's own predictions
that is 0.9357 correlation against 0.9985 here, which is two orders of
magnitude more than the f0 extraction differs by, so it is not a reasonable
substitute when the point is reproducing the paper.
"""
from __future__ import annotations
import json
from pathlib import Path
import numpy as np
# Keys FeaturesToMelConv's constructor accepts. model_config.json carries
# training settings alongside them, and the model calls it mel_dim rather
# than mel_bins and takes no max_len -- so guessing the set does not work.
CONSTRUCTOR_KEYS = {
"input_dim", "d_model", "nhead", "num_encoder_layers",
"num_decoder_layers", "encoder_kernel_size", "conv_kernel_size",
"dim_feedforward", "dropout", "loss_type", "length_ratio",
"gradient_checkpointing", "decoder_only", "attn_window",
"mel_dim",
}
class Conformer:
"""Predicts a mel spectrogram from 200 Hz prosody features."""
def __init__(self, checkpoint: Path | str, cpu: bool = False) -> None:
"""`checkpoint` is a directory holding model_config.json and weights.
The weights may sit beside the config or in a `checkpoint-*`
subdirectory, which is how the training runs leave them.
"""
import torch
from safetensors.torch import load_file
from .conformer_model import FeaturesToMelConv
root = Path(checkpoint)
config_path = root / "model_config.json"
if not config_path.exists():
raise FileNotFoundError(
f"{config_path} not found. The conformer needs the directory "
f"holding model_config.json and its weights."
)
config = json.loads(config_path.read_text())
kwargs = {k: v for k, v in config.items() if k in CONSTRUCTOR_KEYS}
weights = sorted(root.glob("model.safetensors"))
weights += sorted(root.glob("checkpoint-*/model.safetensors"))
if not weights:
raise FileNotFoundError(f"no model.safetensors under {root}")
self._model = FeaturesToMelConv(**kwargs)
# strict=False because the checkpoint carries training-only buffers
# the inference classes do not declare. Checked against the ONNX
# export's output, which is what catches a genuinely wrong load.
self._model.load_state_dict(load_file(str(weights[-1])), strict=False)
self._model.eval()
self._device = torch.device(
"cpu" if cpu or not torch.cuda.is_available() else "cuda"
)
self._model.to(self._device)
self.mel_bins: int = int(config.get("mel_dim", config.get("mel_bins", 12)))
self.length_ratio: float = float(config.get("length_ratio", 0.5))
self._source = str(weights[-1])
@property
def provider(self) -> str:
return f"torch ({self._device.type})"
@property
def max_duration_s(self) -> float:
"""No fixed window: generate() sizes its output from the input."""
return float("inf")
def predict(self, features: np.ndarray) -> np.ndarray:
"""[T, 2] normalised features -> [T*length_ratio, mel_bins] mel.
The padding mask is all-false: a single phrase has nothing padded,
and passing the mask is what lets generate() size the output from
the real length rather than from a fixed window.
"""
import torch
features = np.asarray(features, dtype=np.float32)
if features.ndim != 2 or features.shape[1] != 2:
raise ValueError(f"expected [T, 2] features, got {features.shape}")
if features.size == 0:
return np.zeros((0, self.mel_bins), dtype=np.float32)
with torch.no_grad():
batch = torch.from_numpy(features).unsqueeze(0).to(self._device)
mask = torch.zeros(
(1, batch.shape[1]), dtype=torch.bool, device=self._device
)
out = self._model.generate(batch, input_padding_mask=mask)
mel = (out[0] if isinstance(out, tuple) else out)[0]
return mel.detach().cpu().numpy().astype(np.float32)
|