Endless music with MusicGen

MusicGen is Meta's text-to-music model: you give it a description like "Glam House with glittery synth sparkles" and it gives you back a clip of music. A clip can be 30 seconds at most, though, and each one is a piece on its own, with its own tempo.

These tracks are made by chaining clips together. Each new chunk is generated as a continuation of the last bar of the one before, so the music carries on from where it was, and every chunk is checked and corrected so the tempo stays the same however long it runs. The prompt changes every few chunks, so the music wanders from one style to the next. Let's see how it works.

Using the library

The generation is done by Ai.Audio.Continuous, a function in a small Node.js library. It takes the first prompt, the settings, a function that's given each new chunk of audio, and a function that's asked for the next prompt every few chunks. Here's the script that made stream 16:

const settings = {
    _is: "Ai.Audio.Continuous",
    //  Tempo in beats per minute
    bpm: 85.0,
    beats_per_chunk: 16,
    // Reduces sampling to the k most likely tokens
    top_k: 250,
    // Reduces sampling to tokens with cumulative probability of p. When set to  `0` (default), top_k sampling is used.
    top_p: 1.0,
    //  Controls the 'conservativeness' of the sampling process. Higher temperature means more diversity.
    temperature: 1.0,
    //  Increases the influence of inputs on the output. Higher values produce lower-varience outputs that adhere more closely to inputs.
    cfg_coef: 3.0,
    buffer: 4,
    device: 'cpu',
    accuracy: 'large',
}
function randomPop() {
    let randomIndex = Math.floor(Math.random() * hiphop_prompts.length);
    return hiphop_prompts.splice(randomIndex, 1)[0];
}
const hiphop_prompts = [
  "A smooth, mellow beat with deep bass and a touch of jazz influence",
  "Hip hop instrumental with high-energy trap influences and tight snares",
  "Chilled lo-fi hip hop beat, featuring soft piano and vinyl scratches",
...
]
const seed_prompt = randomPop();
let chunkNumber = 1;

Ai.Audio.Continuous(
    seed_prompt,
    settings,
    (newChunkEncoded) => {
            console.log("chunk received")
        Security.String(newChunkEncoded)
        //  Save new chunk
        const audioBuffer = Buffer.from(newChunkEncoded.split(',')[1], 'base64');
        fs.writeFileSync(`/root/stream/${chunkNumber++}.wav`, audioBuffer);
    },
    () => {
        console.log("prompt requested")
        //  ...
        //  Perform some logic to get next prompt, this can be a call to AI
        const new_prompt = randomPop()
        return new_prompt ?? ""
    }
);

Each chunk arrives as a WAV file in a data URL, data:audio/wav;base64,…, and the script saves it as the next numbered file. randomPop takes a prompt out of the list at random, so each one is used once, and when the list is empty it returns undefined, the script answers with an empty string, and that ends the run. Ten prompts at four chunks each made stream 16's 40 chunks.

The settings are:

The earlier house tracks, streams 08 to 10, used max_duration: 10 instead of a number of beats: ten seconds of new music from each generation, at 140 bpm.

The server

MusicGen is written in Python, so the model runs in a server of its own. The first time Ai.Audio.Continuous is called, it builds and starts a Docker container: Python 3.9, PyTorch, Meta's AudioCraft (which is MusicGen's library), the BeatNet beat tracker and the Rubber Band time-stretcher, with the small, medium and large models downloaded into the image as it's built. Then it talks to the server over a WebSocket, on port 3003.

The conversation is short. The client sends the first prompt and the settings, and the server sends back chunks as it makes them. Every buffer chunks, the chunk comes as progress rather than chunk, which means "here's a chunk, and what's the next prompt?":

                ws.on('message', function (data) {
                    const {success, chunk, progress} = JSON.parse(data);
                    if (success !== undefined) {
                        console.log("Server finished")
                    } else if (chunk !== undefined) {
                        saveChunk(chunk)
                    } else if (progress !== undefined) {
                        const newPrompt = getPrompt()
                        saveChunk(progress)
                        console.log("sending new chunk - ", {
                            finished: newPrompt==="",
                            prompt: newPrompt
                        })
                        ws.send(JSON.stringify({
                            finished: newPrompt==="",
                            prompt: newPrompt
                        }))
                    } else {
                        console.log(data)
                        throw new Error(`Unrecognized server response`);
                    }
                });

The server waits for the answer before it makes the next chunk, so the script can take as long as it likes to choose, or ask another AI for a prompt, as the comment in the script suggests.

Continuing from the last bar

