Metadata-Version: 2.4
Name: faustax
Version: 0.0.1
Summary: Differentiable audio effects for JAX, written in Faust and compiled to Flax NNX modules
Author-email: David Braun <braun@ccrma.stanford.edu>
License-Expression: MIT
Project-URL: Homepage, https://github.com/DBraun/faustax
Project-URL: Documentation, https://dbraun.github.io/faustax/
Project-URL: Repository, https://github.com/DBraun/faustax
Project-URL: Issues, https://github.com/DBraun/faustax/issues
Project-URL: Changelog, https://github.com/DBraun/faustax/blob/main/CHANGELOG.md
Keywords: audio,dsp,faust,jax,differentiable
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Programming Language :: Python :: 3.14
Classifier: Topic :: Multimedia :: Sound/Audio
Classifier: Topic :: Software Development :: Libraries
Requires-Python: >=3.11
Description-Content-Type: text/markdown
License-File: LICENSE
License-File: LICENSE-APACHE-2.0
License-File: NOTICE
Requires-Dist: flax>=0.12.7
Requires-Dist: jax>=0.10.2
Requires-Dist: librosa>=0.10.1
Requires-Dist: numpy>=1.24
Requires-Dist: safetensors>=0.4
Requires-Dist: scipy>=1.10
Provides-Extra: audiotree
Requires-Dist: audiotree>=1.0.0; extra == "audiotree"
Provides-Extra: realtime
Requires-Dist: sounddevice>=0.4; extra == "realtime"
Provides-Extra: vst-datasets
Requires-Dist: audiotree>=1.0.0; extra == "vst-datasets"
Requires-Dist: dawdreamer>=0.9; extra == "vst-datasets"
Provides-Extra: viz
Requires-Dist: matplotlib>=3.7; extra == "viz"
Dynamic: license-file

# Faustax

Faustax supplies batched, differentiable audio processors for JAX.
The applications are parameter estimation, style transfer, automatic mixing, and data augmentation.

