The Ai module

Every AI model I use runs in a Docker container of its own, each one started by a service in my Docker module (the models are in AI models in Docker containers). The Ai module is the layer on top: one function for each task, which starts the model's container if it isn't already running, sends it the request, and hands back the result.

Making a piece of music is one call:

const Ai = require("Ai");
const fs = require("fs");

const prompt = "Drum and bass with electric guitar";
const settings = {
    _is: "Ai.Audio.Audiocraft",
    duration: 20.0
}

Ai.Audio.Audiocraft(
    prompt,
    settings,
    (data) => {
        const audioBuffer = Buffer.from(data.split(',')[1], 'base64');
        fs.writeFileSync(`./resources/audiocraft.wav`, audioBuffer);
    },
    (msg) => {
        console.log("Ai server failed with following message:\n")
        console.log(msg)
    }
);

The first time this runs, the MusicGen container is built and started, which takes a while. After that the container is already running with the model in memory, and every script on the machine that asks for MusicGen shares it.

How a call works

Every function has the same shape: what to work on, a settings object, a function for the result and a function for errors. Inside, each one does the same five things. Here's the music function:

Audiocraft: (prompt, settings, onSuccess, onError) => {
    const defaultSettings = {
        top_k: 250,
        top_p: 0.0,
        temperature: 1.0,
        duration: 30.0,
        cfg_coef: 3.0,
    }
    const serviceName = `ai-audiocraft`;
    const serverPort = 8001;
    const serverUri = `ws://localhost:${serverPort}`;

    Security.String(prompt);
    Security.Check(prompt.length > 0, `Prompt is too short, please provide more context`);
    Security.Object("Ai.Audio.Audiocraft")(settings);
    Security.Callback(onSuccess);
    Security.Callback(onError);

    for (let key in defaultSettings) {
        if (!settings.hasOwnProperty(key)) {
            settings[key] = defaultSettings[key]
        }
    }

    const isStarted = Docker.Container.Exists(serviceName, true);
    const connHandler = () => {
        const ws = new WebSocket(serverUri);
        ws.on('open', function () {
            ws.send(JSON.stringify({ started: true, prompt: prompt, settings }));
        });
        ws.on('message', function (data) {
            const {success, error, progress} = JSON.parse(data);
            if (success !== undefined) {
                onSuccess(success[0]);
            } else if (error !== undefined) {
                onError(error);
            } else if (progress !== undefined) {
                console.info(`Progress information received - ${progress}% done`);
            }
        });
        ws.on('error', function (error) {
            setTimeout(connHandler, 5000)
        });
    }
    if (!isStarted) {
        Docker.Service.AI.AUDIOCRAFT_LARGE(serviceName, serverPort);
        setTimeout(connHandler, 10000);
    }
    else {
        connHandler()
    }
}

• Check the arguments. The Security module checks every argument before anything else happens. The settings object carries its type in _is, so passing the settings for one function to another is caught straight away.
• Fill in the defaults. Any setting that wasn't given takes the model's default, so most calls only set the one or two things they care about.
• Start the container if it isn't running. The container's name is fixed for each model, so it's only ever started once.
• Connect, and keep trying. A model takes a while to load once its container starts, so the first connection waits a few seconds, and a failed connection tries again every five until the server answers.
• Send one message and wait for the answer. The server sends progress messages while it works, then either the result or an error, and the matching callback is called.

Results that are files, like audio and images, come back as base64 data URLs, ready to be saved or put straight into a web page.

The functions

The module is organised by the kind of media each function makes or works on: Ai.Audio, Ai.Image, Ai.Text, Ai.Code and Ai.Object, with Ai.OpenAi for the functions that call OpenAI's online API instead of a local model.

Audio

Images

Text

The LLaMA models run with llama.cpp, from quantised weights downloaded from Hugging Face. Ai.Text['30B'] and the other sizes are a shorthand for the same models: Ai.Text['30B'].WizardLM is Ai.Text.LLaMa['30B'].WizardLM.

Code

Each one is a template around a LLaMA model: the description finishes a sentence like "Write a Go program that…", so Ai.Code.Go.Program("serves static files from /public", …) asks for "Write a Go program that serves static files from /public". The settings are passed on to the model, so they can set the temperature or the length of the answer.

Program: (description, settings, onSuccess, onError) => {
    Security.String(description);
    Security.Object("Ai.Code.Go.Program")(settings);
    Security.Callback(onSuccess);
    Security.Callback(onError);

    const prompt = `Write a Go program that ${description}`;
    Ai.Text['30B'].WizardLM(prompt, { ...settings, _is: "Ai.Text.LLaMa" }, onSuccess, onError);
}

3D objects

OpenAI

These don't start any containers. They call OpenAI's API with fetch, using your own API key, which goes in the settings as apiKey (the examples read it from the OPENAI_API_KEY environment variable), and they give back OpenAI's response as it comes. OpenAI has since withdrawn the Codex models, so Ai.OpenAi.Codex no longer has anything to call.

Building on each other

Because every function has the same shape, the bigger ones are made from the smaller ones. Translating speech is a transcript followed by a translation:

ToString: (promptEncoded, fromLanguage, toLanguage, onSuccess, onError) => {
    Security.String(promptEncoded);
    Security.String(fromLanguage);
    Security.String(toLanguage);
    Security.Callback(onSuccess);
    Security.Callback(onError);

    Ai.Audio.Transcript(promptEncoded, (transcript) => {
        const transcriptText = transcript[0].transcription;
        Ai.Text.Translate(transcriptText, fromLanguage, toLanguage, onSuccess, onError);
    }, onError);
}

Each step starts its own model if it needs to, so the first time this runs it brings up two containers, one after the other. Translating text into speech is the same thing the other way round, a translation followed by Ai.Audio.Speech.

The same goes for scripts built on the module. The music in Endless music with MusicGen is made by Ai.Audio.Continuous, with a small script around it choosing the prompts.

What it needs

The module only needs two of my other modules, both of which it uses in nearly every file:

• Docker, to check whether a model's container is running and to start it with the model's service.
• Security, to check every argument and settings object.

Apart from those it uses ws for the WebSocket connections, kept in the module's own node_modules, and Node's own fs. Riffusion's server takes an HTTP request rather than a WebSocket, and that and the OpenAI functions use fetch, which is built into Node from version 18.

The models themselves aren't part of the module, and nothing needs downloading in advance. Each model's Docker service downloads its weights while it builds the image, from wherever the model is published: Hugging Face, the authors' GitHub repositories, or for the LLaMA models, the URL of the quantised weights that the Ai module passes to the service. So the first call to a model builds its image, model and all, and every call after that uses the image that's there.

The examples folder has a script for each model, about 40 in all, from GPT-2 to Stable Diffusion and Demucs.

Things to know

• The interface is asynchronous in the callback sense: each function returns straight away, and the result arrives later through onSuccess or onError. They don't return promises.
• The first call to a model is slow: its image is built if it hasn't been, the container starts, and the model loads. Later calls go straight to the running container.
• A model's server works on one request at a time, so calls to the same model from several scripts are answered one after the other.
• Ai.Object.Point_e and Ai.Object.Shap_e use the same port, so only one of them can be running at a time.
• Most models run on the processor by default. Where a function takes a device setting, "cuda" or "gpu" starts its container with the graphics card.