@audiofool/moondog 0.3.0-beta.2 → 0.3.0-beta.4

package/src/integrations/stable-audio/local-generator.mjs

@@ -1,7 +1,8 @@
 // Makes music on this Mac with a local Stable Audio 3 install (Stability AI's optimized MLX
 // build). The listener turns it on once, in Settings or with /generation; until then nothing
 // runs. Each track is a WAV and a receipt in Moondog's state directory, and nothing leaves the
-// machine: weights are never downloaded here, and the receipt keeps no absolute path.
+// machine: a run never downloads anything, and the receipt keeps no absolute path. Moondog's own
+// install (setup.mjs) downloads only when the listener asks, from the pinned official sources.
 import { spawn } from "node:child_process";
 import { randomBytes, randomInt, randomUUID } from "node:crypto";
 import { unwatchFile, watchFile } from "node:fs";
@@ -12,10 +13,12 @@ import path from "node:path";
 import { resolveMoondogStateDirectory } from "../../core/state-directory.mjs";
 import { streamVariation } from "../../profile/taste-compiler.mjs";
 import { resolveMoondogAuthFile } from "../../runtime/pi/persistent-credential-store.mjs";
+import { readLoudness } from "./cover.mjs";
+import { DEFAULT_STABLE_AUDIO_MODEL, isStableAudioModel, SHARED_WEIGHTS, STABLE_AUDIO_MODELS } from "./sources.mjs";
 
-// The weights Stable Audio 3 Medium needs for text to music: the text encoder, the model and its decoder.
-const MEDIUM_WEIGHTS = ["t5gemma_f16.npz", "dit_medium_f16.npz", "same_l_decoder_f32.npz"];
-// Medium was trained on up to 380 seconds; very short clips come out rough.
+// The weights a model needs for text to music: the text encoder, the model and its decoder.
+export const stableAudioWeights = (model = DEFAULT_STABLE_AUDIO_MODEL) => [...SHARED_WEIGHTS, ...STABLE_AUDIO_MODELS[model].weights];
+// Medium and Small Music were trained on up to 380 seconds; very short clips come out rough.
 export const GENERATION_SECONDS = Object.freeze({ min: 15, max: 380, default: 120 });
 // A cold start reads about 4.8 GB of weights before the first step; a 380-second track then takes about 35 seconds.
 const generationTimeoutMs = 10 * 60_000;
@@ -29,6 +32,11 @@ export function defaultStableAudioRoot() {
   return path.join(homedir(), "Developer", "stable-audio-3", "optimized", "mlx");
 }
 
