example.py

# Copyright (c) 2026 David Sherlock
#
# Use of this source code is governed by an MIT license that can be found in
# the LICENSE file or at https://opensource.org/licenses/MIT
#
# This applies to this file. The model weights it loads are licensed separately
# by Stability AI — see LICENSE.md and NOTICE alongside the .aimodel.
"""Generate audio with Stable Audio 3 through Apple Core AI.

Self-contained: this file, one `.aimodel` bundle and a `tokenizer/` folder beside it.
No stable-audio-tools, no PyTorch.

    python3 example.py "loud crackling campfire" out.wav --seconds 24
    python3 example.py "a dog barking" out.wav --negative "music" --cfg-scale 3
    python3 example.py "rain" out.wav --init-audio in.wav --regenerate 8:15

Three settings are not optional, and getting them wrong gives audibly broken output
rather than subtly worse output:

  * ping-pong sampling (these are `rf_denoiser` models; Euler overshoots full scale)
  * at least MIN_FRAMES latent frames, i.e. 23.8 s
  * 8 steps at CFG scale 1.0

All three are the defaults here.
"""

from __future__ import annotations

import argparse
import asyncio
import wave
from pathlib import Path
from typing import Any, Iterable

import numpy as np
import coreai.runtime as rt
from tokenizers import Tokenizer

SR = 44100                  # sample rate the models were trained at
LATENT_DIM = 256            # channels in the latent the DiT operates on
DOWNSAMPLE = 4096           # audio samples per latent frame
MAX_TOKENS = 256            # the conditioner's fixed prompt length
STEPS = 8                   # the distilled models' step count

# The model's LogSNRShift, whose `rate` is 0 — so the schedule does not vary with
# sequence length and these two constants fully determine it.
LOGSNR_START = -6.2
LOGSNR_END = 2.0

# distribution_shift_options.min_length. Below this the models emit gross
# high-frequency content: 16-27% of energy above 10 kHz, against ~1% when correct.
MIN_FRAMES = 256


def frames_for(seconds: float) -> int:
    """Latent frames needed for `seconds` of audio, floored at the model's minimum."""
    return max(int(round(seconds * SR / DOWNSAMPLE)), MIN_FRAMES)


def schedule(steps: int = STEPS, sigma_max: float = 1.0) -> np.ndarray:
    """Timestep schedule, evenly spaced in log-SNR with the endpoints pinned.

    Returns `steps + 1` values descending from `sigma_max` to 0.0. `sigma_max` below
    1.0 starts partway down the trajectory, which is how variations work: less noise
    added means less departure from the source audio.
    """
    t = np.linspace(sigma_max, 0.0, steps + 1)
    logsnr = LOGSNR_END - t * (LOGSNR_END - LOGSNR_START)
    shifted = 1.0 / (1.0 + np.exp(logsnr))          # sigmoid(-logsnr)
    shifted[t <= 0] = 0.0
    shifted[t >= 1] = 1.0
    shifted[0] = sigma_max
    return shifted


def f32(array: Any) -> rt.NDArray:
    """Wrap `array` as a float32 Core AI tensor."""
    return rt.NDArray(np.ascontiguousarray(np.asarray(array, dtype=np.float32)))


def i32(array: Any) -> rt.NDArray:
    """Wrap `array` as an int32 Core AI tensor."""
    return rt.NDArray(np.ascontiguousarray(np.asarray(array, dtype=np.int32)))


def read_wav(path: str | Path) -> tuple[np.ndarray, int]:
    """Read a 16-bit WAV as `([1, 2, samples] float32, sample_rate)`, mono upmixed."""
    with wave.open(str(path)) as w:
        count, rate, channels = w.getnframes(), w.getframerate(), w.getnchannels()
        raw = w.readframes(count)
    audio = np.frombuffer(raw, dtype=np.int16).astype(np.float32).reshape(-1, channels).T / 32768
    if channels == 1:
        audio = np.repeat(audio, 2, axis=0)
    return audio[None], rate


def write_wav(path: str | Path, audio: np.ndarray, sr: int = SR) -> None:
    """Peak-normalise `[channels, samples]` audio and write it as a 16-bit WAV."""
    audio = audio / max(float(np.abs(audio).max()), 1e-9) * 0.99
    pcm = (np.clip(audio.T, -1, 1) * 32767).astype("<i2")
    with wave.open(str(path), "wb") as w:
        w.setnchannels(pcm.shape[1])
        w.setsampwidth(2)
        w.setframerate(sr)
        w.writeframes(pcm.tobytes())


