Metadata-Version: 2.4
Name: neural_tilde
Version: 0.0.2
Summary: Exporting PyTorch neural audio synthesis models in ExecuTorch for neural_tilde.
Author-email: Jasper Shuoyang Zheng <jasper.zheng.u@gmail.com>
License-Expression: CC-BY-NC-4.0
Project-URL: Homepage, https://github.com/jasper-zheng/neural_tilde
Project-URL: Repository, https://github.com/jasper-zheng/neural_tilde
Keywords: executorch,pytorch,neural audio,max msp,audio synthesis,rave
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Operating System :: OS Independent
Classifier: Topic :: Multimedia :: Sound/Audio
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.12.0
Provides-Extra: coreml
Requires-Dist: coremltools>=9.0; extra == "coreml"
Provides-Extra: dev
Requires-Dist: build; extra == "dev"
Requires-Dist: twine; extra == "dev"
Dynamic: license-file

# `neural`: neural audio synthesis in Max/MSP

**Max externals for running neural synthesis models realtime/offline.**   
Models are exported from PyTorch with [ExecuTorch](https://pytorch.org/executorch/) and run inside Max with hardware acceleration on CPU / GPU / ANE (Apple Neural Engine).



This package has two families of Max objects:

**Live models:** streaming realtime DSP models (e.g. neural vocoder, neural codecs...):

| object | use it for |
| --- | --- |
| `neural.live~` | a single model instance on one audio stream |
| `mc.neural.live~` | the same model broadcast across the channels of an `mc` signal |
| `mcs.neural.live~` | a fixed batch of streams processed in one pass (takes a batch-size argument) |

**Generative models:** one-shot offline generators (e.g. latent-diffusion text-to-audio, audio-to-audio...):

| object | use it for |
| --- | --- |
| `neural.gen~` | run a generative model and write the output into a `buffer~` |

**Utility**:

| object | use it for |
| --- | --- |
| `neural.tokenizer` | turn a text prompt into token ids and attention mask |
| `neural.gaussianize` | map uniform `jit.noise` to Gaussian noise |

**Supported models:**  
 - [Live] [RAVE-CoreML](https://github.com/jasper-zheng/RAVE-CoreML) A streamable codec for timbre transfer / latent-space audio manipulation
 - [Live] [SAME-S](https://github.com/jasper-zheng/streamable-same-s): A streamable codec for audio encoding/decoding
 - [Gen] [Stable Audio 3](https://github.com/jasper-zheng/stable-audio-3-tilde): Latent diffusion transformer for text-to-audio, audio-to-audio, inpainting
 - [Gen/Live] [CodiCodec-Flow](https://github.com/jasper-zheng/codicodec-flow) [Available soon]

**Note:** Currently only available for MaxMSP on Apple Silicon, macOS. For Windows/CUDA, it's on the todo list. 

**Acknowledgement:** `neural.live~` and its `mc`, `mcs` variants reused an extensive amount of code from [`nn~`](https://github.com/acids-ircam/nn_tilde), migrated from TorchScript/libtorch to more modern ExecuTorch. `nn~` is the work by Antoine Caillon & Axel Chemla--Romeu-Santos (acids-ircam), licensed **CC BY-NC 4.0**.


## How It Works

 - **Export**: Export a PyTorch neural audio model using [ExecuTorch](https://docs.pytorch.org/executorch/main/getting-started.html), with a target hardware backend (see below). This results in a `.pte` (the runnable program + weights) file.  
 - **Config**: Create a `.json` sidecar that details inlet/outlet, ratio, attributes, condition shapes, etc. 
 - **Load**: Hand `neural.live~` / `neural.gen~` the `.pte` and `.json` files in Max.

### Supported back-ends (chosen when exporting the model):

- **[Core ML](https://developer.apple.com/documentation/coreml)** - Leverage Apple Neural Engine (ANE) for hardware acceleration. Needs **macOS 15+** at runtime.
- **[MLX](https://github.com/ml-explore/mlx)** - Apple-Silicon GPU; Good for one-shot / buffer processing.
- **[XNNPACK](https://github.com/google/XNNPACK)** - CPU inference with optimized kernels.
- **portable** - plain, unoptimized C++ kernels, maximum compatibility
- ~~CUDA~~ - Windows CUDA in progress


## Install the externals

<!-- Download the pre-released externals. To install:

1. Unzip and copy the `neural_tilde` folder into your Max **Packages** folder, e.g.
   `~/Documents/Max 9/Packages/`.
2. Restart Max. The streaming objects (`neural.live~`, `mc.neural.live~`, `mcs.neural.live~`) and the
   generative objects (`neural.gen~`, `neural.tokenizer`) are now available.

If macOS quarantines the externals (objects fail to instantiate), clear the quarantine flag or ad-hoc re-sign them. In Terminal, type:

```bash
cd ~/Documents/Max\ 9/Packages/neural_tilde
xattr -dr com.apple.quarantine externals/*.mxo
# or:  codesign --force --deep -s - externals/*.mxo
``` -->

Compiled externals / Max help patches will be available soon. For now, please build from source (see [Build.md](Build.md)).


## Objects

`neural.live~` and `neural.gen~` are the main objects for running neural synthesis models. They are for different use cases: `neural.live~` is for realtimestreaming models (e.g. neural vocoder, neural codec, etc.), while `neural.gen~` is for offline generative models (e.g. one-shot text-to-audio, audio inpainting).

| `neural.live~` | `neural.gen~` |
| --- | --- |
| ![neural.live~ minimal patch](assets/live-minimal.png) | ![neural.gen~ minimal patch](assets/gen-minimal.png) |

<!-- Alt text for the above images for neural.gen~:
```
[prompt lofi house loop]              <- message to the tokenizer
       |
[neural.tokenizer my_tokenizer]
       |        |     <- outlet 0 = {token_ids}, outlet 1 = {attention_mask}
       |        |
[neural.gen~ my_diffusion @seed 0]
       |               (inlet 0 = token_ids, inlet 1 = attention_mask)
   (done bang)        ... writes into [buffer~ gen 176400 2]
```

1. Create `[buffer~ gen 176400 2]` (stereo 4s at 44.1 kHz). Send `neural.gen~` the message `set gen` to target it.
2. Send the prompt to **`neural.tokenizer`** (`prompt lofi house loop`); wire its two outlets into the matching `neural.gen~` condition inlets.
3. Send `generate` message (or `bang`) to `neural.gen~`. The `done` outlet produce a bang once generation finishes, and you can play it with `groove~ gen` / `play~`.

You can pass `neural.tokenizer` the `*.tokenizer.json`. 
-->

### `neural.live~`:

- **Arguments:** `[neural.live~ <model.pte> <method>]`. 
- **Signal inlets / outlets** are created according to the `.json` sidecar. 
- **Conditions / attributes / noise** (if any) are created according to the `.json` sidecar (see [Input Roles](#input-roles)).
- **Turn computation on:** the `enable_model` attribute defaults to **off**, so the
  object outputs silence until you set `enable_model` to 1
- **Buffer size** is read from the model's `.json` sidecar, this is **fixed when the model is exported**.


### `neural.gen~`:

- **Arguments:** `[neural.gen~ <model.pte> <method> <buffer>]`.
- **Buffer**: Create a `buffer~` to hold the generated outputs, set its name to `neural.gen~` by the third argument, or with the `set <buffer>` message. 
- **Conditions / attributes / noise** (if any) are created according to the `.json` sidecar (see [Input Roles](#input-roles)).

> **Notes.** A `buffer~` carries Max's project sample-rate while the mode might have a fixed rate. If your patch runs at a different rate, the data is correct but plays back at the wrong pitch.

### `mc.neural.live~` / `mcs.neural.live~`:
Use `mc.neural.live~` to run one model across every channel of an `mc.` patch cord, or
`mcs.neural.live~` to process a fixed batch of streams in a single forward pass.


<img src="assets/mcs.png" width="380" alt="mcs.neural.live~" />


### `neural.tokenizer`:

Use `neural.tokenizer` to turn a prompt into `token_ids` + `attention_mask`, emitting each as a single-key dictionary on its two outlets. Both outputs are used as **condition inputs** to `neural.gen~` or `neural.live~`.  

Tokenizers are configured by `*.tokenizer.json` and `*.tokenizer.config.json`:

<img src="assets/tokenizer-minimal.png" width="600" alt="neural.tokenizer" />


### `neural.gaussianize`:

Diffusion models often requires Gaussian noise as inputs. However, `jit.noise` / `random` produces uniformed distribution. Therefore, `neural.gaussianize` helps map a uniform distribution to a Gaussian one.

<img src="assets/noise.png" width="500" alt="neural.gaussianize" />



## Use a pre-trained model

### Synthesis Model

`neural.live~` and `neural.gen~` run model exported from PyTorch, it needs:
 - `*.pte` the model weights + program
 - `*.json` the sidecar config defining model's inlets/outlets, attributes, conditions, etc.

They should have **the same** filename, put together in a folder. 

**Make sure to add your model's path in Max:** In Max, open **"Options - File Preferences"**, click **"Add Path"** on the bottom left corner, add your model's folder. 


| file | what it is |
| --- | --- |
| `my_diffusion.pte` | the ExecuTorch program + weights |
| `my_diffusion.json` | model sidecar: typed inputs/output (no tokenizer settings) |

### Tokenizer

If your model uses a prompt-based inputs (e.g. Stable Audio 3), the exported model also has a tokenizer bundle:

| file | what it is |
| --- | --- |
| `my_tokenizer.tokenizer.json` | the HuggingFace fast tokenizer |
| `my_tokenizer.tokenizer.config.json` | tokenizer settings `neural.tokenizer` reads |


## Methods

A model may expose **several methods** with different inputs (e.g. `prompt2audio` and `audio2audio` for diffusion models, or `encode` and `decode` for codecs). `neural.gen/live~` selects only one when initialized. 

Methods are declared in the model's `.json` sidecar, see [PROTOCOL.md - 2.3 Sidecar JSON template](PROTOCOL.md#23-sidecar-json-template) for details.

## Input Roles

A method may take **several inputs**, each with a **role**. The role is declared in the model's `.json` sidecar, see [PROTOCOL.md - 2.3 Sidecar JSON template](PROTOCOL.md#23-sidecar-json-template) for details.  

**Five input roles are supported:**

| role | kinds | in `neural.gen/live~` | examples |
|---|---|---|---|
| `attribute` | live / gen | A scalar number exposed as a Max attribute. Its current value is fed to the model in every run. | temperature, cfg scale, duration, etc. |
| `noise` | live / gen | By default, all noises are derived from the "seed" attribute. Alternatively, if a noise input has a shape of [planes × H × W], an inlet will be created, customized noises from **Jitter matrices** are allowed.| initial noise for diffusion, latent noise for codec, etc. |
| `condition` | live / gen | Additional control vector supplied from inlets, as a `list` (by position) or a `dictionary` (by name). | token IDs, attention mask, inpainting mask, time-varying controls, etc. |
| `signal` | live only | Multi-channel signal from signal inlets. Supports time-domain compression (e.g., encoder/decoder in a codec that downsamples audio-rate 44.1 kHz to latent-rate 21.5 Hz). Exactly one `signal`-role input per method. | 
| `buffer` | gen only | An init waveform read from a Max `buffer~` for audio-to-audio. Resampled and cropped to the declared shape. | 


**Two output roles:**

| role | kinds | in `neural.gen/live~` |
|---|---|---|
| `signal` | live | Signal output from the model, up/downsampled if required. |
| `audio` | gen | Write `channels × length` into a `buffer~` at `sample_rate`. |


## Export your own neural synthesis model

To export your own neural synthesis model to `neural.live/gen~`, use the helper in [`python_tools/`](python_tools/). It provides the `neural_tilde` Python module. It can be installed with pip:

```bash
pip install neural-tilde
# Install the `neural_tilde` module, cached_conv, executorch, coremltools
```

### Usage:

 - Subclass `neural_tilde.LiveModule` / `.GenModule`
 - Build your model
 - **Register inputs:** give each method extra inputs using `register_attribute` / `register_noise` / `register_condition`
 - **Register methods:** use `register_method(...)` register each method and their inputs / outputs. Registered roles are listed in order via the `inputs=[...]` argument; the method takes them after the audio tensor (e.g., `forward(self, x, gain, bias)`). 
 - **Export model:** Use `export_to_pte(..., delegate="coreml")` to export, which will result in the model weights `.pte` and the sidecar `.json`


### Live Example:

A complete `live~` example is in
[`examples/export_live_example.py`](examples/export_live_example.py):

```python
import cached_conv as cc
from neural_tilde import LiveModule

cc.use_cached_conv(True)              # Use cached_conv for streaming

class MyLiveModel(LiveModule):
    def __init__(self, hidden: int = 16, kernel_size: int = 3):
        super().__init__()
        pad = cc.get_padding(kernel_size)
        self.net = cc.CachedSequential(
            cc.Conv1d(2, hidden, kernel_size, padding=pad),
            nn.GELU(),
            cc.Conv1d(hidden, 2, kernel_size, padding=pad)
        )
    def forward(self,
                x: torch.Tensor,        # [batch, 2, time]   (signal)
                gain: torch.Tensor,     # [1]                (attribute)
                bias: torch.Tensor      # [batch, 1, 1]      (condition)
                ) -> torch.Tensor:
        x = x + bias.reshape(-1, 1, 1)
        y = self.net(x)                 # [batch, 2, time]
        return y * gain.reshape(1, 1, 1)

model = MyLiveModel()

# Register the model's inputs and outputs:
model.register_attribute("gain", default=1.0, minimum=0.0, maximum=2.0)
model.register_condition("bias", shape=[1, 1], dtype="float32")
model.register_method("forward", in_channels=2, in_ratio=1,
                      out_channels=2, out_ratio=1,
                      inputs=["gain", "bias"]
                      )

# Export the model with Core ML delegate:
model.export_to_pte("mini-live", delegate="coreml", buffer_size=4096)
```

To run the above example:
```bash
python examples/export_live_example.py mini-live
```

Results in:

<img src="assets/mini-live.png" width="600" alt="export_live_example.py output" />


<!-- **Why this streams:** `cached_conv` replaces zero-padding with *cached* padding (it
keeps each conv's previous-block tail in a buffer). On the Core ML delegate those cache
buffers are taken over as native Core ML **state** that persists across blocks, so
`neural.live~` runs the model click-free in real time. (MLX does not persist this state, that's why we perfer Core ML instead of MLX) -->

### Gen Example:

A complete `gen~` example is in [`examples/export_gen_example.py`](examples/export_gen_example.py):

```python
from neural_tilde import GenModule

VOCAB = 32              # toy vocabulary
LATENT = 128            # latent channels
STEPS = 32              # latent time steps
LENGTH = STEPS * STEPS  # 1024 audio samples

class MyGenModel(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.embed = nn.Embedding(VOCAB, LATENT)
        self.up = nn.ConvTranspose1d(LATENT, 2, kernel_size=STEPS, stride=STEPS)

    def _decode(self, z: torch.Tensor) -> torch.Tensor:
        return torch.tanh(self.up(z))           # [1, 2, LENGTH]

    def prompt2audio(self,
                     vocab_ids: torch.Tensor,   # [1, 64]             (condition)
                     noise_x: torch.Tensor,     # [1, LATENT, STEPS]  (noise)
                     ) -> torch.Tensor:
        emb = self.embed(vocab_ids)             # [1, 64, LATENT]
        cond = emb.mean(1)                      # [1, LATENT]
        z = noise_x + cond.unsqueeze(-1)        # [1, LATENT, STEPS]
        return self._decode(z)

model = GenModule(TinyGen())
model.register_condition("vocab_ids", [1, 64], "int64")
model.register_noise("noise_x", [1, LATENT, STEPS])

model.register_method("prompt2audio",
                      inputs=["vocab_ids", "noise_x"],
                      out_channels=2, out_length=LENGTH,
                      out_sample_rate=44100
                      )

path = model.export_to_pte(out, delegate="coreml")
```
To run the above example:
```bash
python examples/export_gen_example.py mini-gen
```

Results in: 

<img src="assets/mini-gen.png" width="600" alt="export_gen_example.py output" />

### Tokenizer Example:

The `register_tokenizer(...)` method registers the tokenizer if the model has one. If you have a custom tokenizer, you can use `neural_tilde.Tokenizer` to export it to `.tokenizer.json` and `.tokenizer.config.json`:

A complete example is in [`examples/export_tokenizer_example.py`](examples/export_tokenizer_example.py):

```python 
from neural_tilde import Tokenizer

tok = Tokenizer(tokenizer_file, 
                max_length=256,
                ids_key="input_ids", mask_key="attention_mask")
tok.write_files(out_stem)
```


## Full Protocol

The full input-role / sidecar template is in [PROTOCOL.md](PROTOCOL.md).

## Build the externals from source

Please refer to [Build.md](Build.md)

## Credits

C++ externals and the Python tools derive from `nn~` by **Antoine Caillon** and
**Axel Chemla--Romeu-Santos** (ACIDS-IRCAM). ExecuTorch migration and the `neural`
fork: this repository. License: **CC BY-NC 4.0**.
