AudioCodec-Hub is a Python library for encoding and decoding audio data, supporting various neural audio codec models
from .audio_codec import AudioCodec
audiocodec/__init__.py
ga642381/AudioCodec-Hub

import torch
import json
import numpy as np
import soundfile as sf
import librosa
from tqdm import tqdm
from pathlib import Path
from torch.utils.data import DataLoader
from .dataset import SpeechDataset
from .codec_models import DAC, Encodec, AudioDec

class AudioCodec:
    def __init__(self, name):
        """
        Model Specifications:
        ----------------------

        Encodec (encodec_24khz):
        +-------+----+---+---+----+----+
        | n_q   |  2 |  4|  8| 16 | 32 |
        +-------+----+---+---+----+----+
        | kbps  | 1.5| 3 | 6 | 12 | 24 |
        +-------+----+---+---+----+----+

        DAC (dac_24khz):
        +-------+----+---+---+----+----+
        | n_q   |  2 |  4|  8| 16 | 32 |
        +-------+----+---+---+----+----+
        | kbps  | 1.5| 3 | 6 | 12 | 24 |
        +-------+----+---+---+----+----+

        AudioDec (audiodec_24khz, audiodec_48khz):
        +-------+----+
        | n_q   |  8 |
        +-------+----+
        """
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.model = self._load_codec_model(name, self.device)
        self.model_type = name.split("_")[0]
        assert self.model_type in ["encodec", "dac", "audiodec"]

    def encode_file(self, in_file, out_file, n_q=None, codebook_offset=False):
        """
        Encode the audio file into codes and save the resulting codes to an output file.

        Args:
            in_file (str): Path to the input audio file.
            out_file (str): Path to the output file where encoded codes will be saved.
            n_q (int, optional): Number of codes to select from the encoded codes. If provided, only the first
                'n_q' codes will be kept. Default is None, meaning all codes will be kept.
            codebook_offset (bool, optional): Whether to apply a codebook offset. Default is False.
        """
        x, _ = librosa.load(in_file, sr=self.model.sample_rate)
        x = torch.tensor(x).to(self.device).unsqueeze(0).unsqueeze(0) # [1, 1, L]
        codes = self.model.encode_tensor(x, n_q=n_q)

        # wrapup codes
        codes = codes.cpu().squeeze(0).numpy()

        # save codes
        self._save_file(codes, out_file, codebook_offset)

    def encode_dir(self, in_dir, out_dir, n_q=None, batch_size=2, codebook_offset=False):
        """
        Encode audio files from the input directory and save the encoded files to the output directory.

        Args:
            in_dir (str or Path): The input directory containing the audio files to be encoded.
            out_dir (str or Path): The output directory where the encoded files will be saved.
        """
        in_dir = Path(in_dir)
        out_dir = Path(out_dir)
        out_dir.mkdir(parents=True, exist_ok=True)
        speech_dataloader = self._load_dataloader(in_dir, batch_size)
        for speech_batch, masks_batch, file_paths in tqdm(speech_dataloader):
            speech_batch = speech_batch.to(self.device)
            codes_batch = self.model.encode_tensor(speech_batch, masks_batch, n_q=n_q)

            codes_list= self._wrapup_codes(codes_batch, masks_batch)
            for codes, file_path in zip(codes_list, file_paths):
                out_file = out_dir / f"{file_path.stem}.json"

                codes = codes.cpu().squeeze(0).numpy()
                self._save_file(codes, out_file, codebook_offset)

    def decode_file(self, in_file, out_file, codebook_offset):
        """
        Decode codes from an input file and save the audio waveform to an output file.

        Args:
            in_file (str): Path to the input file containing encoded codes.
            out_file (str): Path to the output file where the decoded audio waveform will be saved.
            codebook_offset (bool): Whether a codebook offset was applied during encoding.
        """
        codes = self._load_file(in_file, codebook_offset)
        codes = torch.tensor(codes).unsqueeze(0).to(self.device)
        wav = self.model.decode_tensor(codes)
        wav = wav.cpu().numpy().squeeze(0).squeeze(0)
        sf.write(out_file, wav, self.model.sample_rate, "PCM_16")

    def decode_dir(self, in_dir, out_dir, codebook_offset=False):
        """
        Decode codes from input files in a directory and save the audio waveforms to output files.

        Args:
            in_dir (str): Path to the directory containing encoded code files.
            out_dir (str): Path to the directory where the decoded audio waveforms will be saved.
            codebook_offset (bool): Whether a codebook offset was applied during encoding.
        """
        in_dir = Path(in_dir)
        out_dir = Path(out_dir)
        out_dir.mkdir(parents=True, exist_ok=True)
        
        for in_file in in_dir.glob("*.json"):
            out_file = out_dir / in_file.with_suffix(".wav").name
            self.decode_file(in_file, out_file, codebook_offset)

    def _load_codec_model(self, name, device):
        model_type, model_sr = name.split("_")
        assert model_type in ['dac', 'encodec', 'audiodec']
        assert model_sr in ['24khz']

        if model_type == "dac":
            model = DAC(model_sr, device)
        elif model_type == "encodec":
            model = Encodec(model_sr, device)
        elif model_type == 'audiodec':
            model = AudioDec(model_sr, device)
        return model
    

    def _load_dataloader(self, data_dir, batch_size):
        """
        Load and return a dataloader for a given dataset directory.

        Args:
            data_dir (str): Path to the directory containing audio waveform files.
            batch_size (int): Batch size for the dataloader.

        Returns:
            DataLoader: The dataloader for the given dataset.
        """
        wav_fns = [Path(x) for x in librosa.util.find_files(data_dir)]
        speech_dataset = SpeechDataset(wav_fns, target_sr=self.model.sample_rate)
        dataloader = DataLoader(speech_dataset, batch_size=batch_size, shuffle=False, collate_fn=speech_dataset.collate_fn)
        print(f"Number of wav files found: {len(speech_dataset)}")
        return dataloader

    def _save_file(self, codes, out_file, codebook_offset):
        """
        Save a numpy array of codes to a file in JSON format.

        Parameters:
            codes (np array): Numpy array of codes corresponding to one audio file.
            out_file (str): Output file path.
            codebook_offset (bool): Whether to apply codebook offset.

        The function shifts codes if codebook_offset is True, converts the codes to a dictionary,
        and writes the dictionary to the specified JSON file.
        """
        # shift
        if codebook_offset:
            codes = self._shift_codes(codes)

        # make dict
        codes_dict = {}
        for i , codes_i in enumerate(codes):
            codes_i = [str(c) for c in codes_i]
            codes_dict[str(i)] = " ".join(codes_i)

        # write
        with open(out_file, "w") as f:
            json.dump(codes_dict, f, indent=4)

    def _load_file(self, in_file, codebook_offset):
        """
        Load encoded codes from an input file and optionally apply a codebook offset.

        Args:
            in_file (str): Path to the input file containing encoded codes in JSON format.
            codebook_offset (bool): Whether to apply a codebook offset to the loaded codes.

        Returns:
            np.ndarray: Loaded encoded codes as a NumPy array.
        """
        with open(in_file, "r") as f:
            codes_dict = json.load(f)
            codes = []
            for i in range(len(codes_dict)):
                codes_i_str = codes_dict[str(i)]
                codes_i = np.array([int(c) for c in codes_i_str.split()])
                codes.append(codes_i)
            codes = np.stack(codes)
        

        if self.model_type == "audiodec":
            # These kinds of models expect the codes are offset
            if not codebook_offset:
                codes = self._shift_codes(codes, "add")
        elif self.model_type in ["encodec", "dac"]:
            # These kinds of models expect the codes not offset
            if codebook_offset:
                codes = self._shift_codes(codes, "sub")
        else:
            raise NotImplementedError
        
        return codes
    
    def _wrapup_codes(self, codes_batch, masking_batch):
        """
        Process encoded codes and apply masking.
        
        Parameters:
            codes (np array): Encoded codes of shape B x n_q x T_code.
            masking: Masking array of shape B x T_audio.
            
        Returns:
            List of numpy arrays, each of shape n_q x T_code.
        """
        codes_list = []
        for codes, masks in zip(codes_batch, masking_batch):
            c_len = len(masks.nonzero()) // self.model.downsample_rate
            real_codes = codes[:, :c_len]
            codes_list.append(real_codes)
        return codes_list
    
    def _shift_codes(self, codes, mode="add"):
        """
        codes: np array of codes corresponding to one audio
        """
        for i in range(len(codes)):
            offset_const = int(i * self.model.codebook_size)
            if mode == "add":
                codes[i] += offset_const
            elif mode == "sub":
                codes[i] -= offset_const
        return codes