MusicGen doesn't work on audio directly. A compression model, EnCodec, turns 32kHz audio into tokens, 50 frames of them a second, and MusicGen is a language model over those tokens: given the text and the tokens so far, it predicts the next ones, one frame at a time. You can see this in the server's log as it works through the frames. Stream 16's first chunk is 16 seconds, 800 frames:

Target duration: 16.0s
     1 /    800     2 /    800     3 /    800     4 /    800     5 /    800 ...

Because it's predicting what comes next, MusicGen can carry on from some audio as well as start from nothing. That's generate_continuation: the audio is turned into tokens, and the model generates on from the end of it, under the text prompt. The output starts with the prompt audio, so the server cuts that off and keeps the rest.

For the prompt, the server finds the beats in the previous chunk and takes the end of it, finishing on its last downbeat, the first beat of a bar. In stream 16 that's the last bar, about 2.8 seconds, so each continuation is 941 frames: 141 for the prompt and 800 of new music.

def estimate_beats(wav):
    # resample to BeatNet's sample rate
    beatnet_input = librosa.resample(
        wav,
        orig_sr=sample_rate,
        target_sr=beatnet.sample_rate,
    )
    return beatnet.process(beatnet_input)


#  Gets audio prompt from given loop
def get_audio_prompt(loop, beats_per_prompt):
    beats = estimate_beats(loop)
    bar_times = beats[:, 0][beats[:, 1] == 1]
    num_bars = len(bar_times)
    start_loop_time = beats[0][0]
    end_loop_time = beats[len(beats) - 1][0]
    num_beats = len(beats[(beats[:, 0] >= start_loop_time) & (beats[:, 0] < end_loop_time)])

    print("Prev audio sample contains "+str(num_bars)+" bars")

    if num_beats < beats_per_prompt:
        raise ValueError("Not enough beats for audio prompt")


    end_time = bar_times[len(bar_times)-1]
    last_nbeats = beats[beats[:, 0] <= end_time][-beats_per_prompt:]

    audio_prompt_start_time = last_nbeats[0][0]
    audio_prompt_end_time = last_nbeats[-1][0]
    audio_prompt_start_sample = int(audio_prompt_start_time * sample_rate)
    audio_prompt_end_sample = int(audio_prompt_end_time * sample_rate)

    audio_prompt = torch.tensor(loop[audio_prompt_start_sample:audio_prompt_end_sample])[None]
    audio_prompt_duration = audio_prompt_end_sample - audio_prompt_start_sample
    audio_prompt_seconds = audio_prompt_end_time - audio_prompt_start_time

    return audio_prompt, audio_prompt_duration, audio_prompt_seconds

BeatNet returns one row per beat: its time in seconds, and its position in the bar, where 1 is a downbeat. So bar_times is the time of every downbeat, and the prompt is the last beats_per_prompt beats up to the last one. The prompt runs from the first of those beats to the last, so five beats make one bar of four.

The new music is generated for the length asked for plus the length of the prompt, and the prompt's samples are sliced off the front:

def get_loop(audio_prompt, prompt, model, max_duration, audio_prompt_duration, audio_prompt_seconds, top_k, top_p, temperature, classifier_free_guidance):
    model.set_generation_params(
        duration=max_duration+audio_prompt_seconds,
        top_k=top_k,
        top_p=top_p,
        temperature=temperature,
        cfg_coef=classifier_free_guidance,
    )
    return model.generate_continuation(
        prompt=audio_prompt,
        descriptions=[prompt],
        prompt_sample_rate=sample_rate,
        progress=True,
    ).cpu().numpy()[0, 0, audio_prompt_duration:]

When the chunk length is given in beats, the length asked for is half as much again as the beats need, np.floor((60 / bpm) * beats_per_chunk * 1.5), to leave room to find a downbeat to start on. At 85 bpm, 16 beats is 11.3 seconds, which makes 16 seconds.

Keeping the beat

Each prompt has the tempo added to the end, so stream 16's prompts go to MusicGen as, for example, "Old school hip hop beat with heavy 808s and soulful samples, 85.0 bpm". That helps, but MusicGen only gets roughly the right tempo, and over a long track a small difference in every chunk adds up to music that drifts out of time.

So every chunk is measured and corrected. The server finds the beats in the new music, cuts it from the first downbeat, works out the tempo it actually has, and then either throws it away or stretches it:

#  Normalize new loops
new_loop_normalized = new_loop / np.abs(new_loop).max()

#  Find beats and downbeats
beats = estimate_beats(new_loop_normalized)
downbeat_times = beats[:, 0][beats[:, 1] == 1]

