audiotree integration

Install the audiotree extra. Then faustax.audiotree.FaustFx makes each processor an audiotree random transform. This one wrapper replaces a manual port of each effect.

In this document, a batch is an AudioTree with a waveform of shape (batch, channels, samples). In training, the batch comes from a data pipeline that uses audiotree and grain. For experiments, stack file excerpts along the batch axis:

import librosa
from audiotree import AudioTree

batch = AudioTree.batch(
    [AudioTree.from_file(librosa.example("brahms"), duration=10, sample_rate=44100)] * 4
)  # waveform: (4, 1, 441000)

The transform draws physical parameter values for each batch item from distribution tuples. The audiotree transform options (prob, scope, output_key, split_seed) also apply:

from audiotree import AudioTree
from faustax.audiotree import FaustFx

transform = FaustFx(
    Compressor(sample_rate=44100),
    param_dists={
        "threshold_db": ("uniform", -40.0, -6.0),
        "ratio": ("choice", [2.0, 4.0, 8.0]),
        "attack_ms": ("const", 5.0),
    },
    prob=0.5,
)

tree = AudioTree(
    waveform=jax.random.normal(jax.random.key(0), (2, 1, 2048)) * 0.5,
    sample_rate=44100,
)
out = transform.random_map(tree, jax.random.key(1))
print(out.waveform.shape)

(2, 1, 2048)

The distribution tuples follow the audiotools convention: ("const", value), ("uniform", low, high), or ("choice", [values...]). Parameters that you omit use the Faust slider defaults. The constructor validates the bounds against the slider ranges. Thus an incorrect range causes a failure before a pipeline runs.

Know these two behaviors:

  • The transform always clears the loudness metadata. The effect changes the spectral content. Thus the pre-step of the transform invalidates the cached lufs / lufs_windows values. This occurs also when the probabilistic branch does not run, because the two branches of the prob selection must share one pytree structure.

  • The sample rates must be equal. You construct the processor at a fixed sample rate. If you apply the processor to an AudioTree with a different sample_rate, the transform raises an error.

In a grain data pipeline

Chain the FaustFx transforms onto a grain dataset with ds.seed(n) and ds.random_map(...), the same as the audiotree transforms. The position of the transforms relative to the multiprocessing stage is important. Put the Faustax transforms after .batch() and .mp_prefetch():

ds = create_audio_dataset(sources=sources, sample_rate=44100, duration=2.0, mono=True)
ds = ds.seed(42)

# CPU stage: NumPy-backend audiotree transforms, per item, executed
# inside the grain worker processes (everything before mp_prefetch).
ds = ds.random_map(volume_norm(min_db=-24.0, max_db=-18.0))

ds = ds.to_iter_dataset(
    read_options=grain.ReadOptions(num_threads=0, prefetch_buffer_size=0)
)
ds = ds.batch(batch_size=batch_size, drop_remainder=True, batch_fn=AudioTree.batch)
ds = ds.mp_prefetch(options=grain.MultiprocessingOptions(num_workers=worker_count))

# JAX stage: Faustax transforms go HERE, after mp_prefetch — they run in
# the main process, vmapped over the whole batch.
ds = ds.seed(43)
ds = ds.random_map(compressor_transform)
ds = ds.random_map(reverb_transform)

Put the transforms after mp_prefetch, and not before, for these three reasons. We measured all three reasons:

  1. A transform cannot cross the spawn boundary. The grain workers are spawned processes, and cloudpickle sends data to them. The generated NNX modules contain a jax.custom_vjp (the magic-clamp gradient). cloudpickle cannot reconstruct this jax.custom_vjp. Thus a FaustFx before mp_prefetch fails when num_workers > 0.

  2. The post-batch position is approximately 6x faster. The per-sample scan is sequential in time, but it is parallel across the batch. Thus one vmapped call on a batch of 8 has approximately the same cost as one item (0.89 ms/item post-batch vs 5.6 ms/item per-item, Compressor at 44.1 kHz).

  3. The workers keep their most efficient tasks. The workers do the file I/O, the decoding, and the NumPy-backend audiotree transforms. This work overlaps with the JAX stage.