├── .github
    └── workflows
    │   └── pypi-release.yml
├── .gitignore
├── CITATION.cff
├── LICENSE
├── README.md
├── img
    └── snac.png
├── requirements.txt
├── setup.py
└── snac
    ├── __init__.py
    ├── attention.py
    ├── layers.py
    ├── snac.py
    └── vq.py


/.github/workflows/pypi-release.yml:
--------------------------------------------------------------------------------
 1 | name: Publish Python package
 2 | 
 3 | on:
 4 |   release:
 5 |     types: [ published ]
 6 | 
 7 | permissions:
 8 |   contents: read
 9 | 
10 | jobs:
11 |   publish:
12 |     runs-on: ubuntu-latest
13 |     steps:
14 |       - uses: actions/checkout@v4
15 |       - name: Set up Python
16 |         uses: actions/setup-python@v5
17 |         with:
18 |           python-version: "3.x"
19 |       - name: Install dependencies
20 |         run: |
21 |           python -m pip install --upgrade pip
22 |           pip install setuptools wheel
23 |       - name: Build a binary wheel
24 |         run: >-
25 |           python setup.py sdist bdist_wheel
26 |       - name: Publish to PyPI
27 |         uses: pypa/gh-action-pypi-publish@release/v1
28 |         with:
29 |           password: ${{ secrets.PYPI_API_TOKEN }}
30 | 


--------------------------------------------------------------------------------
/.gitignore:
--------------------------------------------------------------------------------
  1 | # Byte-compiled / optimized / DLL files
  2 | __pycache__/
  3 | *.py[cod]
  4 | *$py.class
  5 | 
  6 | # C extensions
  7 | *.so
  8 | 
  9 | # Distribution / packaging
 10 | .Python
 11 | build/
 12 | develop-eggs/
 13 | dist/
 14 | downloads/
 15 | eggs/
 16 | .eggs/
 17 | lib/
 18 | lib64/
 19 | parts/
 20 | sdist/
 21 | var/
 22 | wheels/
 23 | share/python-wheels/
 24 | *.egg-info/
 25 | .installed.cfg
 26 | *.egg
 27 | MANIFEST
 28 | 
 29 | # PyInstaller
 30 | #  Usually these files are written by a python script from a template
 31 | #  before PyInstaller builds the exe, so as to inject date/other infos into it.
 32 | *.manifest
 33 | *.spec
 34 | 
 35 | # Installer logs
 36 | pip-log.txt
 37 | pip-delete-this-directory.txt
 38 | 
 39 | # Unit test / coverage reports
 40 | htmlcov/
 41 | .tox/
 42 | .nox/
 43 | .coverage
 44 | .coverage.*
 45 | .cache
 46 | nosetests.xml
 47 | coverage.xml
 48 | *.cover
 49 | *.py,cover
 50 | .hypothesis/
 51 | .pytest_cache/
 52 | cover/
 53 | 
 54 | # Translations
 55 | *.mo
 56 | *.pot
 57 | 
 58 | # Django stuff:
 59 | *.log
 60 | local_settings.py
 61 | db.sqlite3
 62 | db.sqlite3-journal
 63 | 
 64 | # Flask stuff:
 65 | instance/
 66 | .webassets-cache
 67 | 
 68 | # Scrapy stuff:
 69 | .scrapy
 70 | 
 71 | # Sphinx documentation
 72 | docs/_build/
 73 | 
 74 | # PyBuilder
 75 | .pybuilder/
 76 | target/
 77 | 
 78 | # Jupyter Notebook
 79 | .ipynb_checkpoints
 80 | 
 81 | # IPython
 82 | profile_default/
 83 | ipython_config.py
 84 | 
 85 | # pyenv
 86 | #   For a library or package, you might want to ignore these files since the code is
 87 | #   intended to run in multiple environments; otherwise, check them in:
 88 | # .python-version
 89 | 
 90 | # pipenv
 91 | #   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
 92 | #   However, in case of collaboration, if having platform-specific dependencies or dependencies
 93 | #   having no cross-platform support, pipenv may install dependencies that don't work, or not
 94 | #   install all needed dependencies.
 95 | #Pipfile.lock
 96 | 
 97 | # poetry
 98 | #   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
 99 | #   This is especially recommended for binary packages to ensure reproducibility, and is more