audiocodec/audio_codec.py
ga642381/AudioCodec-Hub
import dac
import torch
import sys
import os
import git
import requests
import zipfile
from tqdm import tqdm
from pathlib import Path
from transformers import EncodecModel

class CodecModel:
    def __init__(self):
        self.model = None

    def load_model(self):
        pass
    
    @torch.no_grad()
    def encode_tensor(self, x):
        """
        Args:
            "x" : Tensor[B x 1 x T_wav]

        Returns:
            "codes" : Tensor[B x n_q x T_code]
                Codebook indices for each codebook
                (quantized discrete representation of input)
        """
        pass
    
    @torch.no_grad()
    def decode_tensor(self):
        """
        Args:
            codes (Tensor): Encoded codes to be decoded.

        Returns:
            Tensor: Decoded audio waveform.
        """
        pass

    @property
    def codebook_size(self):
        pass

    @property
    def sample_rate(self):
        pass

    @property
    def downsample_rate(self):
        """
        Sometimes called `stride factor` and `hop size`.
        T_audio = downsample_rate * T_code
        """
        pass

class DAC(CodecModel):
    def __init__(self, model_sr, device):
        super().__init__()
        self.load_model(model_sr, device)
        
    def load_model(self, model_sr, device):
        model_path = dac.utils.download(model_type=model_sr)
        self.model = dac.DAC.load(model_path).to(device)
    
    @torch.no_grad()
    def encode_tensor(self, x: torch.Tensor, padding_mask=None, n_q=None) -> torch.Tensor:
        zq, codes, _, _, _ = self.model.encode(x)
        if n_q is not None:
            codes = codes[:, :n_q, :]
        return codes
    
    @torch.no_grad()
    def decode_tensor(self, codes: torch.Tensor):
        zq, _, codes = self.model.quantizer.from_codes(codes)
        y = self.model.decode(zq)
        return y

    @property
    def codebook_size(self):
        return self.model.codebook_size
    
    @property
    def downsample_rate(self):
        return 320

    @property
    def sample_rate(self):
        return self.model.sample_rate

