Spaces:
Sleeping
Sleeping
File size: 2,318 Bytes
f5b74a2 458a0e7 f5b74a2 458a0e7 f5b74a2 458a0e7 | 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 | """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", device_map="auto"
)
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",
)
device = self._model.device
inputs = {k: v.to(device) for k, v in inputs.items()}
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())
|