100 | #   commonly ignored for libraries.
101 | #   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
102 | #poetry.lock
103 | 
104 | # pdm
105 | #   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
106 | #pdm.lock
107 | #   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
108 | #   in version control.
109 | #   https://pdm.fming.dev/#use-with-ide
110 | .pdm.toml
111 | 
112 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
113 | __pypackages__/
114 | 
115 | # Celery stuff
116 | celerybeat-schedule
117 | celerybeat.pid
118 | 
119 | # SageMath parsed files
120 | *.sage.py
121 | 
122 | # Environments
123 | .env
124 | .venv
125 | env/
126 | venv/
127 | ENV/
128 | env.bak/
129 | venv.bak/
130 | 
131 | # Spyder project settings
132 | .spyderproject
133 | .spyproject
134 | 
135 | # Rope project settings
136 | .ropeproject
137 | 
138 | # mkdocs documentation
139 | /site
140 | 
141 | # mypy
142 | .mypy_cache/
143 | .dmypy.json
144 | dmypy.json
145 | 
146 | # Pyre type checker
147 | .pyre/
148 | 
149 | # pytype static type analyzer
150 | .pytype/
151 | 
152 | # Cython debug symbols
153 | cython_debug/
154 | 
155 | # PyCharm
156 | #  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
157 | #  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
158 | #  and can be added to the global gitignore or merged into this file.  For a more nuclear
159 | #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
160 | .idea/


--------------------------------------------------------------------------------
/CITATION.cff:
--------------------------------------------------------------------------------
 1 | cff-version: 1.2.0
 2 | message: "If you use this software, please cite it as below."
 3 | authors:
 4 |   - family-names: Siuzdak
 5 |     given-names: Hubert
 6 |   - family-names: Grötschla
 7 |     given-names: Florian
 8 |   - family-names: Lanzendörfer
 9 |     given-names: Luca A.
10 | title: 'SNAC: Multi-Scale Neural Audio Codec'
11 | url: 'https://github.com/hubertsiuzdak/snac'
12 | preferred-citation:
13 |   type: conference-paper
14 |   authors:
15 |     - family-names: Siuzdak
16 |       given-names: Hubert
17 |     - family-names: Grötschla
18 |       given-names: Florian
19 |     - family-names: Lanzendörfer
20 |       given-names: Luca A.
21 |   collection-title: 'Audio Imagination: NeurIPS 2024 Workshop AI-Driven Speech, Music, and Sound Generation'
22 |   title: 'SNAC: Multi-Scale Neural Audio Codec'
23 |   year: 2024


--------------------------------------------------------------------------------
/LICENSE:
--------------------------------------------------------------------------------
 1 | MIT License
 2 | 
 3 | Copyright (c) 2023, Descript
 4 | Copyright (c) 2024, Hubert Siuzdak
 5 | 
 6 | Permission is hereby granted, free of charge, to any person obtaining a copy
 7 | of this software and associated documentation files (the "Software"), to deal
 8 | in the Software without restriction, including without limitation the rights
 9 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
10 | copies of the Software, and to permit persons to whom the Software is
11 | furnished to do so, subject to the following conditions:
12 | 
13 | The above copyright notice and this permission notice shall be included in all
14 | copies or substantial portions of the Software.
15 | 
16 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
17 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
18 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
19 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
20 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
21 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
22 | SOFTWARE.


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
 1 | # SNAC 🍿
 2 | 
 3 | Multi-**S**cale **N**eural **A**udio **C**odec (SNAC) compresses audio into discrete codes at a low bitrate. For more information, read the paper: https://arxiv.org/abs/2410.14411
 4 | 
 5 | | 🎸 Music samples                                                                                         | 🗣️ Speech samples                                                                                       |
 6 | |----------------------------------------------------------------------------------------------------------|----------------------------------------------------------------------------------------------------------|
 7 | | <video src='https://github.com/hubertsiuzdak/snac/assets/35269911/e8adac68-d3f1-4fc1-8cf9-f48d9bcd95ed'> | <video src='https://github.com/hubertsiuzdak/snac/assets/35269911/65ac2547-c711-49d4-8a5d-64d52e6d6ba1'> |
 8 | 
 9 | 🎧 More audio samples available at https://hubertsiuzdak.github.io/snac/