class Encodec(CodecModel):
    def __init__(self, model_sr, device):
        super().__init__()
        self.load_model(model_sr, device)
        self.nq_bw = {2: 1.5, 4: 3.0, 8: 6.0, 16: 12.0, 32: 24.0}

    def load_model(self, model_sr, device):
        model = EncodecModel.from_pretrained(f"facebook/encodec_{model_sr}")
        self.model = model.to(device)

    @torch.no_grad()
    def encode_tensor(self, x: torch.Tensor, padding_mask=None, n_q=None) -> torch.Tensor:
        bw = self.nq_bw[n_q] if n_q is not None else None
        model_output = self.model.encode(x, padding_mask=padding_mask, bandwidth=bw)
        codes = model_output['audio_codes'].squeeze(0)
        # B x n_q x T
        return codes
    
    @torch.no_grad()
    def decode_tensor(self, codes: torch.Tensor):
        codes = codes.unsqueeze(0) # B x n_q x T -> 1 x B x n_q x T
        y = self.model.decode(codes, audio_scales=[None])['audio_values']
        return y

    @property
    def codebook_size(self):
        return self.model.config.codebook_size
    
    @property
    def downsample_rate(self):
        return 320

    @property
    def sample_rate(self):
        return self.model.config.sampling_rate

class AudioDec(CodecModel):
    def __init__(self, model_sr, device):
        super().__init__()
        self.repo_path = Path(__file__).parent /  "AudioDec"
        sys.path.append(str(self.repo_path))
        self.load_model(model_sr, device)      
        
    def load_model(self, model_sr, device):
        """
        Load the AudioDec model.

        1. Initialize the AudioDec Git submodule if not already initialized.
        2. Download pre-trained model files if not present.
        3. Load the selected model based on sample rate.
        """
        # init git module
        cur_repo = git.Repo(Path(__file__).parent.parent)
        sm = cur_repo.submodules[0]
        if not sm.module_exists():
            print("[INFO] Init AudioDec repo")
            cur_repo.git.submodule('update', '--init')

        # download exp dir
        exp_dir = self.repo_path / "exp"
        if len(list(exp_dir.rglob('*.pkl'))) == 0:
            url = "https://github.com/facebookresearch/AudioDec/releases/download/pretrain_models/exp.zip"
            zip_file_name = "exp.zip"
            response = requests.get(url, stream=True)
            total_size = int(response.headers.get('content-length', 0))
            progress_bar = tqdm(total=total_size, unit='B', unit_scale=True, desc="Downloading exp dir from AudioDec repo")
            
            with open(zip_file_name, 'wb') as zip_file:
                for data in response.iter_content(chunk_size=1024):
                    progress_bar.update(len(data))
                    zip_file.write(data)
            
            # Extract the ZIP file
            with zipfile.ZipFile(zip_file_name, 'r') as zip_ref:
                zip_ref.extractall(exp_dir.parent)
            os.remove(zip_file_name)

        from .AudioDec.utils.audiodec import AudioDec as AudioDecModel, assign_model
        if model_sr == "24khz":
            self.sr, encoder_checkpoint, decoder_checkpoint = assign_model('libritts_v1')
        elif model_sr == "48khz":
            self.sr, encoder_checkpoint, decoder_checkpoint = assign_model('vctk_v1')
        
        cwd = os.getcwd()
        os.chdir(self.repo_path) 
        audiodec = AudioDecModel(tx_device=device, rx_device=device)
        audiodec.load_transmitter(encoder_checkpoint)
        audiodec.load_receiver(encoder_checkpoint, decoder_checkpoint)
        self.model = audiodec
        os.chdir(cwd)

    @torch.no_grad()
    def encode_tensor(self, x: torch.Tensor, padding_mask=None, n_q=None) -> torch.Tensor:
        """
        return:
        z: B x D x T
        idx: nq x T
        zq: B x T x D
        """
        self.model.tx_encoder.reset_buffer()
        z = self.model.tx_encoder.encode(x)

        # see Quantizer.encode
        zq, codes = self.model.tx_encoder.quantizer.codebook.forward_index(z.transpose(2, 1), flatten_idx=False)
        if len(codes.shape) == 2:
            # forward index do squeeze batch size.
            # To match our interface, we unsqueeze it if the batch size is 1
            codes = codes.unsqueeze(1)
        codes = codes.transpose(0, 1)
        return codes
    
    @torch.no_grad()
    def decode_tensor(self, codes: torch.Tensor):
        codes = codes.squeeze(0)
        zq = self.model.rx_encoder.lookup(codes)
        y = self.model.decoder.decode(zq)
        # y = y.squeeze(1).transpose(1, 0).cpu().detach().numpy() # (T, C)
        return y

    @property
    def codebook_size(self):
        return self.model.rx_encoder.quantizer.codebook.codebook_size
    
    @property
    def downsample_rate(self):
        # Warning for future compatibility
        return 300

    @property
    def sample_rate(self):
        return self.sr
