Spaces:
Paused
Paused
File size: 2,629 Bytes
20857b0 | 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 85 86 | import torch
from .base import AudioTokenizer
class EnCodecAudioCodec(AudioTokenizer):
def __init__(self, device: str = "cpu", bandwidth: float = 6.0):
self.device = torch.device(device)
self._bandwidth = bandwidth
self._model = None
self._loaded = False
def _lazy_load(self):
if self._loaded:
return
try:
from encodec import EncodecModel
self._model = EncodecModel.encodec_model_24khz()
self._model.set_target_bandwidth(self._bandwidth)
self._model.to(self.device)
self._model.eval()
self._loaded = True
except ImportError:
raise ImportError("encodec package not installed; run: pip install encodec")
except Exception as e:
raise RuntimeError(f"failed to load EnCodec: {e}")
@torch.inference_mode()
def encode(self, audio: torch.Tensor) -> torch.Tensor:
self._lazy_load()
audio = audio.to(self.device)
if audio.dim() == 1:
audio = audio.unsqueeze(0).unsqueeze(0)
elif audio.dim() == 2:
audio = audio.unsqueeze(1)
from encodec.utils import convert_audio
audio = convert_audio(audio, self._model.sample_rate, self._model.sample_rate, self._model.channels)
frames = self._model.encode(audio)
codes = torch.cat([f[0] for f in frames], dim=-1)
return codes
@torch.inference_mode()
def decode(self, tokens: torch.Tensor) -> torch.Tensor:
self._lazy_load()
if tokens.dim() == 2:
tokens = tokens.unsqueeze(0)
frames = self._model.decode([(tokens, None)])
return frames[0]
@property
def sample_rate(self) -> int:
return 24000
@property
def num_codebooks(self) -> int:
return 8
@property
def name(self) -> str:
return f"encodec-{self._bandwidth}kbps"
class DummyAudioCodec(AudioTokenizer):
def __init__(self, device: str = "cpu"):
self.device = torch.device(device)
def encode(self, audio: torch.Tensor) -> torch.Tensor:
B = audio.shape[0] if audio.dim() > 1 else 1
return torch.randint(0, 2048, (4, B * 10), dtype=torch.int32, device=self.device)
def decode(self, tokens: torch.Tensor) -> torch.Tensor:
length = tokens.shape[-1] * 320
return torch.randn(1, length, device=self.device)
@property
def sample_rate(self) -> int:
return 32000
@property
def num_codebooks(self) -> int:
return 4
@property
def name(self) -> str:
return "dummy"
|