10 | 
11 | ## Overview
12 | 
13 | SNAC encodes audio into hierarchical tokens similarly to SoundStream, EnCodec, and DAC (see the image
14 | on the left). However, SNAC introduces a simple change where coarse tokens are sampled less frequently,
15 | covering a broader time span (see the image on the right).
16 | 
17 | This can not only save on bitrate, but more importantly this might be very useful for language modeling approaches to
18 | audio generation. E.g. with coarse tokens of ~10 Hz and a context window of 2048 you can effectively model a
19 | consistent structure of an audio track for ~3 minutes.
20 | 
21 | ![snac.png](img%2Fsnac.png)
22 | 
23 | ## Pretrained models
24 | 
25 | Currently, all models support only single audio channel (mono).
26 | 
27 | | Model                                                                       | Bitrate   | Sample Rate | Params | Recommended use case     | 
28 | |-----------------------------------------------------------------------------|-----------|-------------|--------|--------------------------|
29 | | [hubertsiuzdak/snac_24khz](https://huggingface.co/hubertsiuzdak/snac_24khz) | 0.98 kbps | 24 kHz      | 19.8 M | 🗣️ Speech               | 
30 | | [hubertsiuzdak/snac_32khz](https://huggingface.co/hubertsiuzdak/snac_32khz) | 1.9 kbps  | 32 kHz      | 54.5 M | 🎸 Music / Sound Effects | 
31 | | [hubertsiuzdak/snac_44khz](https://huggingface.co/hubertsiuzdak/snac_44khz) | 2.6 kbps  | 44 kHz      | 54.5 M | 🎸 Music / Sound Effects |
32 | 
33 | ## Usage
34 | 
35 | Install it using:
36 | 
37 | ```bash
38 | pip install snac
39 | ```
40 | 
41 | To encode (and decode) audio with SNAC in Python, use the following code:
42 | 
43 | ```python
44 | import torch
45 | from snac import SNAC
46 | 
47 | model = SNAC.from_pretrained("hubertsiuzdak/snac_32khz").eval().cuda()
48 | audio = torch.randn(1, 1, 32000).cuda()  # placeholder for actual audio with shape (B, 1, T)
49 | 
50 | with torch.inference_mode():
51 |     codes = model.encode(audio)
52 |     audio_hat = model.decode(codes)
53 | ```
54 | 
55 | You can also encode and reconstruct in a single call:
56 | 
57 | ```python
58 | with torch.inference_mode():
59 |     audio_hat, codes = model(audio)
60 | ```
61 | 
62 | ⚠️ Note that `codes` is a list of token sequences of variable lengths, each corresponding to a different temporal
63 | resolution.
64 | 
65 | ```
66 | >>> [code.shape[1] for code in codes]
67 | [12, 24, 48, 96]
68 | ```
69 | 
70 | ## Acknowledgements
71 | 
72 | Module definitions are adapted from the [Descript Audio Codec](https://github.com/descriptinc/descript-audio-codec).
73 | 
74 | ## Citation
75 | 
76 | If this code contributes to your research, please cite our work:
77 | 
78 | ```
79 | @inproceedings{siuzdak2024snac,
80 |   title={SNAC: Multi-Scale Neural Audio Codec},
81 |   author={Siuzdak, Hubert and Gr{\"o}tschla, Florian and Lanzend{\"o}rfer, Luca A},
82 |   booktitle={Audio Imagination: NeurIPS 2024 Workshop AI-Driven Speech, Music, and Sound Generation},
83 |   year={2024}
84 | }
85 | ```


--------------------------------------------------------------------------------
/img/snac.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/hubertsiuzdak/snac/8f79a718f1ad71f94f79999f0071348227aff22e/img/snac.png


--------------------------------------------------------------------------------
/requirements.txt:
--------------------------------------------------------------------------------
1 | torch
2 | numpy
3 | einops
4 | huggingface_hub
5 | 


--------------------------------------------------------------------------------
/setup.py:
--------------------------------------------------------------------------------
 1 | import io
 2 | import os
 3 | 
 4 | from setuptools import find_packages, setup
 5 | 
 6 | for line in open("snac/__init__.py"):
 7 |     line = line.strip()
 8 |     if "__version__" in line:
 9 |         context = {}
10 |         exec(line, context)
11 |         VERSION = context["__version__"]
12 | 
13 | 
14 | def read(*paths, **kwargs):
15 |     with io.open(os.path.join(os.path.dirname(__file__), *paths), encoding=kwargs.get("encoding", "utf8")) as open_file:
16 |         content = open_file.read().strip()
17 |     return content
18 | 
19 | 
20 | def read_requirements(path):
21 |     return [line.strip() for line in read(path).split("\n") if not line.startswith(('"', "#", "-", "git+"))]
22 | 
23 | 
24 | setup(
25 |     name="snac",
26 |     version=VERSION,
27 |     author="Hubert Siuzdak",
28 |     author_email="hubert.siuzdak@gmail.com",
29 |     description="Multi-Scale Neural Audio Codec",
30 |     url="https://github.com/hubertsiuzdak/snac",
31 |     long_description=read("README.md"),
32 |     long_description_content_type="text/markdown",
33 |     packages=find_packages(),
34 |     install_requires=read_requirements("requirements.txt"),
35 | )
36 | 


--------------------------------------------------------------------------------
/snac/__init__.py:
--------------------------------------------------------------------------------
1 | from .snac import SNAC
2 | 
3 | __version__ = "1.2.1"
4 | 


--------------------------------------------------------------------------------
/snac/attention.py:
--------------------------------------------------------------------------------
 1 | import torch
 2 | from einops import rearrange
 3 | from torch import nn
 4 | 
 5 | 
 6 | class LocalMHA(nn.Module):
 7 |     def __init__(self, dim=1024, window_size=32, dim_head=64, use_rotary_pos_emb=True):
 8 |         super().__init__()
 9 |         self.norm = nn.LayerNorm(dim)
10 |         self.heads = dim // dim_head
11 |         self.window_size = window_size
12 |         self.to_qkv = nn.Linear(dim, dim * 3, bias=False)
13 |         if use_rotary_pos_emb:
14 |             self.rel_pos = SinusoidalEmbeddings(dim_head, scale_base=window_size // 2)
15 |         else:
16 |             self.rel_pos = None
17 |         self.to_out = nn.Linear(dim, dim, bias=False)
18 | 
19 |     def forward(self, x):
20 |         B, C, T = x.shape
21 |         residual = x
22 |         x = self.norm(x.transpose(1, 2))
23 |         windows = T // self.window_size
24 |         q, k, v = self.to_qkv(x).chunk(3, dim=-1)
25 |         q, k, v = map(lambda t: rearrange(t, "b (w n) (h d) -> b h w n d", w=windows, h=self.heads), (q, k, v))
26 |         if self.rel_pos is not None:
27 |             pos_emb, scale = self.rel_pos(k)
28 |             q, k = apply_rotary_pos_emb(q, k, pos_emb, scale)
29 |         out = torch.nn.functional.scaled_dot_product_attention(q, k, v)
30 |         out = rearrange(out, "b h w n d -> b (w n) (h d)")
31 |         out = self.to_out(out)
32 |         return out.transpose(1, 2) + residual
33 | 
34 | 
35 | class SinusoidalEmbeddings(nn.Module):
36 |     def __init__(self, dim, scale_base=None, use_xpos=False):
37 |         super().__init__()
38 |         inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))
39 |         self.register_buffer("inv_freq", inv_freq)
40 |         # xpos related
41 |         self.use_xpos = use_xpos
42 |         self.scale_base = scale_base
43 |         assert not (use_xpos and scale_base is None), "scale base must be defined if using xpos"
44 |         scale = (torch.arange(0, dim, 2) + 0.4 * dim) / (1.4 * dim)
45 |         self.register_buffer("scale", scale, persistent=False)
46 | 
47 |     def forward(self, x):
48 |         seq_len, device = x.shape[-2], x.device
49 |         t = torch.arange(seq_len, device=x.device).type_as(self.inv_freq)
50 |         freqs = torch.einsum("i , j -> i j", t, self.inv_freq)
51 |         freqs = torch.cat((freqs, freqs), dim=-1)
52 |         if not self.use_xpos:
53 |             return freqs, torch.ones(1, device=device)
54 |         power = (t - (seq_len // 2)) / self.scale_base
55 |         scale = self.scale ** rearrange(power, "n -> n 1")
56 |         scale = torch.cat((scale, scale), dim=-1)
57 | 
58 |         return freqs, scale
59 | 
60 | 
61 | def rotate_half(x):
62 |     x = rearrange(x, "b ... (r d) -> b ... r d", r=2)
63 |     x1, x2 = x.unbind(dim=-2)
64 |     return torch.cat((-x2, x1), dim=-1)
65 | 
66 | 
67 | def apply_rotary_pos_emb(q, k, freqs, scale=1):
68 |     q_len = q.shape[-2]
69 |     q_freqs = freqs[..., -q_len:, :]
70 |     inv_scale = scale**-1
71 |     if scale.ndim == 2:
72 |         scale = scale[-q_len:, :]
73 |     q = (q * q_freqs.cos() * scale) + (rotate_half(q) * q_freqs.sin() * scale)
74 |     k = (k * freqs.cos() * inv_scale) + (rotate_half(k) * freqs.sin() * inv_scale)
75 |     return q, k
76 | 


--------------------------------------------------------------------------------
/snac/layers.py:
--------------------------------------------------------------------------------
  1 | import math
  2 | 
  3 | import torch
  4 | import torch.nn as nn
  5 | from torch.nn.utils.parametrizations import weight_norm
  6 | 
  7 | from .attention import LocalMHA
  8 | 
  9 | 
 10 | class Encoder(nn.Module):
 11 |     def __init__(
 12 |         self,
 13 |         d_model=64,
 14 |         strides=[3, 3, 7, 7],
 15 |         depthwise=False,
 16 |         attn_window_size=32,
 17 |     ):
 18 |         super().__init__()
 19 |         layers = [WNConv1d(1, d_model, kernel_size=7, padding=3)]
 20 |         for stride in strides:
 21 |             d_model *= 2
 22 |             groups = d_model // 2 if depthwise else 1
 23 |             layers += [EncoderBlock(output_dim=d_model, stride=stride, groups=groups)]
 24 |         if attn_window_size is not None:
 25 |             layers += [LocalMHA(dim=d_model, window_size=attn_window_size)]
 26 |         groups = d_model if depthwise else 1
 27 |         layers += [
 28 |             WNConv1d(d_model, d_model, kernel_size=7, padding=3, groups=groups),
 29 |         ]
 30 |         self.block = nn.Sequential(*layers)
 31 | 
 32 |     def forward(self, x):
 33 |         return self.block(x)
 34 | 
 35 | 
 36 | class Decoder(nn.Module):
 37 |     def __init__(
 38 |         self,
 39 |         input_channel,
 40 |         channels,
 41 |         rates,
 42 |         noise=False,
 43 |         depthwise=False,
 44 |         attn_window_size=32,
 45 |         d_out=1,
 46 |     ):
 47 |         super().__init__()
 48 |         if depthwise:
 49 |             layers = [
 50 |                 WNConv1d(input_channel, input_channel, kernel_size=7, padding=3, groups=input_channel),
 51 |                 WNConv1d(input_channel, channels, kernel_size=1),
 52 |             ]
 53 |         else:
 54 |             layers = [WNConv1d(input_channel, channels, kernel_size=7, padding=3)]
 55 | 
 56 |         if attn_window_size is not None:
 57 |             layers += [LocalMHA(dim=channels, window_size=attn_window_size)]
 58 | 
 59 |         for i, stride in enumerate(rates):
 60 |             input_dim = channels // 2**i
 61 |             output_dim = channels // 2 ** (i + 1)
 62 |             groups = output_dim if depthwise else 1
 63 |             layers.append(DecoderBlock(input_dim, output_dim, stride, noise, groups=groups))
 64 | 
 65 |         layers += [
 66 |             Snake1d(output_dim),
 67 |             WNConv1d(output_dim, d_out, kernel_size=7, padding=3),
 68 |             nn.Tanh(),
 69 |         ]
 70 |         self.model = nn.Sequential(*layers)
 71 | 
 72 |     def forward(self, x):
 73 |         x = self.model(x)
 74 |         return x
 75 | 
 76 | 
 77 | class ResidualUnit(nn.Module):
 78 |     def __init__(self, dim=16, dilation=1, kernel=7, groups=1):
 79 |         super().__init__()
 80 |         pad = ((kernel - 1) * dilation) // 2
 81 |         self.block = nn.Sequential(
 82 |             Snake1d(dim),
 83 |             WNConv1d(dim, dim, kernel_size=kernel, dilation=dilation, padding=pad, groups=groups),
 84 |             Snake1d(dim),
 85 |             WNConv1d(dim, dim, kernel_size=1),
 86 |         )
 87 | 
 88 |     def forward(self, x):
 89 |         y = self.block(x)
 90 |         pad = (x.shape[-1] - y.shape[-1]) // 2
 91 |         if pad > 0:
 92 |             x = x[..., pad:-pad]
 93 |         return x + y
 94 | 
 95 | 
 96 | class EncoderBlock(nn.Module):
 97 |     def __init__(self, output_dim=16, input_dim=None, stride=1, groups=1):
 98 |         super().__init__()
 99 |         input_dim = input_dim or output_dim // 2
100 |         self.block = nn.Sequential(
101 |             ResidualUnit(input_dim, dilation=1, groups=groups),
102 |             ResidualUnit(input_dim, dilation=3, groups=groups),
103 |             ResidualUnit(input_dim, dilation=9, groups=groups),
104 |             Snake1d(input_dim),
105 |             WNConv1d(
106 |                 input_dim,
107 |                 output_dim,
108 |                 kernel_size=2 * stride,
109 |                 stride=stride,
110 |                 padding=math.ceil(stride / 2),
111 |             ),
112 |         )
113 | 
114 |     def forward(self, x):
115 |         return self.block(x)
116 | 
117 | 
118 | class NoiseBlock(nn.Module):
119 |     def __init__(self, dim):
120 |         super().__init__()
121 |         self.linear = WNConv1d(dim, dim, kernel_size=1, bias=False)
122 | 
123 |     def forward(self, x):
124 |         B, C, T = x.shape
125 |         noise = torch.randn((B, 1, T), device=x.device, dtype=x.dtype)
126 |         h = self.linear(x)
127 |         n = noise * h
128 |         x = x + n
129 |         return x
130 | 
131 | 
132 | class DecoderBlock(nn.Module):
133 |     def __init__(self, input_dim=16, output_dim=8, stride=1, noise=False, groups=1):
134 |         super().__init__()
135 |         layers = [
136 |             Snake1d(input_dim),
137 |             WNConvTranspose1d(
138 |                 input_dim,
139 |                 output_dim,
140 |                 kernel_size=2 * stride,
141 |                 stride=stride,
142 |                 padding=math.ceil(stride / 2),
143 |                 output_padding=stride % 2,
144 |             ),
145 |         ]
146 |         if noise:
147 |             layers.append(NoiseBlock(output_dim))
148 |         layers.extend(
149 |             [
150 |                 ResidualUnit(output_dim, dilation=1, groups=groups),
151 |                 ResidualUnit(output_dim, dilation=3, groups=groups),
152 |                 ResidualUnit(output_dim, dilation=9, groups=groups),
153 |             ]
154 |         )
155 |         self.block = nn.Sequential(*layers)
156 | 
157 |     def forward(self, x):
158 |         return self.block(x)
159 | 
160 | 
161 | def WNConv1d(*args, **kwargs):
162 |     return weight_norm(nn.Conv1d(*args, **kwargs))
163 | 
164 | 
165 | def WNConvTranspose1d(*args, **kwargs):
166 |     return weight_norm(nn.ConvTranspose1d(*args, **kwargs))
167 | 
168 | 
169 | @torch.jit.script
170 | def snake(x, alpha):
171 |     shape = x.shape
172 |     x = x.reshape(shape[0], shape[1], -1)
173 |     x = x + (alpha + 1e-9).reciprocal() * torch.sin(alpha * x).pow(2)
174 |     x = x.reshape(shape)
175 |     return x
176 | 
177 | 
178 | class Snake1d(nn.Module):
179 |     def __init__(self, channels):
180 |         super().__init__()
181 |         self.alpha = nn.Parameter(torch.ones(1, channels, 1))
182 | 
183 |     def forward(self, x):
184 |         return snake(x, self.alpha)
185 | 


--------------------------------------------------------------------------------
/snac/snac.py:
--------------------------------------------------------------------------------
  1 | import json
  2 | import math
  3 | import os
  4 | from typing import List, Tuple
  5 | 
  6 | import numpy as np
  7 | import torch
  8 | from torch import nn
  9 | 
 10 | from .layers import Encoder, Decoder
 11 | from .vq import ResidualVectorQuantize
 12 | 
 13 | 
 14 | class SNAC(nn.Module):
 15 |     def __init__(
 16 |         self,
 17 |         sampling_rate=44100,
 18 |         encoder_dim=64,
 19 |         encoder_rates=[3, 3, 7, 7],
 20 |         latent_dim=None,
 21 |         decoder_dim=1536,
 22 |         decoder_rates=[7, 7, 3, 3],
 23 |         attn_window_size=32,
 24 |         codebook_size=4096,
 25 |         codebook_dim=8,
 26 |         vq_strides=[8, 4, 2, 1],
 27 |         noise=True,
 28 |         depthwise=True,
 29 |     ):
 30 |         super().__init__()
 31 |         self.sampling_rate = sampling_rate
 32 |         self.encoder_dim = encoder_dim
 33 |         self.encoder_rates = encoder_rates
 34 |         self.decoder_dim = decoder_dim
 35 |         self.decoder_rates = decoder_rates
 36 |         if latent_dim is None:
 37 |             latent_dim = encoder_dim * (2 ** len(encoder_rates))
 38 |         self.latent_dim = latent_dim
 39 |         self.hop_length = np.prod(encoder_rates)
 40 |         self.encoder = Encoder(
 41 |             encoder_dim,
 42 |             encoder_rates,
 43 |             depthwise=depthwise,
 44 |             attn_window_size=attn_window_size,
 45 |         )
 46 |         self.n_codebooks = len(vq_strides)
 47 |         self.codebook_size = codebook_size
 48 |         self.codebook_dim = codebook_dim
 49 |         self.vq_strides = vq_strides
 50 |         self.attn_window_size = attn_window_size
 51 |         self.quantizer = ResidualVectorQuantize(
 52 |             input_dim=latent_dim,
 53 |             codebook_size=codebook_size,
 54 |             codebook_dim=codebook_dim,
 55 |             vq_strides=vq_strides,
 56 |         )
 57 |         self.decoder = Decoder(
 58 |             latent_dim,
 59 |             decoder_dim,
 60 |             decoder_rates,
 61 |             noise,
 62 |             depthwise=depthwise,
 63 |             attn_window_size=attn_window_size,
 64 |         )
 65 | 
 66 |     def preprocess(self, audio_data):
 67 |         length = audio_data.shape[-1]
 68 |         lcm = math.lcm(self.vq_strides[0], self.attn_window_size or 1)
 69 |         pad_to = self.hop_length * lcm
 70 |         right_pad = math.ceil(length / pad_to) * pad_to - length
 71 |         audio_data = nn.functional.pad(audio_data, (0, right_pad))
 72 |         return audio_data
 73 | 
 74 |     def forward(self, audio_data: torch.Tensor) -> Tuple[torch.Tensor, List[torch.Tensor]]:
 75 |         length = audio_data.shape[-1]
 76 |         audio_data = self.preprocess(audio_data)
 77 |         z = self.encoder(audio_data)
 78 |         z_q, codes = self.quantizer(z)
 79 |         audio_hat = self.decoder(z_q)
 80 |         return audio_hat[..., :length], codes
 81 | 
 82 |     def encode(self, audio_data: torch.Tensor) -> List[torch.Tensor]:
 83 |         audio_data = self.preprocess(audio_data)
 84 |         z = self.encoder(audio_data)
 85 |         _, codes = self.quantizer(z)
 86 |         return codes
 87 | 
 88 |     def decode(self, codes: List[torch.Tensor]) -> torch.Tensor:
 89 |         z_q = self.quantizer.from_codes(codes)
 90 |         audio_hat = self.decoder(z_q)
 91 |         return audio_hat
 92 | 
 93 |     @classmethod
 94 |     def from_config(cls, config_path):
 95 |         with open(config_path, "r") as f:
 96 |             config = json.load(f)
 97 |         model = cls(**config)
 98 |         return model
 99 | 
100 |     @classmethod
101 |     def from_pretrained(cls, repo_id, **kwargs):
102 |         from huggingface_hub import hf_hub_download
103 | 
104 |         if not os.path.isdir(repo_id):
105 |             config_path = hf_hub_download(repo_id=repo_id, filename="config.json", **kwargs)
106 |             model_path = hf_hub_download(repo_id=repo_id, filename="pytorch_model.bin", **kwargs)
107 |             model = cls.from_config(config_path)
108 |             state_dict = torch.load(model_path, map_location="cpu")
109 |         else:
110 |             model = cls.from_config(os.path.join(repo_id, "config.json"))
111 |             state_dict = torch.load(os.path.join(repo_id, "pytorch_model.bin"), map_location="cpu")
112 |         model.load_state_dict(state_dict)
113 |         model.eval()
114 |         return model
115 | 


--------------------------------------------------------------------------------
/snac/vq.py:
--------------------------------------------------------------------------------
 1 | from typing import List
 2 | 
 3 | import torch
 4 | import torch.nn as nn
 5 | import torch.nn.functional as F
 6 | from einops import rearrange
 7 | 
 8 | from .layers import WNConv1d
 9 | 
10 | 
11 | class VectorQuantize(nn.Module):
12 |     def __init__(self, input_dim: int, codebook_size: int, codebook_dim: int, stride: int = 1):
13 |         super().__init__()
14 |         self.codebook_size = codebook_size
15 |         self.codebook_dim = codebook_dim
16 |         self.stride = stride
17 | 
18 |         self.in_proj = WNConv1d(input_dim, codebook_dim, kernel_size=1)
19 |         self.out_proj = WNConv1d(codebook_dim, input_dim, kernel_size=1)
20 |         self.codebook = nn.Embedding(codebook_size, codebook_dim)
21 | 
22 |     def forward(self, z):
23 |         if self.stride > 1:
24 |             z = torch.nn.functional.avg_pool1d(z, self.stride, self.stride)
25 | 
26 |         # Factorized codes (ViT-VQGAN) Project input into low-dimensional space
27 |         z_e = self.in_proj(z)  # z_e : (B x D x T)
28 |         z_q, indices = self.decode_latents(z_e)
29 |         z_q = z_e + (z_q - z_e).detach()  # noop in forward pass, straight-through gradient estimator in backward pass
30 | 
31 |         z_q = self.out_proj(z_q)
32 | 
33 |         if self.stride > 1:
34 |             z_q = z_q.repeat_interleave(self.stride, dim=-1)
35 | 
36 |         return z_q, indices
37 | 
38 |     def embed_code(self, embed_id):
39 |         return F.embedding(embed_id, self.codebook.weight)
40 | 
41 |     def decode_code(self, embed_id):
42 |         return self.embed_code(embed_id).transpose(1, 2)
43 | 
44 |     def decode_latents(self, latents):
45 |         encodings = rearrange(latents, "b d t -> (b t) d")
46 |         codebook = self.codebook.weight  # codebook: (N x D)
47 | 
48 |         # L2 normalize encodings and codebook (ViT-VQGAN)
49 |         encodings = F.normalize(encodings)
50 |         codebook = F.normalize(codebook)
51 | 
52 |         # Compute euclidean distance with codebook
53 |         dist = (
54 |             encodings.pow(2).sum(1, keepdim=True)
55 |             - 2 * encodings @ codebook.t()
56 |             + codebook.pow(2).sum(1, keepdim=True).t()
57 |         )
58 |         indices = rearrange((-dist).max(1)[1], "(b t) -> b t", b=latents.size(0))
59 |         z_q = self.decode_code(indices)
60 |         return z_q, indices
61 | 
62 | 
63 | class ResidualVectorQuantize(nn.Module):
64 |     def __init__(
65 |         self,
66 |         input_dim: int = 512,
67 |         codebook_size: int = 1024,
68 |         codebook_dim: int = 8,
69 |         vq_strides: List[int] = [1, 1, 1, 1],
70 |     ):
71 |         super().__init__()
72 |         self.n_codebooks = len(vq_strides)
73 |         self.codebook_dim = codebook_dim
74 |         self.codebook_size = codebook_size
75 |         self.quantizers = nn.ModuleList(
76 |             [VectorQuantize(input_dim, codebook_size, codebook_dim, stride) for stride in vq_strides]
77 |         )
78 | 
79 |     def forward(self, z):
80 |         z_q = 0
81 |         residual = z
82 |         codes = []
83 |         for i, quantizer in enumerate(self.quantizers):
84 |             z_q_i, indices_i = quantizer(residual)
85 |             z_q = z_q + z_q_i
86 |             residual = residual - z_q_i
87 |             codes.append(indices_i)
88 | 
89 |         return z_q, codes
90 | 
91 |     def from_codes(self, codes: List[torch.Tensor]) -> torch.Tensor:
92 |         z_q = 0.0
93 |         for i in range(self.n_codebooks):
94 |             z_p_i = self.quantizers[i].decode_code(codes[i])
95 |             z_q_i = self.quantizers[i].out_proj(z_p_i)
96 |             z_q_i = z_q_i.repeat_interleave(self.quantizers[i].stride, dim=-1)
97 |             z_q += z_q_i
98 |         return z_q
99 | 


--------------------------------------------------------------------------------