feat: add EnCodec and Mimi codec implementations with self-registration

compare_codec/__init__.py

@@ -47,4 +47,6 @@ def get_all() -> dict[str, AudioCodec]:
 
 # Import codec modules so they self-register on startup.
 from compare_codec import dac as _dac  # noqa: E402, F401
+from compare_codec import encodec_codec as _encodec  # noqa: E402, F401
+from compare_codec import mimi_codec as _mimi  # noqa: E402, F401
 from compare_codec import snac_codec as _snac  # noqa: E402, F401

compare_codec/encodec_codec.py

@@ -0,0 +1,79 @@
+"""EnCodec (Meta) — wraps the HuggingFace transformers implementation."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+import torch
+import torchaudio
+
+from compare_codec import CodecConfig, register
+
+_BANDWIDTHS = [1.5, 3.0, 6.0, 12.0, 24.0]
+
+
+class EnCodecCodec:
+    """EnCodec 24kHz codec with lazy model loading."""
+
+    def __init__(self) -> None:
+        self._model = None
+        self._processor = None
+
+    @property
+    def name(self) -> str:
+        return "EnCodec"
+
+    @property
+    def sample_rate(self) -> int:
+        return 24_000
+
+    def configs(self) -> list[CodecConfig]:
+        return [
+            CodecConfig(
+                name=f"{bw:g} kbps",
+                params={"bandwidth": bw, "sample_rate": 24_000},
+            )
+            for bw in _BANDWIDTHS
+        ]
+
+    def _load(self):
+        if self._model is None:
+            from transformers import AutoProcessor, EncodecModel
+
+            self._model = EncodecModel.from_pretrained("facebook/encodec_24khz")
+            self._model.eval()
+            self._processor = AutoProcessor.from_pretrained("facebook/encodec_24khz")
+
+    @torch.no_grad()
+    def encode_decode(self, audio_path: Path, config: CodecConfig) -> np.ndarray:
+        self._load()
+        bandwidth: float = config.params["bandwidth"]
+        target_sr: int = config.params["sample_rate"]
+
+        wav, sr = torchaudio.load(str(audio_path))
+        if wav.shape[0] > 1:
+            wav = wav.mean(dim=0, keepdim=True)
+        if sr != target_sr:
+            wav = torchaudio.functional.resample(wav, sr, target_sr)
+
+        inputs = self._processor(
+            raw_audio=wav.squeeze(0).numpy(),
+            sampling_rate=target_sr,
+            return_tensors="pt",
+        )
+        enc = self._model.encode(
+            inputs["input_values"],
+            inputs["padding_mask"],
+            bandwidth=bandwidth,
+        )
+        audio_out = self._model.decode(
+            enc.audio_codes,
+            enc.audio_scales,
+            padding_mask=inputs["padding_mask"],
+        )[0]
+
+        return audio_out.squeeze(0).squeeze(0).cpu().numpy()
+
+
+register(EnCodecCodec())

compare_codec/mimi_codec.py

@@ -0,0 +1,72 @@
+"""Mimi (Kyutai) — wraps the HuggingFace transformers implementation."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+import torch
+import torchaudio
+
+from compare_codec import CodecConfig, register
+
+
+class MimiCodec:
+    """Mimi codec with lazy model loading."""
+
+    def __init__(self) -> None:
+        self._model = None
+        self._fe = None
+
+    @property
+    def name(self) -> str:
+        return "Mimi"
+
+    @property
+    def sample_rate(self) -> int:
+        return 24_000
+
+    def configs(self) -> list[CodecConfig]:
+        return [
+            CodecConfig(
+                name="1.1 kbps",
+                params={"sample_rate": 24_000},
+            )
+        ]
+
+    def _load(self):
+        if self._model is None:
+            from transformers import AutoFeatureExtractor, MimiModel
+
+            self._model = MimiModel.from_pretrained("kyutai/mimi")
+            self._model.eval()
+            self._fe = AutoFeatureExtractor.from_pretrained("kyutai/mimi")
+
+    @torch.no_grad()
+    def encode_decode(self, audio_path: Path, config: CodecConfig) -> np.ndarray:
+        self._load()
+        target_sr: int = config.params["sample_rate"]
+
+        wav, sr = torchaudio.load(str(audio_path))
+        if wav.shape[0] > 1:
+            wav = wav.mean(dim=0, keepdim=True)
+        if sr != target_sr:
+            wav = torchaudio.functional.resample(wav, sr, target_sr)
+
+        original_len = wav.shape[-1]
+        inputs = self._fe(
+            raw_audio=wav.squeeze(0).numpy(),
+            sampling_rate=target_sr,
+            return_tensors="pt",
+        )
+        enc = self._model.encode(inputs["input_values"], inputs["padding_mask"])
+        audio_out = self._model.decode(
+            enc.audio_codes, inputs["padding_mask"]
+        )[0]
+
+        # Trim to original length (Mimi may pad).
+        audio_out = audio_out.squeeze(0).squeeze(0).cpu().numpy()[:original_len]
+        return audio_out
+
+
+register(MimiCodec())