+/** The install Moondog downloads and manages itself; the MLX build is in its mlx folder. */
+export function resolveStableAudioManagedRoot(environment = process.env) {
+  return path.join(resolveMoondogStateDirectory(environment), "stable-audio-3");
+}
+
 // Kept next to the playback preference, apart from runtime settings.
 export function resolveGenerationSettingsFile(environment = process.env) {
   return path.join(path.dirname(resolveMoondogAuthFile(environment)), "generation.json");
@@ -47,41 +55,92 @@ function literalFolder(value) {
   return path.isAbsolute(expanded) ? path.resolve(expanded) : null;
 }
 
+export async function readSavedSettings(environment) {
+  try {
+    const document = JSON.parse(await readFile(resolveGenerationSettingsFile(environment), "utf8"));
+    return document?.version === 1 ? document : {};
+  } catch {
+    return {}; // Not set up yet.
+  }
+}
+
 /**
- * The listener's choice: { enabled, path, source }. Generation is off until they turn it on;
- * MOONDOG_SA3_DIR can point at another install but never turns it on.
+ * The listener's choice: { enabled, path, source, model, folder }. Generation is off until they
+ * turn it on; MOONDOG_SA3_DIR can point at another install but never turns it on. source is
+ * "managed" for Moondog's own install, "settings" or "default" for the listener's own folder
+ * (folder, which Moondog only reads), or "environment". model picks the weights in either.
  * A missing or unreadable file reads as off, so an older or newer file never stops the room.
  */
-export async function readGenerationSettings(environment = process.env) {
-  let document = null;
-  try {
-    document = JSON.parse(await readFile(resolveGenerationSettingsFile(environment), "utf8"));
-  } catch { /* Not set up yet. */ }
-  const saved = document?.version === 1 ? document : {};
+export async function readGenerationSettings(environment = process.env, { platform = process.platform, arch = process.arch } = {}) {
+  const saved = await readSavedSettings(environment);
   const fromEnvironment = literalFolder(environment.MOONDOG_SA3_DIR);
   const fromFile = literalFolder(saved.path);
+  const managed = !fromEnvironment && saved.source === "managed";
+  const folder = fromFile ?? defaultStableAudioRoot();
+  let model = isStableAudioModel(saved.model) ? saved.model : DEFAULT_STABLE_AUDIO_MODEL;
+  // A folder Moondog doesn't manage (MOONDOG_SA3_DIR, a typed folder, ~/Developer) runs the saved model
+  // only if it has it, otherwise the first model it can run; a model chosen for Moondog's own install
+  // waits until that install is chosen again. A folder that runs nothing keeps the saved model.
+  if (!managed) {
+    for (const name of [...new Set([model, DEFAULT_STABLE_AUDIO_MODEL, ...Object.keys(STABLE_AUDIO_MODELS)])]) {
+      if ((await checkStableAudioInstall(fromEnvironment ?? folder, { model: name, platform, arch })).state === "ready") { model = name; break; }
+    }
+  }
   return {
     enabled: saved.enabled === true,
-    path: fromEnvironment ?? fromFile ?? defaultStableAudioRoot(),
-    source: fromEnvironment ? "environment" : fromFile ? "settings" : "default",
+    path: fromEnvironment ?? (managed ? path.join(resolveStableAudioManagedRoot(environment), "mlx") : folder),
+    source: fromEnvironment ? "environment" : managed ? "managed" : fromFile ? "settings" : "default",
+    model,
+    folder,
   };
 }
 
-export async function writeGenerationSettings({ enabled, path: folder } = {}, environment = process.env) {
-  const current = await readGenerationSettings(environment);
+/**
+ * Saves what the listener changed: enabled; path, their own folder, which also switches to it;
+ * source "managed" or "settings" to switch between Moondog's install and their folder; model.
+ * A model chosen without a folder or a source switches to Moondog's install, which holds it.
+ * Older Moondogs read the same version 1 file and see only enabled and path.
+ */
+export async function writeGenerationSettings({ enabled, path: folder, source, model } = {}, environment = process.env, { platform = process.platform, arch = process.arch } = {}) {
+  // The file's own view: a folder MOONDOG_SA3_DIR pins in this process is never written into it.
+  const current = await readGenerationSettings({ ...environment, MOONDOG_SA3_DIR: "" }, { platform, arch });
+  const saved = await readSavedSettings(environment);
   const typed = folder === undefined ? undefined : literalFolder(folder);
   if (folder !== undefined && folder !== null && !typed) {
     throw Object.assign(new Error("Give the Stable Audio 3 folder as a full path, or one starting with ~/."), { code: "generation_path_invalid" });
   }
+  if (model !== undefined && !isStableAudioModel(model)) {
+    throw Object.assign(new Error("Choose medium or small-music."), { code: "generation_model_invalid" });
+  }
+  const savedFolder = typed ?? literalFolder(saved.path);
+  const managed = typed ? false
+    : source === "managed" ? true
+      : source === "settings" ? false
+        : model !== undefined ? true
+          : saved.source === "managed";
+  // Switching to the listener's own folder without a model keeps the model only if that folder has it;
+  // otherwise the first model it can run, so a working install is never checked for weights it lacks.
+  let chosen = model ?? current.model;
+  if (model === undefined && !managed && (typed || source === "settings")) {
+    const folderRoot = savedFolder ?? current.folder;
+    const candidates = [current.model, DEFAULT_STABLE_AUDIO_MODEL, ...Object.keys(STABLE_AUDIO_MODELS)];
+    // A folder that runs nothing yet gets Medium, the model the install steps Moondog shows download.
+    chosen = DEFAULT_STABLE_AUDIO_MODEL;
+    for (const name of [...new Set(candidates)]) {
+      if ((await checkStableAudioInstall(folderRoot, { model: name, platform, arch })).state === "ready") { chosen = name; break; }
+    }
+  }
   const document = {
     version: 1,
     enabled: enabled === undefined ? current.enabled : enabled === true,
-    ...(typed ? { path: typed } : current.source === "settings" ? { path: current.path } : {}),
+    ...(savedFolder ? { path: savedFolder } : {}),
+    ...(managed ? { source: "managed" } : {}),
+    model: chosen,
   };
   const file = resolveGenerationSettingsFile(environment);
   await mkdir(path.dirname(file), { recursive: true, mode: 0o700 });
   await writeFile(file, `${JSON.stringify(document, null, 2)}\n`, { encoding: "utf8", mode: 0o600 });
-  return readGenerationSettings(environment);
+  return readGenerationSettings(environment, { platform, arch });
 }
 
 const isFile = async (file, { executable = false, minimumBytes = 0 } = {}) => {
@@ -97,16 +156,16 @@ const isFile = async (file, { executable = false, minimumBytes = 0 } = {}) => {
 };
 
 /**
- * Whether an install can make music now, without running it: "ready", "unsupported" (MLX needs
- * Apple silicon), "missing" (no install at the folder) or "incomplete" (no Python environment
- * or no Medium weights, which Stability's ./install.sh fetches).
+ * Whether an install can make music now with a model, without running it: "ready", "unsupported"
+ * (MLX needs Apple silicon), "missing" (no install at the folder) or "incomplete" (no Python
+ * environment or not that model's weights, which Moondog's setup or Stability's ./install.sh fetches).
  */
-export async function checkStableAudioInstall(root, { platform = process.platform, arch = process.arch } = {}) {
+export async function checkStableAudioInstall(root, { model = DEFAULT_STABLE_AUDIO_MODEL, platform = process.platform, arch = process.arch } = {}) {
   if (platform !== "darwin" || arch !== "arm64") return { state: "unsupported" };
   if (!root || !path.isAbsolute(root)) return { state: "missing" };
   if (!(await isFile(path.join(root, "scripts", "sa3_mlx.py")))) return { state: "missing" };
   if (!(await isFile(path.join(root, ".venv", "bin", "python"), { executable: true }))) return { state: "incomplete", missing: "environment" };
-  for (const weight of MEDIUM_WEIGHTS) {
+  for (const weight of stableAudioWeights(isStableAudioModel(model) ? model : DEFAULT_STABLE_AUDIO_MODEL)) {
     if (!(await isFile(path.join(root, "models", "mlx", weight), { minimumBytes: 1024 }))) return { state: "incomplete", missing: "weights" };
   }
   return { state: "ready" };
@@ -171,7 +230,7 @@ export function createLocalGenerator({
   spawnImpl = spawn,
   now = () => new Date(),
 } = {}) {
-  let current = { state: "unknown", enabled: false, install: "unknown", path: null, source: null };
+  let current = { state: "unknown", enabled: false, install: "unknown", path: null, source: null, model: DEFAULT_STABLE_AUDIO_MODEL, managed: false };
   let running = null;
   // The tracks made here, newest last, with what happened each time they played: the loop back
   // from listening reads them. Each track's outcomes.jsonl keeps them; the receipt never changes.
@@ -196,8 +255,8 @@ export function createLocalGenerator({
   const check = () => {
     checking ??= (async () => {
       try {
-        const settings = await readGenerationSettings(environment);
-        const install = await checkStableAudioInstall(settings.path, { platform, arch });
+        const settings = await readGenerationSettings(environment, { platform, arch });
+        const install = await checkStableAudioInstall(settings.path, { model: settings.model, platform, arch });
         return publish({
           state: !settings.enabled ? "off" : install.state,
           enabled: settings.enabled,
@@ -205,6 +264,8 @@ export function createLocalGenerator({
           ...(install.missing ? { missing: install.missing } : {}),
           path: settings.path,
           source: settings.source,
+          model: settings.model,
+          managed: settings.source === "managed",
         });
       } finally {
         checking = null;
@@ -233,13 +294,13 @@ export function createLocalGenerator({
   };
 
   const stopped = () => Object.assign(new Error("Making the track was stopped."), { name: "AbortError", code: "generation_stopped" });
-  const run = ({ root, prompt, seconds, seed, out }, { signal, onProgress }) => new Promise((resolve, reject) => {
+  const run = ({ root, model, prompt, seconds, seed, out }, { signal, onProgress }) => new Promise((resolve, reject) => {
     // A stop that came while the install was checked starts nothing.
     if (signal?.aborted) return reject(stopped());
     // The Python script itself, not Stability's ./sa3 wrapper: the wrapper wants uv on PATH and a
     // terminal, and its exec would leave a shell between this process and the one to stop.
     const child = spawnImpl(path.join(root, ".venv", "bin", "python"), [
-      "scripts/sa3_mlx.py", "--prompt", prompt, "--dit", "medium", "--decoder", "same-l",
+      "scripts/sa3_mlx.py", "--prompt", prompt, "--dit", STABLE_AUDIO_MODELS[model].dit, "--decoder", STABLE_AUDIO_MODELS[model].decoder,
       "--seconds", String(seconds), "--seed", String(seed), "--steps", "8", "--out", out,
     ], {
       cwd: root,
@@ -414,7 +475,7 @@ export function createLocalGenerator({
     busy: () => running !== null,
     check,
     async save(settings) {
-      await writeGenerationSettings(settings, environment);
+      await writeGenerationSettings(settings, environment, { platform, arch });
       return check();
     },
     subscribe(listener) {
@@ -469,10 +530,13 @@ export function createLocalGenerator({
         await mkdir(temporary, { mode: 0o700 });
         const wav = path.join(temporary, "track.wav");
         onProgress?.({ stage: "starting" });
-        const measured = await run({ root: status.path, prompt, seconds, seed, out: wav }, { signal, onProgress });
+        const model = isStableAudioModel(status.model) ? status.model : DEFAULT_STABLE_AUDIO_MODEL;
+        const measured = await run({ root: status.path, model, prompt, seconds, seed, out: wav }, { signal, onProgress });
         const audio = await wavInfo(wav).catch(() => null);
         if (!audio || audio.duration_seconds <= 0) throw generationError("generation_failed", "Stable Audio 3 finished but wrote no playable track.");
         await chmod(wav, 0o600);
+        // How loud it runs from start to end, which its cover draws (cover.mjs); kept after a stream drops the audio.
+        const loudness = await readLoudness(wav).catch(() => null);
         const elapsed = Date.now() - started;
         const document = {
           receipt_version: "moondog-generated-track/1",
@@ -482,11 +546,11 @@ export function createLocalGenerator({
           title,
           ...receipt,
           generator: {
-            adapter: "stable-audio-3", adapter_version: 1, model: "sa3-medium-ARC", decoder: "same-l",
+            adapter: "stable-audio-3", adapter_version: 1, model: STABLE_AUDIO_MODELS[model].receipt, decoder: STABLE_AUDIO_MODELS[model].decoder,
             prompt, seconds, seed, steps: 8, cfg: 1, elapsed_ms: elapsed, ...(measured ? { measured } : {}),
             license: "Stability AI Community License; outputs belong to the listener",
           },
-          audio: { file: "track.wav", ...audio },
+          audio: { file: "track.wav", ...audio, ...(loudness ? { loudness } : {}) },
         };
         await writeFile(path.join(temporary, "receipt.json"), `${JSON.stringify(document, null, 2)}\n`, { encoding: "utf8", mode: 0o600, flag: "wx" });
         await rename(temporary, path.join(folder, id));
@@ -509,6 +573,8 @@ export function createLocalGenerator({
 // background. Stopping the player, other music, or the time limit ends it. One stream per room, on
 // the room's one generator; only its newest tracks keep their audio, and every receipt stays.
 const STREAM_SKIP_MS = 30_000;
+// How long a later track waits for its own name before it takes a numbered one.
+const NAMING_TIMEOUT_MS = 10_000;
 export function createAtmosphereStream({ generator, player, now = Date.now, keepAudio = 3 }) {
   let stream = null;
   // What the surfaces see (view()) and hear (subscribe): the phase, and the end with its reason.
@@ -562,9 +628,30 @@ export function createAtmosphereStream({ generator, player, now = Date.now, keep
   };
   const trackFor = (ending, position) => {
     const variation = streamVariation(ending.compiled, position, ending.leaveOut);
-    return { variation, prompt: stableAudioPrompt(variation), title: `${ending.title} · ${position + 1}`,
+    // The first track carries the stream's own name; a later one is named as it is made (nameFor), or numbered.
+    return { variation, prompt: stableAudioPrompt(variation), title: position === 0 ? ending.title : `${ending.title} · ${position + 1}`,
       receipt: { ...ending.receipt, shaped_by: variation.descriptors, ...(variation.bpm ? { bpm: variation.bpm } : {}),
-        stream: { started_at: new Date(ending.startedAt).toISOString(), position: position + 1, featured: variation.featured } } };
+        stream: { name: ending.title, started_at: new Date(ending.startedAt).toISOString(), position: position + 1, featured: variation.featured } } };
+  };
+  // A name of its own for a later track, asked of the stream's namer within a few seconds. One that
+  // fails, is slow, or repeats a name this stream gave keeps the numbered name.
+  const nameFor = async (ending, variation, numbered) => {
+    if (typeof ending.nameTrack !== "function") return numbered;
+    const timeout = AbortSignal.timeout(NAMING_TIMEOUT_MS);
+    try {
+      const name = await Promise.race([
+        ending.nameTrack({
+          stream: ending.title, featured: variation.featured,
+          sound: variation.descriptors.map((descriptor) => descriptor.term).filter((term) => !variation.featured.includes(term)),
+          ...(variation.bpm ? { bpm: variation.bpm } : {}), earlier: [...ending.names],
+        }, { signal: AbortSignal.any([timeout, ending.controller.signal]) }),
+        new Promise((resolve) => { timeout.addEventListener("abort", () => resolve(null), { once: true }); }),
+      ]);
+      const text = typeof name === "string" ? name.trim() : "";
+      return text && !ending.names.some((earlier) => earlier.toLocaleLowerCase() === text.toLocaleLowerCase()) ? text : numbered;
+    } catch {
+      return numbered;
+    }
   };
   // The tracks that started playing keep their audio, the newest few; the rest keep their receipts.
   const prune = (current) => {
@@ -608,7 +695,10 @@ export function createAtmosphereStream({ generator, player, now = Date.now, keep
     current.making = true;
     changed();
     try {
-      const { prompt, title, receipt, variation } = trackFor(current, current.made.length);
+      const { prompt, title: numbered, receipt, variation } = trackFor(current, current.made.length);
+      const title = await nameFor(current, variation, numbered);
+      if (stream !== current) return;
+      current.names.push(title);
       const made = await generator.generate({ prompt, seconds: current.seconds, title, receipt }, { signal: current.controller.signal });
       if (stream !== current) return;
       const file = await generator.trackFile(made.id);
@@ -697,10 +787,10 @@ export function createAtmosphereStream({ generator, player, now = Date.now, keep
      * play. Returns the first track. A stream already running gives way to the new one, and its track
      * plays on until the first new one takes over, so nothing goes silent while that is made.
      */
-    async start({ compiled, title, seconds, minutes, receipt = {} }, { signal, onProgress, onTrack, beforePlay } = {}) {
+    async start({ compiled, title, seconds, minutes, receipt = {} }, { signal, onProgress, onTrack, beforePlay, nameTrack } = {}) {
       end("replaced");
       const current = { compiled, title, seconds, receipt, onTrack, startedAt: now(), endsAt: now() + minutes * 60_000,
-        made: [], leaveOut: new Set(), current: null, next: null, waiting: false, making: true, controller: new AbortController() };
+        made: [], names: [title], nameTrack, leaveOut: new Set(), current: null, next: null, waiting: false, making: true, controller: new AbortController() };
       stream = current;
       changed();
       const stopped = () => Object.assign(new Error("The stream was stopped before it began."), { code: "generation_stopped" });