| """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 |
|
|
| |
| |
| |
| 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) |
| |
| |
| |
| 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) |
|
|
