FM Synthesizer Audio-Parameter Shared Embeddings

import numpy as np
from audiotree import AudioTree
from dexed import DexedSynth, Patch
from jax import numpy as jnp
from synapse import SynAPSE

SAMPLE_RATE = 44_100

# "held-out" was evaluated on 8 DX7 algorithms it never saw in training;
# "80-10-10" saw all 32. Both pair a PANNs audio encoder with the DX7-GNN.
model = SynAPSE.from_pretrained("davidbraun/fm-synapse-dx7-gnn", subfolder="held-out")

# Any DX7 cartridge (.syx) works. This one ships with the repo, so clone it or
# point load_bank at a bank of your own.
syx_filepath = "tests/syx/DX7_AllTheWeb_Aminet_2.syx"
presets = [patch.to_preset() for patch in Patch.load_bank(syx_filepath)]

# Render one voice the way the model was trained: 4 s at 44.1 kHz,
# C4 at velocity 85, a 3 s note plus 1 s of release.
target = 0
synth = DexedSynth(sample_rate=SAMPLE_RATE)
synth.load_preset(presets[target])
audio = synth.render(midi_note=60, velocity=85, note_duration=3.0, render_duration=4.0)

query = AudioTree.create(audio, SAMPLE_RATE)
gallery = AudioTree.create(
    None,  # waveform not needed for parameter encoding
    SAMPLE_RATE,
    extras={
        "params": jnp.stack([p.to_array() for p in presets]),
        "algorithm": jnp.asarray([p.algorithm for p in presets], dtype=jnp.int32),
    },
)

# arm1 is the waveform/pre-computed mel encoder.
# arm2 is the parameter encoder.
q_audio = model.arm1(query).prediction     # [1, 384]
q_preset = model.arm2(gallery).prediction  # [N, 384]

# Both arms L2-normalize, so a dot product is already the cosine score.
ranking = np.argsort(-np.asarray(q_audio @ q_preset.T)[0])
print(f"rendered preset {target}; model ranked it #{list(ranking).index(target) + 1} of {len(presets)}")

# the single render paired with its own preset
batch = AudioTree(
    jnp.asarray(audio)[None, None, :],
    sample_rate=SAMPLE_RATE,
    extras={
        "params": jnp.asarray(presets[target].to_array())[None],
        "algorithm": jnp.asarray([presets[target].algorithm], dtype=jnp.int32),
    },
)

out = model(batch)                     # SynAPSEOutput dataclass
z_A = out.audio_output.projection      # [B, 384] audio projection
z_P = out.parameter_output.projection  # [B, 384] preset projection
q_A = out.audio_output.prediction      # [B, 384] audio prediction
q_P = out.parameter_output.prediction  # [B, 384] preset prediction
# cosine similarity between q_A and q_P is the retrieval score used at eval time.

model.save_safetensors("model.safetensors")           # weights only (Param/BatchStat)
model.load_safetensors("model.safetensors")           # in place, into a same-shape model
model.save_orbax("ckpt/"); model.load_orbax("ckpt/")  # Orbax checkpoint dir

model.save_pretrained("out/", argbind_config={...})   # argbind_params.yml + model.safetensors
model = SynAPSE.from_pretrained("out/")               # local dir, or a HF repo id ([hf] extra)

# One repo can hold several models; name the one you want.
model = SynAPSE.from_pretrained("davidbraun/fm-synapse-dx7-gnn", subfolder="80-10-10")
model.train()  # switch to training mode (pretrained models default to eval)