def _pad_to(array: np.ndarray, length: int, axis: int) -> np.ndarray:
    """Zero-pad `array` up to `length` along `axis`; return it unchanged if long enough."""
    short = length - array.shape[axis]
    if short <= 0:
        return array
    shape = list(array.shape)
    shape[axis] = short
    return np.concatenate([array, np.zeros(shape, np.float32)], axis=axis)


class Model:
    """A loaded `.aimodel` bundle and the tokenizer beside it.

    The bundle holds several graphs as named functions: `condition`, `denoise`, and
    ladders of `decode_N` / `encode_N` at fixed lengths. `condition` and `denoise`
    take any batch size and the denoiser any length; the codecs are fixed at batch 1
    and one length each, which is why the ladder exists.
    """

    def __init__(self, asset: rt.AIModel, tokenizer: Tokenizer) -> None:
        self._asset = asset
        self._tokenizer = tokenizer
        self.condition_fn = asset.load_function("condition")
        self.denoise_fn = asset.load_function("denoise")

    @classmethod
    async def load(cls, model_dir: str | Path) -> "Model":
        """Load the bundle in `model_dir` (or the `.aimodel` path itself) onto the GPU."""
        path = Path(model_dir)
        asset_path = path if path.suffix == ".aimodel" else next(path.glob("*.aimodel"))
        folder = path.parent if path.suffix == ".aimodel" else path

        options = rt.SpecializationOptions.from_preferred_compute_unit_kind(
            rt.ComputeUnitKind.gpu())
        asset = await rt.AIModel.load(str(asset_path), options)

        tokenizer = Tokenizer.from_file(str(folder / "tokenizer" / "tokenizer.json"))
        tokenizer.enable_truncation(max_length=MAX_TOKENS)
        tokenizer.enable_padding(length=MAX_TOKENS, pad_id=0, pad_token="<pad>")
        return cls(asset, tokenizer)

    def rung(self, kind: str, need: int):
        """Smallest `decode_N`/`encode_N` that fits `need` frames, with its length.

        A fixed-length graph costs its full size on every call, since shorter input
        has to be padded up to it — so taking the smallest that fits is worth doing.
        """
        available = sorted(
            int(name.split("_")[1]) for name in self._asset.function_names
            if name.startswith(f"{kind}_"))
        fits = [length for length in available if length >= need]
        if not fits:
            raise SystemExit(
                f"{need} frames ({need * DOWNSAMPLE / SR:.1f}s) exceeds this model's maximum "
                f"of {available[-1]} ({available[-1] * DOWNSAMPLE / SR:.1f}s)")
        return self._asset.load_function(f"{kind}_{fits[0]}"), fits[0]

    async def condition(self, prompt: str, seconds: float,
                        takes: int = 1) -> tuple[Any, Any, Any]:
        """Encode `prompt` and duration into `(cross_attn_cond, cross_attn_mask, global_cond)`.

        The graphs take a batch dimension, so `takes` above 1 conditions that many
        clips at once — the whole point being that one denoise call then produces
        several independent takes of the same prompt.
        """
        encoded = self._tokenizer.encode(prompt)
        out = await self.condition_fn(inputs={
            "input_ids": i32([encoded.ids] * takes),
            "attention_mask": i32([encoded.attention_mask] * takes),
            "seconds": f32([seconds] * takes),
        })
        return out["cross_attn_cond"], out["cross_attn_mask"], out["global_cond"]

    async def velocity(self, x: np.ndarray, t: float, cond: tuple[Any, Any, Any],
                       local: rt.NDArray) -> np.ndarray:
        """One DiT evaluation: the velocity field at timestep `t`."""
        cross, mask, global_cond = cond
        out = await self.denoise_fn(inputs={
            "x": f32(x), "t": f32([t] * x.shape[0]),
            "cross_attn_cond": cross, "cross_attn_mask": mask,
            "global_cond": global_cond, "local_add_cond": local,
        })
        return out["v"].numpy().astype(np.float32)


def guide(x: np.ndarray, t: float, v_cond: np.ndarray, v_uncond: np.ndarray,
          cfg_scale: float, apg_scale: float) -> np.ndarray:
    """Blend conditional and unconditional velocities into a guided one.

    Guidance is defined on the *denoised* estimate rather than the velocity, so this
    converts across and back. With `apg_scale` at 0 that reduces to vanilla CFG,
    `v_cond + (scale - 1)(v_cond - v_uncond)`, which is verified to match the
    reference exactly. Above 0 it is adaptive projected guidance: the component of
    the difference parallel to the conditional estimate is attenuated, which lets
    higher guidance scales be used without the output over-saturating.
    """
    cond = x - t * v_cond
    uncond = x - t * v_uncond
    diff = cond - uncond

    if apg_scale != 0.0:
        norm = np.sqrt((cond ** 2).sum(axis=(-2, -1), keepdims=True))
        reference = cond / np.maximum(norm, 1e-8)
        parallel = (diff * reference).sum(axis=(-2, -1), keepdims=True) * reference
        diff = apg_scale * (diff - parallel) + (1.0 - apg_scale) * diff

    guided = cond + (cfg_scale - 1.0) * diff
    return (x - guided) / t


