File size: 1,086 Bytes
4f5f309 | 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 | from __future__ import annotations
from pathlib import Path
import numpy as np
import onnxruntime as ort
from .ort_opts import session_options
class MimiDecoder:
def __init__(self, path: str | Path, threads: int | None = None) -> None:
self.path = Path(path)
self.opts = session_options(threads)
self.session = ort.InferenceSession(str(self.path), self.opts, providers=["CPUExecutionProvider"])
def decode(self, codes: np.ndarray) -> np.ndarray:
"""
Input:
codes: (16, F) int64
Returns:
wav: (F * 1920,) float32
"""
codes = np.asarray(codes, dtype=np.int64)
if codes.ndim == 2:
codes_input = codes[None, :, :] # (1, 16, F)
elif codes.ndim == 3:
codes_input = codes
else:
raise ValueError(f"Expected codes to have 2 or 3 dims, got {codes.shape}")
outputs = self.session.run(None, {"codes": codes_input})
wav = outputs[0] # (1, 1, samples)
wav = np.squeeze(wav)
return wav.astype(np.float32)
|