Each effect is a small program in [Faust](https://faust.grame.fr), built mostly from [Faust Libraries](https://faustlibraries.grame.fr/).
The Faust NNX backend compiles each program ahead-of-time to a [Flax NNX](https://flax.readthedocs.io/en/latest/nnx/index.html) module that works with `jax.jit` and `jax.vmap`.
Behind the scenes, the module generates one sample of audio at a time with `nnx.scan`.
One intentional exception exists: the noise-shaped reverb applies a synthesized 65536-tap FIR, which is hand-written JAX code, not Faust.

**End users do not need the Faust compiler.**
The repository contains the generated Python modules in `src/faustax/_generated/`.
The runtime dependencies are only `jax`, `flax`, `numpy`, `librosa`, and `safetensors`.

**Documentation: <https://dbraun.github.io/faustax/>**

## Installation

```bash
pip install faustax  # or: uv add faustax
```

Faustax needs Python 3.11 or later.
The base installation runs on CPU JAX and doesn't need the Faust compiler.
Each extra adds the dependencies of one optional feature:

| Install                 | Adds                                                                   |
|:------------------------|:-----------------------------------------------------------------------|
| `faustax[audiotree]`    | `audiotree`, for the transform adapter (`faustax.audiotree`)           |
| `faustax[realtime]`     | `sounddevice`, for the duplex stream in `faustax.realtime`             |
| `faustax[viz]`          | `matplotlib`, for the `--plot` option of the fitting examples          |
| `faustax[vst-datasets]` | `audiotree` and `dawdreamer`, for the `faustax-vst-*` dataset commands |

Ask for several extras at the same time:

```bash
pip install "faustax[realtime,viz]"
```

The CUDA wheels of JAX are not an extra of Faustax, because the wheel to install depends on your CUDA version.
Install them next to Faustax:

```bash
pip install faustax "jax[cuda13]"
```

## Usage

```python
from flax import nnx
import jax.numpy as jnp
from faustax import Compressor

comp = Compressor(sample_rate=44100)

x = jnp.zeros((4, 2, 44100))  # (batch, channels, samples)

# Physical parameters (scalars broadcast; arrays of shape (batch,) vary per item)
y = comp.process(x, threshold_db=-24.0, ratio=4.0, attack_ms=10.0, release_ms=100.0)

# dasp-style: normalized parameters on [0, 1], shape (batch, num_params),
# e.g. straight from a neural network controller
rngs = nnx.Rngs(0)
params = rngs.uniform((4, comp.num_params))
y = comp.process_normalized(x, params)
```

Every processor exposes `param_ranges`.
The library introspects these ranges from the Faust slider declarations; no person maintains them by hand.
Every processor also exposes the underlying NNX module as `.module`.
Use `.module` for streaming (`process_block`) and for NNX-native training workflows.

## Effects

| Processor                                                       | Source                                   |
|:----------------------------------------------------------------|:-----------------------------------------|
| `Gain`                                                          | `src/faustax/dsp/gain.dsp`               |
| `Distortion` (soft-clipping tanh with drive)                    | `src/faustax/dsp/distortion.dsp`         |
| `ParametricEQ` (RBJ low shelf, 4 peaking bands, RBJ high shelf) | `src/faustax/dsp/parametric_eq.dsp`      |
| `Compressor` (feed-forward, soft knee, makeup gain)             | `src/faustax/dsp/compressor.dsp`         |
| `Freeverb` (Schroeder reverb with dry/wet mix)                  | `src/faustax/dsp/freeverb.dsp`           |
| `NoiseShapedReverb` (12-band noise shaping, WASPAA 2021)        | `src/faustax/reverb.py` (JAX, not Faust) |
| `StereoPanner` (equal-power mono-to-stereo pan)                 | `src/faustax/dsp/stereo_panner.dsp`      |
| `StereoWidener` (mid-side width control)                        | `src/faustax/dsp/stereo_widener.dsp`     |
| `functional.stereo_bus` (per-track sends, stereo sum)           | pure JAX (dynamic track count)           |
| `diffvox.EQ` (2 peaks, 2 shelves, low/high pass)                | `src/faustax/dsp/diffvox/eq.dsp`         |
| `diffvox.Compressor` (compressor-expander w/ lookahead)         | `src/faustax/dsp/diffvox/compressor.dsp` |
| `diffvox.PingPongDelay` (cross-fed stereo delay send)           | `src/faustax/dsp/diffvox/pingpong.dsp`   |
| `diffvox.FDN` (6-line reverb send with decay FIRs)              | `src/faustax/dsp/diffvox/fdn.dsp`        |

Faustax has the role of [dasp-pytorch](https://github.com/csteinmetz1/dasp-pytorch) in the JAX ecosystem, with a different implementation strategy and different performance characteristics.
The first block of the table is the dasp-pytorch parity set.
Every processor that dasp-pytorch implements exists here with the same parameter names and the same semantics.
(The `expander` of dasp-pytorch is unimplemented upstream.)
Two upgrades are intentional.
The EQ and the compressor are exact per-sample recurrences, not frequency-sampled approximations as in dasp-pytorch.
The EQ is coefficient-exact against the RBJ formulas of dasp-pytorch.
The Faustax compressor applies `release_ms`; dasp-pytorch accepts `release_ms` but ignores it.
The Faustax reverb keys its shaping noise with `rng=`; dasp-pytorch uses unseeded `torch.randn`.
This means Faustax runs are reproducible by default.
[Performance versus dasp-pytorch](#performance-versus-dasp-pytorch) below gives the measured ratios for both libraries.

The four `faustax.diffvox` processors port the vocal effects chain of [DiffVox](https://github.com/SonyResearch/diffvox) (Yu et al., DAFx25).
Each processor matches the real-time reference implementation of diffvox to float32 precision.
`faustax.diffvox` loads its two curated preset datasets directly from a diffvox checkout.
The datasets contain 385 internal presets and 70 MedleyDB presets, fit to real vocal productions.
The datasets include the Gaussian parameter prior, so you can sample new vocal-chain settings.
`diffvox.Chain` renders the full chain: the EQ, the compressor, the panned direct signal, and the delay and reverb sends.

```python
from faustax.diffvox import Chain, load_preset_dataset

ds = load_preset_dataset("/path/to/diffvox/presets/internal")
chain = Chain(sample_rate=44100)
wet_stereo = chain.process(dry_mono, ds[7])   # (batch, 1, T) -> (batch, 2, T)
```

All effects share these semantics.
Every effect is zero-latency and causal: it has no lookahead and no latency compensation.
The output length always equals the input length.
Thus, the effect truncates a reverb or delay tail at the end of the excerpt.
Zero-pad the input if you need the decay.
Deterministic effects ignore the `rng` argument.
The `rng` argument exists for stochastic DSPs: the noise-shaped reverb, and Faust programs that call `random_*` foreign functions.
For double precision, construct the processor with `faust_float=jnp.float64` and enable `jax_enable_x64`.
The whole state carry then follows that `dtype`.

## Documentation

|                                                                                              |                                                                                                               |
|:---------------------------------------------------------------------------------------------|:--------------------------------------------------------------------------------------------------------------|
| [Introduction](https://dbraun.github.io/faustax/introduction.html)                           | Install, the `Processor` API, normalized parameters                                                           |
| [Status](https://dbraun.github.io/faustax/status.html)                                       | The known defect, the self-checking vectorizer, and what does not differentiate                               |
| [Parameter estimation](https://dbraun.github.io/faustax/parameter-estimation.html)           | Recover effect settings, response curves and instrument physics by gradient descent; saving what you trained  |
| [Learnable soundfiles and menus](https://dbraun.github.io/faustax/learnable-parameters.html) | Trainable wavetables (`[param:1]`) and categorical `nentry` menus as Gumbel-softmax                           |
| [audiotree integration](https://dbraun.github.io/faustax/audiotree.html)                     | Any processor as a batched random transform in a grain data pipeline                                          |
| [argbind configuration](https://dbraun.github.io/faustax/argbind.html)                       | Configure every slider from YAML                                                                              |
| [VST reference datasets](https://dbraun.github.io/faustax/vst-datasets.html)                 | Sweep a VST3 plugin's parameters and export (dry, wet) pairs to fit against                                   |
| [Real-time deployment](https://dbraun.github.io/faustax/realtime.html)                       | Streaming `process_block` and the sounddevice callback                                                        |
| [Performance](https://dbraun.github.io/faustax/performance.html)                             | CPU/GPU benchmarks, `unroll`, data parallelism across cores                                                   |
| [Development](https://dbraun.github.io/faustax/development.html)                             | Repository layout, regenerating modules, adding an effect                                                     |
| [API reference](https://dbraun.github.io/faustax/api.html)                                   | Every public module, class and function                                                                       |
| [Custom gradient primitives](https://dbraun.github.io/faustax/ops.html)                      | `faustax.ops` — memory-light custom-VJP recursive filters                                                     |
| [NNX backend notes](https://dbraun.github.io/faustax/nnx-backend-notes.html)                 | What makes a generated module fast or slow: carry layout, `unroll`, GPU scan latency                          |
| [Future work](https://dbraun.github.io/faustax/future-work.html)                             | Scoped but unstarted tasks, and the reasons behind them                                                       |

The [`examples/`](examples/README.md) directory contains runnable scripts.

## Performance versus dasp-pytorch

The two libraries have different performance characteristics because their implementations are different.
An exact per-sample recurrence is sequential in time.
The frequency-domain approximations of dasp-pytorch are a small number of large batched tensor operations.

On CPU, the Faustax EQ costs approximately 4x the dasp-pytorch forward time and 6-8x the dasp-pytorch gradient time.
This cost is the cost of coefficient-exact IIR output.
The Faustax compressor is at parity with the dasp-pytorch compressor.
The Faustax FFT-convolution reverb trains 20-30x faster than the direct convolution of dasp-pytorch.

On GPU, a sequential scan costs approximately 0.5-1 s per call at any batch size.
Thus, dasp-pytorch is faster at small batch sizes.
The Faustax wall time stays almost constant as the batch size grows.
The gap for recursive effects decreases to approximately 5x at batch size 256.
The Faustax reverb is faster than the dasp-pytorch reverb at every batch size.
Run the Faustax scan on CPU when the number of parallel lanes is below approximately 64.

Use Faustax when you need exactness, streaming parity with deployment, fast reverb or dynamics training, or wide-batch or CPU throughput.
Use dasp-pytorch frequency sampling when small-batch EQ gradients on GPU are the most important factor.
The [Performance](https://dbraun.github.io/faustax/performance.html) page shows the measured tables for each claim in this section.

## Development

```bash
uv sync
uv run pytest
```

`uv sync` is the only necessary setup step.
The test suite runs without a Faust compiler.
The tests that require a Faust compiler skip automatically.
See [CONTRIBUTING.md](CONTRIBUTING.md) for how to add effects, how to get a Faust with the NNX backend, and the licensing rules for contributions.

## Status

**Alpha (0.0.x).**
The public API can change between releases.
One known defect limits what you can depend on: **Freeverb diverges from Faust's C++ backend** at sample 1116 and after, the length of its shortest comb delay.
Every other effect matches the C++ backend to < 2e-5 max absolute error.
[Status](https://dbraun.github.io/faustax/status.html) lists that defect in full, and three more facts to know before you depend on the library.

## License

MIT — see [LICENSE](LICENSE).

Some components carry additional third-party terms.
[NOTICE](NOTICE) records all of these terms:

- **Generated modules** (`src/faustax/_generated/`) are Faust compiler output.
  Their scaffolding comes from GRAME's Faust architecture file, whose grant explicitly permits redistribution under terms of your choice.
  Their DSP body is a translation of Faust standard library code that carries an LGPL exception, and the exception grants the same freedom.
  Individual library functions declare their own MIT or MIT-style STK-4.3 terms.
  `NOTICE` reproduces the copyright notices of those terms.
- **`faustax.ops` and `faustax.dynamics`** port gradient rules and conventions from Chin-Yun Yu's MIT-licensed torchlpc, philtorch and torchcomp.
- **The dasp-pytorch parity layer** (`faustax.reverb` and the parity `.dsp` sources) ports Apache-2.0-licensed [dasp-pytorch](https://github.com/csteinmetz1/dasp-pytorch).
  A copy of that license ships as [LICENSE-APACHE-2.0](LICENSE-APACHE-2.0).
- **`fdn_toolbox`** (optional, `dev-fdn` group) is GPL-3.0.
  It is not required and is not vendored.
  It is currently a private repository.
  The two files that use it skip without it.

`tools/collect_attribution.py --check` verifies that `NOTICE` covers every library that the generated modules use.
CI runs this check on every PR.

## Citing

[CITATION.cff](CITATION.cff) contains the Faustax citation; GitHub's "Cite this repository" feature renders it.
The file also contains a machine-readable reference list for the work that Faustax is based on.
If your research relies on a particular layer, cite its upstream work together with Faustax:

- **Any Faustax processor** — [Faust](https://faust.grame.fr)
  ([GitHub](https://github.com/grame-cncm/faust)): Orlarey, Letz & Fober,
  *Faust: an Efficient Functional Approach to DSP Programming*, in *New
  Computational Paradigms for Computer Music*, Delatour, 2009, pp. 65–96 —
  the compiler and standard libraries the modules are generated from.
- **The dasp-pytorch parity set** —
  [dasp-pytorch](https://github.com/csteinmetz1/dasp-pytorch) and
  Steinmetz, Bryan & Reiss, *Style Transfer of Audio Effects with
  Differentiable Signal Processing*, JAES 70(9), 2022; for
  `NoiseShapedReverb`, Steinmetz, Ithapu & Calamia, *Filtered Noise Shaping
  for Time Domain Room Impulse Response Estimation from Reverberant
  Speech*, WASPAA 2021; for the compressor design, Giannoulis, Massberg &
  Reiss, *Digital Dynamic Range Compressor Design—A Tutorial and Analysis*,
  JAES 60(6), 2012.
- **`faustax.ops` / `faustax.dynamics` / `faustax.filters`** — Yu et al.,
  *Differentiable All-pole Filters for Time-varying Audio Systems*,
  DAFx 2024 (torchlpc / torchcomp), and
  [philtorch](https://github.com/yoyolicoris/philtorch).
- **`faustax.diffvox`** — Yu et al., *DiffVox: A Differentiable Model for
  Capturing and Analysing Vocal Effects Distributions*, DAFx 2025.
- **Learnable normalized parameters** — Ben Hayes,
  [Magic Clamp](https://github.com/ben-hayes/magic-clamp), 2025.