def step(x: np.ndarray, v: np.ndarray, t_now: float, t_next: float,
         rng: np.random.Generator, sampler: str) -> np.ndarray:
    """Advance the latent one sampler step.

    `pingpong` is what these distilled models were trained for and the only one that
    reliably sounds right; `euler` is here because the reference offers it, and it
    overshoots full scale on this family.
    """
    if sampler == "euler":
        return x + (t_next - t_now) * v
    denoised = x - t_now * v
    return (1 - t_next) * denoised + t_next * rng.standard_normal(x.shape).astype(np.float32)


def _context_mask(frames: int, keep: float | None, regenerate: str | None) -> np.ndarray:
    """Build the inpaint mask: 1 keeps a frame as context, 0 regenerates it."""
    if regenerate is not None:
        start, end = (float(v) for v in regenerate.split(":"))
        lo, hi = (min(max(int(round(v * SR / DOWNSAMPLE)), 0), frames) for v in (start, end))
        mask = np.ones((1, 1, frames), np.float32)
        mask[:, :, lo:hi] = 0.0                     # a hole in the middle
        return mask
    kept = frames if keep is None else min(int(round(keep * SR / DOWNSAMPLE)), frames)
    mask = np.zeros((1, 1, frames), np.float32)
    mask[:, :, :kept] = 1.0                         # keep a prefix, continue from it
    return mask


async def _encode_latent(model: Model, frames: int, path: str) -> np.ndarray:
    """Encode a WAV to a `[1, LATENT_DIM, frames]` latent, padding or trimming to fit."""
    encoder, _ = model.rung("encode", frames)
    want = encoder.desc.input_descriptor("audio").shape[2]
    source, rate = read_wav(path)
    if rate != SR:
        raise SystemExit(f"{path} is {rate} Hz; resample to {SR} first")
    source = _pad_to(source[:, :, :want], want, axis=2)
    latent = (await encoder(inputs={"audio": f32(source)}))["latent"].numpy().astype(np.float32)
    return _pad_to(latent[:, :, :frames], frames, axis=2)


async def _local_conditioning(model: Model, frames: int, init_audio: str | None,
                              keep: float | None, regenerate: str | None) -> rt.NDArray:
    """`local_add_cond` for the DiT: `[inpaint_mask ; inpaint_masked_input]` on dim 1.

    All zeros means plain text-to-audio. With `init_audio`, the clip is encoded to a
    latent and masked, so kept regions condition the generation.
    """
    if init_audio is None:
        return f32(np.zeros((1, LATENT_DIM + 1, frames)))

    latent = await _encode_latent(model, frames, init_audio)
    mask = _context_mask(frames, keep, regenerate)
    return f32(np.concatenate([mask, latent * mask], axis=1))