audiocodec/codec_models.py
ga642381/AudioCodec-Hub
import librosa
from torch.utils.data import Dataset
from typing import List
from transformers import EncodecFeatureExtractor

class SpeechDataset(Dataset):
    def __init__(self, file_paths:List, target_sr):
        self.file_paths = file_paths
        self.sr = target_sr
        self.processor = EncodecFeatureExtractor(sampling_rate=self.sr)

    def __len__(self):
        return len(self.file_paths)
    
    def __getitem__(self, idx):
        path = self.file_paths[idx]
        wav, _ = librosa.load(path, sr=self.sr)
        return wav, path
    
    def collate_fn(self, data):
        audios = [d[0] for d in data]
        paths = [d[1] for d in data]
        inputs = self.processor(raw_audio=audios, sampling_rate=self.sr, return_tensors="pt")
        return inputs['input_values'], inputs['padding_mask'], paths

audiocodec/dataset.py
ga642381/AudioCodec-Hub
import unittest
from audiocodec.audio_codec import AudioCodec
import os
import shutil
from pathlib import Path
from itertools import product

def test_model(model, model_name, test_results_dir, test_wav_file, test_wav_dir, n_q, codebook_offset):
    model_test_results_dir = os.path.join(test_results_dir)
    
    test_code = Path(f"{model_test_results_dir}/test_codes/nq_{n_q}_offset_{codebook_offset}/121_121726_000000_000000.json")
    test_rewav = Path(f"{model_test_results_dir}/test_rewavs/nq_{n_q}_offset_{codebook_offset}/121_121726_000000_000000.wav")

    for path in [test_code.parent, test_rewav.parent]:
        path.mkdir(parents=True, exist_ok=True)

    print(f"Testing with n_q = {n_q} and codebook_offset = {codebook_offset}")
    
    # Test encode_file and decode_file for the given model
    print(f"\nTesting encode_file and decode_file with {model_name} model...")
    model.encode_file(test_wav_file, out_file=test_code, n_q=n_q, codebook_offset=codebook_offset)
    model.decode_file(test_code, out_file=test_rewav, codebook_offset=codebook_offset)
    
    # Test encode_dir for the given model
    print(f"\nTesting directory encoding with {model_name} model...")
    test_code_dir = Path(f"{model_test_results_dir}/test_codes/nq_{n_q}_offset_{codebook_offset}")
    model.encode_dir(test_wav_dir, test_code_dir, n_q=n_q, batch_size=2, codebook_offset=codebook_offset)
    
    # Test decode_dir for the given model
    print(f"\nTesting directory decoding with {model_name} model...")
    test_rewav_dir = Path(f"{model_test_results_dir}/test_rewavs/nq_{n_q}_offset_{codebook_offset}")
    model.decode_dir(test_code_dir, test_rewav_dir, codebook_offset=codebook_offset)