#  Find best cutting points to beats i.e cut from downbeat to downbeat
min_loop_time = downbeat_times[0]
max_loop_time = downbeat_times[len(downbeat_times)-1]

#  Create new beats array always ends on downbeat-1 beat - 1, 2, 3, 4, 1, 2, 3, 4
trim_index_start = np.where(beats[:, 0] == min_loop_time)[0][0]
trim_index_end = np.where(beats[:, 0] == max_loop_time)[0][0]
processed_beats = beats[trim_index_start:trim_index_end]

#  Beat Sanity check
if len(processed_beats) < beats_per_chunk:
    raise ValueError("Failed to produce loop with good beat count, retrying...")

#  Now select commanded number of beats from new array
start_loop_time = processed_beats[0][0]
end_loop_time = processed_beats[beats_per_chunk-1][0]

#  Finally truncate sample
start_sample = int(start_loop_time * sample_rate)
end_sample = int(end_loop_time * sample_rate)
new_loop_truncated = new_loop_normalized[start_sample:end_sample]

#  Find bpm of adjusted beats array, time stretch sample accordingly to fit within given bpm
num_beats = len(processed_beats[(processed_beats[:, 0] >= start_loop_time) & (processed_beats[:, 0] < end_loop_time)])
duration = end_loop_time - start_loop_time
actual_bpm = num_beats / duration * 60
print("BPM count: " + str(actual_bpm)+" vs "+str(bpm))
if abs(actual_bpm - bpm) > 10:
    raise ValueError("Failed to produce loop in good bpm range, retrying...")

loop_final = pyrb.time_stretch(new_loop_truncated, sample_rate, bpm / actual_bpm)

If the tempo is more than 10 bpm out, the chunk is generated again with the same prompt. If it's close enough, Rubber Band stretches it by bpm / actual_bpm, which changes its speed without changing its pitch, so it lands exactly on the tempo. Here's how stream 19 started, at 80 bpm:

Most of the chunks that are kept are within 1 bpm and only need the slightest stretch; the furthest any published track needed was 6.5 bpm. The ones that are thrown away are at about double the tempo: 57 of the 58 in these logs measured between 159 and 173 bpm, and the other one 70.

Throwing chunks away costs time, and it varies a lot between styles. The hip hop prompts at 80 bpm were the hardest:

The end of each chunk is the time of beat number beats_per_chunk, the 16th, so a chunk holds 15 beats rather than 16: every chunk of stream 16 is exactly 10.588 seconds, 15 beats at 85 bpm, and every chunk of stream 19 is exactly 11.25 seconds. Each chunk starts on a downbeat, so where one chunk joins the next, the bar is a beat short. The stateless version of the library, below, cuts from downbeat to downbeat instead.

There's also an unguided setting, which leaves the tempo out of the prompt and skips all of this, for music without a steady beat. The orchestral and cinematic runs used it.

Changing the prompt

MusicGen didn't respond well to changing prompts when you wanted to be precise. If you wanted an instrument to come in at a particular point, a new prompt wouldn't do it: the continuation carries on from what was playing, and the new words only nudge where it goes from there. But it worked well at drifting through its latent space, from one style into the next, and that's what these tracks are. With a new prompt every four chunks, stream 16 moves from gritty underground hip hop to southern hip hop with brass, then a mellow beat with a touch of jazz, then lo-fi piano, and on through all ten prompts in about seven minutes.

To try for more control there's a second version, Ai.Audio.Stateless. It takes a list of prompts, one per chunk, and makes each chunk with a separate request that carries the previous chunk's audio with it, so the server keeps nothing between chunks. A track can be planned prompt by prompt, four chunks of each style say, and with recover_from a run can carry on from a saved WAV file.

Picking the good ones

Letting it run is the easy part. MusicGen can drift to a point where it's making unpleasant sounds rather than music, and it doesn't find its way back, because every chunk continues from the one before. High-pitched noise and white noise compound in the same way: a little hiss in one chunk is in the prompt for the next, and over many generations it builds up until it takes over.

So some discernment was needed. There were 28 runs, and 6 were good enough to publish. The others tried the small, medium and large models, chunks of 8, 16 and 32 beats with prompts from 4 to 32 beats (a 4-beat chunk from a 4-beat prompt wouldn't generate at all), orchestral and cinematic prompts with no tempo, stream 10's prompts on an old and a new version of the server, and two runs of over 20 minutes.

The covers

The covers are from Midjourney. They all come from one prompt, [CHECK: the Midjourney prompt], zoomed out over and over with the same prompt: each zoom shrinks the picture to half its size and paints a new one around it. A cover is a frame picked from along the way. The Experimental AI music page has the tracks as tiles with their covers.