async def generate(model_dir: str | Path, prompt: str, seconds: float, *, seed: int = 42,
                   steps: int = STEPS, init_audio: str | None = None,
                   keep: float | None = None, negative: str | None = None,
                   cfg_scale: float = 1.0, regenerate: str | None = None,
                   apg_scale: float = 0.0, sampler: str = "pingpong",
                   variation: float | None = None, takes: int = 1) -> np.ndarray:
    """Generate `[2, samples]` audio at 44.1 kHz.

    Args:
        model_dir: folder holding the `.aimodel` and `tokenizer/`, or the asset itself.
        prompt: what to generate.
        seconds: requested duration; raised to 23.8 s if shorter, since the models
            degrade badly below that.
        seed: fixed seed, so the same inputs give the same audio.
        init_audio: 44.1 kHz WAV to continue from or inpaint into.
        keep: seconds of `init_audio` to hold as context, continuing after it.
        regenerate: `"START:END"` in seconds to replace, keeping everything else.
        negative: prompt to steer away from; needs `cfg_scale` above 1.
        cfg_scale: guidance strength. 1.0 disables it.
        apg_scale: 0 gives vanilla CFG, 1 gives adaptive projected guidance.
        sampler: `pingpong` (correct for these models) or `euler` (for comparison).
        variation: with `init_audio` and no `keep`/`regenerate`, how far to depart
            from the source. 0 keeps it, 1 ignores it. Unset generates from noise.
        takes: how many independent clips to generate in one pass.

    Returns:
        `[takes, 2, samples]` at 44.1 kHz.
    """
    frames = frames_for(seconds)
    model = await Model.load(model_dir)
    decoder, decoder_frames = model.rung("decode", frames)

    duration = frames * DOWNSAMPLE / SR
    cond = await model.condition(prompt, duration, takes)
    uncond = await model.condition(negative, duration, takes) if (
        negative is not None and cfg_scale != 1.0) else None
    # Two distinct ways to use `init_audio`, and they must not be combined:
    #   * variation  - start the trajectory from the source latent, no masking
    #   * keep/regenerate - mask part of the source as fixed context (inpainting)
    # Applying an all-ones mask alongside a variation pins the output to the source
    # and makes the noise level do nothing.
    inpainting = keep is not None or regenerate is not None
    local = await _local_conditioning(
        model, frames, init_audio if inpainting else None, keep, regenerate)
    if takes > 1:
        local = f32(np.repeat(np.asarray(local.numpy()), takes, axis=0))

    rng = np.random.default_rng(seed)
    noise = rng.standard_normal((takes, LATENT_DIM, frames)).astype(np.float32)

    if variation is None:
        x, sigma_max = noise, 1.0
    else:
        # Variation: start partway down the trajectory from the source audio rather
        # than from pure noise. Lower `variation` stays closer to the original.
        if init_audio is None:
            raise SystemExit("--variation needs --init-audio to vary from")
        sigma_max = float(np.clip(variation, 0.0, 1.0))
        source = await _encode_latent(model, frames, init_audio)
        x = np.repeat(source, takes, axis=0) * (1.0 - sigma_max) + noise * sigma_max

    t = schedule(steps, sigma_max)

    for i in range(steps):
        v = await model.velocity(x, t[i], cond, local)
        if uncond is not None:
            v_uncond = await model.velocity(x, t[i], uncond, local)
            v = guide(x, t[i], v, v_uncond, cfg_scale, apg_scale)
        x = step(x, v, t[i], t[i + 1], rng, sampler)

    # The chosen rung is a fixed length: pad the latent up to it, trim the audio back.
    # Decode one take at a time. The codecs are exported at batch 1 — their chunking
    # folds the batch dimension into the sequence dimension, so a free batch there
    # yields a graph that silently returns zeros. Decode runs once per take against
    # the denoiser's eight steps, so looping costs little.
    x = _pad_to(x, decoder_frames, axis=2)
    takes_audio = []
    for i in range(x.shape[0]):
        decoded = (await decoder(inputs={"latent": f32(x[i:i + 1])}))["audio"]
        takes_audio.append(decoded.numpy().astype(np.float32)[0])
    return np.stack(takes_audio)[:, :, :frames * DOWNSAMPLE]


def main(argv: Iterable[str] | None = None) -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("prompt")
    parser.add_argument("out")
    parser.add_argument("--seconds", type=float, default=24.0)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--steps", type=int, default=STEPS)
    parser.add_argument("--model-dir", default=".",
                        help="folder holding the .aimodel and tokenizer/")
    parser.add_argument("--init-audio", help="44.1 kHz WAV to continue or inpaint from")
    parser.add_argument("--keep", type=float,
                        help="seconds of --init-audio to keep as context")
    parser.add_argument("--regenerate", metavar="START:END",
                        help="inpaint: regenerate this span (seconds), keep the rest")
    parser.add_argument("--negative", help="negative prompt (needs --cfg-scale above 1)")
    parser.add_argument("--cfg-scale", type=float, default=1.0,
                        help="guidance strength; 1.0 = off")
    parser.add_argument("--apg-scale", type=float, default=0.0,
                        help="0 = vanilla CFG, 1 = adaptive projected guidance")
    parser.add_argument("--sampler", choices=("pingpong", "euler"), default="pingpong",
                        help="pingpong is correct for these models; euler is for comparison")
    parser.add_argument("--variation", type=float, metavar="LEVEL",
                        help="with --init-audio: 0 keeps the source, 1 ignores it")
    parser.add_argument("--takes", type=int, default=1,
                        help="generate this many independent clips in one pass")
    args = parser.parse_args(list(argv) if argv is not None else None)

    audio = asyncio.run(generate(
        args.model_dir, args.prompt, args.seconds, seed=args.seed, steps=args.steps,
        init_audio=args.init_audio, keep=args.keep, negative=args.negative,
        cfg_scale=args.cfg_scale, regenerate=args.regenerate, apg_scale=args.apg_scale,
        sampler=args.sampler, variation=args.variation, takes=args.takes))

    out = Path(args.out)
    for i, take in enumerate(audio):
        path = out if len(audio) == 1 else out.with_name(f"{out.stem}_{i + 1}{out.suffix}")
        write_wav(path, take)
        print(f"wrote {path}: {take.shape[1] / SR:.2f}s")


if __name__ == "__main__":
    main()