Upload conformer.py with huggingface_hub

conformer.py

@@ -0,0 +1,113 @@
+"""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)