class TestDAC(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        # Create test directories if they don't exist
        cls.test_results_dir = "./test_results"
        cls.test_wav_file = "./test_wavs/121_121726_000000_000000.wav"
        cls.test_wav_dir = "./test_wavs"

        os.makedirs(cls.test_results_dir, exist_ok=True)

    def setUp(self):
        # Clean up test_results directory before each test
        shutil.rmtree(self.test_results_dir, ignore_errors=True)
                
    def test_dac(self):
        model_dac = AudioCodec(name="dac_24khz")
        codebook_offset_values = [True, False]
        n_q_values = [2, 4, 8, 16, 32]

        test_combinations = product(n_q_values, codebook_offset_values)
        for n_q, codebook_offset in test_combinations:
            with self.subTest(n_q=n_q, codebook_offset=codebook_offset):
                test_model(model_dac, "encodec", self.test_results_dir, self.test_wav_file, self.test_wav_dir, n_q, codebook_offset)

class TestEncodec(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        # Create test directories if they don't exist
        cls.test_results_dir = "./test_results"
        cls.test_wav_file = "./test_wavs/121_121726_000000_000000.wav"
        cls.test_wav_dir = "./test_wavs"

        os.makedirs(cls.test_results_dir, exist_ok=True)

    def setUp(self):
        # Clean up test_results directory before each test
        shutil.rmtree(self.test_results_dir, ignore_errors=True)

    def test_encodec(self):
        model_encodec = AudioCodec(name="encodec_24khz")

        codebook_offset_values = [True, False]
        n_q_values = [2, 4, 8, 16, 32]

        test_combinations = product(n_q_values, codebook_offset_values)
        for n_q, codebook_offset in test_combinations:
            with self.subTest(n_q=n_q, codebook_offset=codebook_offset):
                test_model(model_encodec, "encodec", self.test_results_dir, self.test_wav_file, self.test_wav_dir, n_q, codebook_offset)

class TestAudioDec(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        # Create test directories if they don't exist
        cls.test_results_dir = "./test_results"
        cls.test_wav_file = "./test_wavs/121_121726_000000_000000.wav"
        cls.test_wav_dir = "./test_wavs"

        os.makedirs(cls.test_results_dir, exist_ok=True)

    def setUp(self):
        # Clean up test_results directory before each test
        shutil.rmtree(self.test_results_dir, ignore_errors=True)
                
    def test_audiodec(self):
        model_audiodec = AudioCodec(name="audiodec_24khz")

        codebook_offset_values = [True, False]
        n_q_values = [8]

        test_combinations = product(n_q_values, codebook_offset_values)
        for n_q, codebook_offset in test_combinations:
            with self.subTest(n_q=n_q, codebook_offset=codebook_offset):
                test_model(model_audiodec, "encodec", self.test_results_dir, self.test_wav_file, self.test_wav_dir, n_q, codebook_offset)

    
if __name__ == "__main__":
    unittest.main()

audiocodec/tests/test_audio_codec.py
ga642381/AudioCodec-Hub
from setuptools import setup, find_packages

with open("README.md") as f:
    long_description = f.read()

setup(
    name='audiocodec-hub',
    version='1.2',
    packages=find_packages(),
    author="Kai-Wei Chang",
    author_email="kaiwei.chang.tw@gmail.com",
    description="AudioCodec-Hub is a Python library for encoding and decoding audio data, supporting various neural audio codec models.",
    long_description=long_description,
    long_description_content_type="text/markdown",
    install_requires=[
        'GitPython',
        'librosa',
        'numpy',
        'PyYAML',
        'scikit-learn',
        'soundfile',
        'tqdm',
        'transformers==4.33.2',
        'requests',
        'descript-audio-codec==1.0.0',
        'torch',
        'torchaudio'
    ],
)
setup.py
ga642381/AudioCodec-Hub