copy paste

.gitignore

@@ -0,0 +1,4 @@
+token.txt
+__pycache__/
+tokens.pt
+out.wav

README.md

@@ -1,3 +1,138 @@
 ---
 license: mit
+tags:
+- DAC
+- Descript Audio Codec
+- PyTorch
 ---
+
+# Descript Audio Codec (DAC)
+DAC is the state-of-the-art audio tokenizer with improvement upon the previous tokenizers like SoundStream and EnCodec.
+
+This model card provides an easy-to-use API for a *pretrained DAC* [1] for 24khz audio whose backbone and pretrained weights are from [its original reposotiry](https://github.com/descriptinc/descript-audio-codec). With this API, you can encode and decode by a single line of code either using CPU or GPU. Furhtermore, it supports chunk-based processing for memory-efficient processing, especially important for GPU processing. 
+
+
+
+
+
+
+
+
+
+### Model variations
+There are three types of model depending on an input audio sampling rate. 
+
+| Model | Input audio sampling rate [khz] |
+| ------------------ | ----------------- |
+| [`hance-ai/descript-audio-codec-44khz`](https://huggingface.co/hance-ai/descript-audio-codec-44khz) | 44.1khz |
+| [`hance-ai/descript-audio-codec-24khz`](https://huggingface.co/hance-ai/descript-audio-codec-24khz) | 24khz |
+| [`hance-ai/descript-audio-codec-16khz`](https://huggingface.co/hance-ai/descript-audio-codec-16khz) | 16khz |
+
+
+
+
+
+
+
+
+
+
+# Usage
+
+### Load
+```python
+from transformers import AutoModel
+
+# device setting
+device = 'cpu'  # or 'cuda:0'
+
+# load
+model = AutoModel.from_pretrained('hance-ai/descript-audio-codec-24khz', trust_remote_code=True)
+model.to(device)
+```
+
+### Encode
+```python
+audio_filename = 'path/example_audio.wav'
+zq, s = model.encode(audio_filename)
+```
+`zq` is discrete embeddings with dimension of (1, num_RVQ_codebooks, token_length) and `s` is a token sequence with dimension of (1, num_RVQ_codebooks, token_length). 
+
+
+### Decode
+```python
+# decoding from `zq`
+waveform = model.decode(zq=zq)  # (1, 1, audio_length); the output has a mono channel.
+
+# decoding from `s`
+waveform = model.decode(s=s)  # (1, 1, audio_length); the output has a mono channel.
+```
+
+### Save a waveform as an audio file
+```python
+model.waveform_to_audiofile(waveform, 'out.wav')
+```
+
+### Save and load tokens
+```python
+model.save_tensor(s, 'tokens.pt')
+loaded_s = model.load_tensor('tokens.pt')
+```
+
+
+
+
+
+
+
+
+# Runtime
+
+To give you a brief idea, the following table reports average runtime on CPU and GPU to encode and decode 10s audio. The runtime is measured in second. The used CPU is Intel Core i9 11900K and GPU is RTX3060.
+```
+|           Task  |   CPU   |   GPU   |
+|-----------------|---------|---------|
+|    Encoding     |   6.71  |  0.19   |
+|    Decoding     |   15.4  |  0.31   |
+```
+The decoding process takes a longer simply because the decoder is larger than the encoder. 
+
+
+
+
+
+
+
+
+
+# Technical Discussion
+
+### Chunk-based Processing
+It's introduced for memory-efficient processing for both encoding and decoding. 
+For encoding, we simply chunk an audio into N chunks and process them iteratively.
+Similarly, for decoding, we chunk a token set into M chunks of token subsets and process them iteratively.
+However, the decoding process with naive chunking causes an artifact in the decoded audio. 
+That is because the decoder reconstructs audio given multiple neighboring tokens (i.e., multiple neighboring tokens for a segment of audio) rather than a single token for a segment of audio. 
+
+To tackle the problem, we introduce overlap between the chunks in the decoding, parameterized by `decoding_overlap_rate` in the model. By default, we introduce 10% of overlap between the chunks. Then, two subsequent chunks reuslt in two segments of audio with 10% overlap, and the overlap is averaged out for smoothing.
+
+The following figure compares reconstructed audio with and without the overlapping.
+<p align="center">
+<img src=".fig/artifact-dac-decoding without overlap.png" alt="" width=50%>
+</p>
+
+
+
+
+
+
+
+# References
+[1] Kumar, Rithesh, et al. "High-fidelity audio compression with improved rvqgan." Advances in Neural Information Processing Systems 36 (2024).
+
+
+
+<!-- contributions 
+- chunk processing 
+- add device parameter in the test notebook
+-->

model.py

@@ -0,0 +1,212 @@
+from typing import Union
+
+import numpy as np
+import torch
+import torchaudio
+import torch.nn as nn
+import torchaudio.transforms as transforms
+from transformers import PretrainedConfig, PreTrainedModel
+
+import dac
+from audiotools import AudioSignal
+
+from utils import freeze
+
+
+class DACConfig(PretrainedConfig):
+    model_type = 'dac'
+
+    def __init__(self, 
+                 model_type_by_sampling_freq:str='44khz',
+                 encoding_chunk_size_in_sec:int=1,
+                 decoding_chunk_rate:float=0.1,
+                 decoding_overlap_rate:float=0.1,
+                 **kwargs):
+        super().__init__(**kwargs)
+        """
+        Initializes the model object.
+        Args:
+            model_type_by_sampling_freq (str, optional): The model type based on the sampling frequency. Defaults to '44khz'. Choose among ['44khz', '24khz', '16khz']
+            encoding_chunk_size_in_sec (int, optional): The size of the encoding chunk in seconds. Defaults to 1.
+            decoding_chunk_rate (float, optional): The decoding chunk rate. Must be between 0 and 1. Defaults to 0.1.
+            decoding_overlap_rate (float, optional): The decoding overlap rate. Must be between 0 and 1. Defaults to 0.1.
+            **kwargs: Additional keyword arguments.
+        Raises:
+            AssertionError: If the model_type_by_sampling_freq is not one of ['44khz', '24khz', '16khz'].
+            AssertionError: If the decoding_chunk_rate is not between 0 and 1.
+            AssertionError: If the decoding_overlap_rate is not between 0 and 1.
+        """
+        self.model_type_by_sampling_freq = model_type_by_sampling_freq
+        self.encoding_chunk_size_in_sec = encoding_chunk_size_in_sec
+        self.decoding_chunk_rate = decoding_chunk_rate
+        self.decoding_overlap_rate = decoding_overlap_rate
+
+        assert model_type_by_sampling_freq.lower() in ['44khz', '24khz', '16khz']
+        assert decoding_chunk_rate > 0 and decoding_chunk_rate <= 1.0, '`decoding_chunk_rate` must be bewteen 0 and 1.'
+        assert decoding_overlap_rate >= 0 and decoding_overlap_rate < 1.0, '`decoding_overlap_rate` must be bewteen 0 and 1.'
+
+
+
+class DAC(PreTrainedModel):
+    config_class = DACConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.model_type_by_sampling_freq = config.model_type_by_sampling_freq.lower()
+        self.model_type_by_sampling_freq_int = {'44khz':44100, '24khz':24000, '16khz':16000}[self.model_type_by_sampling_freq]
+        self.encoding_chunk_size_in_sec = config.encoding_chunk_size_in_sec
+        self.decoding_chunk_rate = config.decoding_chunk_rate
+        self.decoding_overlap_rate = config.decoding_overlap_rate
+
+
+        dac_path = dac.utils.download(model_type=self.model_type_by_sampling_freq)
+        self.dac = dac.DAC.load(dac_path)
+        self.dac.eval()
+        freeze(self.dac)
+
+        self.downsampling_rate = int(np.prod(self.dac.encoder_rates))  # 512
+    
+    def load_audio(self, filename:str):
+        waveform, sample_rate = torchaudio.load(filename)  # waveform: (n_channels, length); sample_rate: const.
+        return waveform, sample_rate
+    
+    def resample_audio(self, waveform:torch.FloatTensor, orig_sr:int, target_sr:int):
+        """
+        - sr: sampling rate
+        - waveform: (n_channels, length)
+        """
+        if orig_sr == target_sr:
+            return waveform
+
+        converter = transforms.Resample(orig_freq=orig_sr, new_freq=target_sr)
+        waveform = converter(waveform)  # (n_channels, new_length)
+        return waveform  # (n_channels, new_length)
+
+    def to_mono_channel(self, waveform:torch.FloatTensor):
+        """
+        - waveform: (n_channels, length)
+        """
+        n_channels = waveform.shape[0]
+        if n_channels > 1:
+            waveform = torch.mean(waveform, dim=0, keepdim=True)  # (1, length)
+        return waveform  # (1, length)
+    
+    @torch.no_grad()
+    def encode(self, audio_fname:str):
+        self.eval()
+
+        waveform, sr = self.load_audio(audio_fname)
+        waveform = self.resample_audio(waveform, orig_sr=sr, target_sr=self.model_type_by_sampling_freq_int)
+        sr = self.model_type_by_sampling_freq_int
+        waveform = self.to_mono_channel(waveform)  # DAC accepts a mono channel only.
+        
+        zq, s = self._chunk_encoding(waveform, sr)
+        return zq, s
+
+    def _chunk_encoding(self, waveform:torch.FloatTensor, sr:int):
+        # TODO: can I make it parallel?
+        """
+        waveform: (c l)
+        """
+        x = waveform  # brief varname
+        x = x.unsqueeze(1)  # (b 1 l); add a null batch dim
+        chunk_size = int(self.encoding_chunk_size_in_sec * sr)
+
+        # adjust `chunk_size` to prevent any padding in `dac.preprocess`, which causes a gap between the mini-batches in the resulting music.
+        remainer = chunk_size % self.dac.hop_length
+        chunk_size = chunk_size-remainer
+
+        # process
+        zq_list, s_list = [], []
+        audio_length = x.shape[-1]
+        for start in range(0, audio_length, chunk_size):
+            end = start + chunk_size
+            chunk = x[:, :, start:end]
+            chunk = self.dac.preprocess(chunk, sr)
+            zq, s, _, _, _ = self.dac.encode(chunk.to(self.device))
+            zq = zq.cpu()
+            s = s.cpu()
+            """
+            "zq" : Tensor[B x D x T]
+                Quantized continuous representation of input
+                = summation of all the residual quantized vectors across every rvq level
+                = E(x) = z = \sum_n^N{zq_n} where N is the number of codebooks
+            "s" : Tensor[B x N x T]
+                Codebook indices for each codebook
+                (quantized discrete representation of input)
+                *first element in the N dimension = first RVQ level
+            """
+            zq_list.append(zq)
+            s_list.append(s)
+            torch.cuda.empty_cache()
+        
+        zq = torch.cat(zq_list, dim=2).float()  # (1, d, length)
+        s = torch.cat(s_list, dim=2).long()  # (1, n_rvq, length)
+
+        return zq, s
+    
+    @torch.no_grad()
+    def decode(self, *, zq:Union[torch.FloatTensor,None]=None, s:Union[torch.IntTensor,None]=None):
+        """
+        zq: (b, d, length)
+        """
+        if isinstance(zq,type(None)) and isinstance(s,type(None)):
+            assert False, 'one of them must be valid.'
+        self.eval()
+
+        if not isinstance(zq,type(None)):
+            waveform = self._chunk_decoding(zq)  # (b, 1, length); output always has a mono-channel.
+        if not isinstance(s,type(None)):
+            zq = self.code_to_zq(s)
+            waveform = self._chunk_decoding(zq)  # (b, 1, length); output always has a mono-channel.
+
+        return waveform
+    
+    def _chunk_decoding(self, zq:torch.FloatTensor):
+        """
+        zq: (b, d, length)
+        """
+        length = zq.shape[-1]
+        chunk_size = round(int(self.decoding_chunk_rate * length))
+        overlap_size = round(self.decoding_overlap_rate * chunk_size)  # overlap size in terms of token length
+        overlap_size_in_data_space = round(overlap_size * self.downsampling_rate)
+        waveform_concat = None
+        for start in range(0, length, chunk_size-overlap_size):
+            end = start + chunk_size
+            chunk = zq[:,:, start:end]  # (b, d, chunk_size)
+            waveform = self.dac.decode(chunk.to(self.device))  # (b, 1, chunk_size*self.downsampling_rate)
+            waveform = waveform.cpu()
+            
+            if isinstance(waveform_concat, type(None)):
+                waveform_concat = waveform.clone()
+            else:
+                if self.decoding_overlap_rate != 0.:
+                    prev_x = waveform_concat[:,:,:-overlap_size_in_data_space]
+                    rest_of_new_x = waveform[:,:,overlap_size_in_data_space:]
+                    overlap_x_from_prev_x = waveform_concat[:,:,-overlap_size_in_data_space:]  # (b, 1, overlap_size_in_data_space)
+                    overlap_x_from_new_x = waveform[:,:,:overlap_size_in_data_space]  # (b, 1, overlap_size_in_data_space)
+                    overlap = (overlap_x_from_prev_x + overlap_x_from_new_x) / 2  # take mean; maybe there's a better strategy but it seems to work fine.
+                    waveform_concat = torch.cat((prev_x, overlap, rest_of_new_x), dim=-1)  # (b, 1, ..)
+                else:
+                    prev_x = waveform_concat
+                    rest_of_new_x = waveform
+                    waveform_concat = torch.cat((prev_x, rest_of_new_x), dim=-1)  # (b, 1, ..)
+        return waveform_concat  # (b, 1, length)
+
+    def code_to_zq(self, s:torch.IntTensor):
+        """
+        s: (b, n_rvq, length)
+        """
+        zq, _, _ = self.dac.quantizer.from_codes(s.to(self.device))  # zq: (b, d, length)
+        zq = zq.cpu()
+        return zq
+
+    def save_tensor(self, tensor:torch.Tensor, fname:str) -> None:
+        torch.save(tensor.cpu(), fname)
+    
+    def load_tensor(self, fname:str):
+        return torch.load(fname)
+    
+    def waveform_to_audiofile(self, waveform:torch.FloatTensor, fname:str) -> None:
+        AudioSignal(waveform, sample_rate=self.model_type_by_sampling_freq_int).write(fname)

requirements.txt

@@ -0,0 +1,4 @@
+torch==2.4.0
+torchaudio==2.4.0
+transformers==4.44.0
+descript-audio-codec==1.0.0

save_model.ipynb

@@ -0,0 +1,206 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "C:\\Users\\dslee\\AppData\\Roaming\\Python\\Python38\\site-packages\\tqdm\\auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+      "  from .autonotebook import tqdm as notebook_tqdm\n"
+     ]
+    }
+   ],
+   "source": [
+    "from model import DACConfig, DAC"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Registering a model with custom code to the auto classes\n",
+    "DACConfig.register_for_auto_class()\n",
+    "DAC.register_for_auto_class()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "C:\\Users\\dslee\\AppData\\Roaming\\Python\\Python38\\site-packages\\audiotools\\ml\\layers\\base.py:172: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
+      "  model_dict = torch.load(location, \"cpu\")\n",
+      "c:\\Users\\dslee\\anaconda3\\envs\\sound_effect_variation_generation\\lib\\site-packages\\torch\\nn\\utils\\weight_norm.py:134: FutureWarning: `torch.nn.utils.weight_norm` is deprecated in favor of `torch.nn.utils.parametrizations.weight_norm`.\n",
+      "  WeightNorm.apply(module, name, dim)\n"
+     ]
+    }
+   ],
+   "source": [
+    "# create instances\n",
+    "config = DACConfig(model_type_by_sampling_freq='24khz')\n",
+    "model = DAC(config)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "c:\\Users\\dslee\\anaconda3\\envs\\sound_effect_variation_generation\\lib\\site-packages\\huggingface_hub\\file_download.py:159: UserWarning: `huggingface_hub` cache-system uses symlinks by default to efficiently store duplicated files but your machine does not support them in C:\\Users\\dslee\\.cache\\huggingface\\hub\\models--hance-ai--descript-audio-codec-24khz. Caching files will still work but in a degraded version that might require more space on your disk. This warning can be disabled by setting the `HF_HUB_DISABLE_SYMLINKS_WARNING` environment variable. For more details, see https://huggingface.co/docs/huggingface_hub/how-to-cache#limitations.\n",
+      "To support symlinks on Windows, you either need to activate Developer Mode or to run Python as an administrator. In order to see activate developer mode, see this article: https://docs.microsoft.com/en-us/windows/apps/get-started/enable-your-device-for-development\n",
+      "  warnings.warn(message)\n",
+      "model.safetensors: 100%|██████████| 299M/299M [00:11<00:00, 26.5MB/s]\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "CommitInfo(commit_url='https://huggingface.co/hance-ai/descript-audio-codec-24khz/commit/cf72b50044750326ebc01f01c3a032adbaf59439', commit_message='Upload DAC', commit_description='', oid='cf72b50044750326ebc01f01c3a032adbaf59439', pr_url=None, pr_revision=None, pr_num=None)"
+      ]
+     },
+     "execution_count": 5,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# push the model to the huggingface\n",
+    "with open('token.txt', 'r') as file:\n",
+    "    token = file.read().strip()\n",
+    "\n",
+    "model.push_to_hub('hance-ai/descript-audio-codec-24khz', token=token)  # put your token"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Recommend git fetching and pulling so that the uploaded model is synced locally too."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "***"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "C:\\Users\\dslee\\AppData\\Roaming\\Python\\Python38\\site-packages\\tqdm\\auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+      "  from .autonotebook import tqdm as notebook_tqdm\n",
+      "c:\\Users\\dslee\\anaconda3\\envs\\sound_effect_variation_generation\\lib\\site-packages\\huggingface_hub\\file_download.py:159: UserWarning: `huggingface_hub` cache-system uses symlinks by default to efficiently store duplicated files but your machine does not support them in C:\\Users\\dslee\\.cache\\huggingface\\hub\\models--hance-ai--descript-audio-codec-24khz. Caching files will still work but in a degraded version that might require more space on your disk. This warning can be disabled by setting the `HF_HUB_DISABLE_SYMLINKS_WARNING` environment variable. For more details, see https://huggingface.co/docs/huggingface_hub/how-to-cache#limitations.\n",
+      "To support symlinks on Windows, you either need to activate Developer Mode or to run Python as an administrator. In order to see activate developer mode, see this article: https://docs.microsoft.com/en-us/windows/apps/get-started/enable-your-device-for-development\n",
+      "  warnings.warn(message)\n",
+      "A new version of the following files was downloaded from https://huggingface.co/hance-ai/descript-audio-codec-24khz:\n",
+      "- model.py\n",
+      ". Make sure to double-check they do not contain any added malicious code. To avoid downloading new versions of the code file, you can pin a revision.\n",
+      "C:\\Users\\dslee\\AppData\\Roaming\\Python\\Python38\\site-packages\\audiotools\\ml\\layers\\base.py:172: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
+      "  model_dict = torch.load(location, \"cpu\")\n",
+      "c:\\Users\\dslee\\anaconda3\\envs\\sound_effect_variation_generation\\lib\\site-packages\\torch\\nn\\utils\\weight_norm.py:134: FutureWarning: `torch.nn.utils.weight_norm` is deprecated in favor of `torch.nn.utils.parametrizations.weight_norm`.\n",
+      "  WeightNorm.apply(module, name, dim)\n"
+     ]
+    }
+   ],
+   "source": [
+    "# load the uploaded model\n",
+    "from transformers import AutoModel\n",
+    "model = AutoModel.from_pretrained('hance-ai/descript-audio-codec-24khz', trust_remote_code=True)\n",
+    "model.to('cpu');"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "zq.shape: torch.Size([1, 1024, 750])\n",
+      "s.shape: torch.Size([1, 32, 750])\n"
+     ]
+    }
+   ],
+   "source": [
+    "# encodeing\n",
+    "import os\n",
+    "from pathlib import Path\n",
+    "\n",
+    "fname = str(Path(os.getcwd()).joinpath('.sample_sound', 'jazz_swing.wav'))\n",
+    "zq, s = model.encode(fname)\n",
+    "print('zq.shape:', zq.shape)\n",
+    "print('s.shape:', s.shape)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "waveform.shape: torch.Size([1, 1, 239904])\n"
+     ]
+    }
+   ],
+   "source": [
+    "# decoding (from zq -- discrete latent vectors)\n",
+    "waveform = model.decode(zq=zq)\n",
+    "print('waveform.shape:', waveform.shape)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "sound_effect_variation_generation",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.8.19"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

test_DAC.ipynb

@@ -0,0 +1,183 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# DAC Audio Tokenizer"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "C:\\Users\\dslee\\AppData\\Roaming\\Python\\Python38\\site-packages\\tqdm\\auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+      "  from .autonotebook import tqdm as notebook_tqdm\n"
+     ]
+    }
+   ],
+   "source": [
+    "import os\n",
+    "from pathlib import Path\n",
+    "\n",
+    "from model import DAC, DACConfig"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# settings\n",
+    "fname = str(Path(os.getcwd()).joinpath('.sample_sound', 'jazz_swing.wav'))\n",
+    "device = 'cpu'\n",
+    "model_type_by_sampling_freq = '24khz'"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "C:\\Users\\dslee\\AppData\\Roaming\\Python\\Python38\\site-packages\\audiotools\\ml\\layers\\base.py:172: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
+      "  model_dict = torch.load(location, \"cpu\")\n",
+      "c:\\Users\\dslee\\anaconda3\\envs\\sound_effect_variation_generation\\lib\\site-packages\\torch\\nn\\utils\\weight_norm.py:134: FutureWarning: `torch.nn.utils.weight_norm` is deprecated in favor of `torch.nn.utils.parametrizations.weight_norm`.\n",
+      "  WeightNorm.apply(module, name, dim)\n"
+     ]
+    }
+   ],
+   "source": [
+    "# load the model\n",
+    "config = DACConfig(model_type_by_sampling_freq=model_type_by_sampling_freq)\n",
+    "model = DAC(config).to(device)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "zq.shape: torch.Size([1, 1024, 750])\n",
+      "s.shape: torch.Size([1, 32, 750])\n"
+     ]
+    }
+   ],
+   "source": [
+    "# encoding\n",
+    "zq, s = model.encode(fname)\n",
+    "print('zq.shape:', zq.shape)\n",
+    "print('s.shape:', s.shape)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "waveform.shape: torch.Size([1, 1, 239904])\n"
+     ]
+    }
+   ],
+   "source": [
+    "# decoding (from zq -- discrete latent vectors)\n",
+    "waveform = model.decode(zq=zq)\n",
+    "print('waveform.shape:', waveform.shape)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "waveform.shape: torch.Size([1, 1, 239904])\n"
+     ]
+    }
+   ],
+   "source": [
+    "# decoding (from s -- tokens)\n",
+    "waveform = model.decode(s=s)\n",
+    "print('waveform.shape:', waveform.shape)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# save waveform into an audio file\n",
+    "model.waveform_to_audiofile(waveform, 'out.wav')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "d:\\projects\\descript-audio-codec-24khz\\model.py:209: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
+      "  return torch.load(fname)\n"
+     ]
+    }
+   ],
+   "source": [
+    "# save and load tokens\n",
+    "model.save_tensor(s, 'tokens.pt')\n",
+    "loaded_s = model.load_tensor('tokens.pt')  # s == loaded_s"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "sound_effect_variation_generation",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.8.19"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

utils.py

@@ -0,0 +1,6 @@
+
+
+def freeze(model):
+    for param in model.parameters():
+        param.requires_grad = False
+