cpu and gpu support; readme update;

.gitignore

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

README.md

@@ -1,3 +1,67 @@
 ---
 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] 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. 
+
+
+# Usage
+
+
+
+<!-- 
+- different models for different khz  
+-->
+
+
+# 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=80%>
+</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

@@ -124,7 +124,7 @@ class DAC(PreTrainedModel):
             end = start + chunk_size
             chunk = x[:, :, start:end]
             chunk = self.dac.preprocess(chunk, sr)
-            zq, s, _, _, _ = self.dac.encode(chunk)
+            zq, s, _, _, _ = self.dac.encode(chunk.to(self.device))
             zq = zq.cpu()
             s = s.cpu()
             """
@@ -175,7 +175,7 @@ class DAC(PreTrainedModel):
         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)  # (b, 1, chunk_size*self.downsampling_rate)
+            waveform = self.dac.decode(chunk.to(self.device))  # (b, 1, chunk_size*self.downsampling_rate)
             waveform = waveform.cpu()
             
             if isinstance(waveform_concat, type(None)):
@@ -198,11 +198,12 @@ class DAC(PreTrainedModel):
         """
         s: (b, n_rvq, length)
         """
-        zq, _, _ = self.dac.quantizer.from_codes(s)  # zq: (b, d, 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, fname)
+        torch.save(tensor.cpu(), fname)
     
     def load_tensor(self, fname:str):
         return torch.load(fname)

save_model.ipynb

@@ -101,7 +101,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 8,
    "metadata": {},
    "outputs": [
     {
@@ -124,6 +124,26 @@
     "model = AutoModel.from_pretrained('hance-ai/descript-audio-codec', trust_remote_code=True)"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "device(type='cpu')"
+      ]
+     },
+     "execution_count": 9,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "model.device"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 6,

test_DAC.ipynb

@@ -11,20 +11,40 @@
    "cell_type": "code",
    "execution_count": 1,
    "metadata": {},
-   "outputs": [],
+   "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",
     "import torch\n",
     "\n",
-    "from model import DAC"
+    "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'"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
    "outputs": [
     {
      "name": "stderr",
@@ -39,22 +59,13 @@
    ],
    "source": [
     "# load the model\n",
-    "dac = DAC('44khz')"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 3,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# settings\n",
-    "fname = str(Path(os.getcwd()).joinpath('.sample_sound', 'jazz_swing.wav'))"
+    "config = DACConfig()\n",
+    "dac = DAC(config).to(device)"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 7,
    "metadata": {},
    "outputs": [
     {
@@ -75,14 +86,22 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 8,
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "waveform.shape: torch.Size([1, 1, 441344])\n"
+      "waveform.shape: torch.Size([1, 1, 441344])\n",
+      "waveform.shape: torch.Size([1, 1, 441344])\n",
+      "waveform.shape: torch.Size([1, 1, 441344])\n",
+      "waveform.shape: torch.Size([1, 1, 441344])\n",
+      "waveform.shape: torch.Size([1, 1, 441344])\n",
+      "waveform.shape: torch.Size([1, 1, 441344])\n",
+      "waveform.shape: torch.Size([1, 1, 441344])\n",
+      "waveform.shape: torch.Size([1, 1, 441344])\n",
+      "15.4 s ± 142 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
      ]
     }
    ],
@@ -94,7 +113,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 9,
    "metadata": {},
    "outputs": [
     {
@@ -113,7 +132,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -123,14 +142,14 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "d:\\projects\\sound_effect_variation_generation\\dac_model.py:181: 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",
+      "d:\\projects\\descript-audio-codec\\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"
      ]
     }
@@ -140,6 +159,13 @@
     "dac.save_tensor(s, 'tokens.pt')\n",
     "loaded_s = dac.load_tensor('tokens.pt')  # s == loaded_s"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {