Music generation tasks

import json
import os
import time
from urllib.request import Request, urlopen

API = "https://api.snapgen.org/v1"
AUTH = {"Authorization": f"Bearer {os.environ['SNAPGEN_API_KEY']}"}
FINAL = {"succeeded", "failed", "expired", "reconciliation_required"}


def submit(body):
    request = Request(
        f"{API}/music/generations",
        data=json.dumps(body).encode(),
        headers={**AUTH, "Content-Type": "application/json"},
    )
    with urlopen(request) as response:
        return json.load(response)["id"]


def wait(task_id):
    while True:
        with urlopen(Request(f"{API}/tasks/{task_id}", headers=AUTH)) as response:
            task = json.load(response)
        if task["status"] in FINAL:
            if task["status"] != "succeeded":
                raise RuntimeError(f"{task_id} {task['status']}: {task['error']}")
            return task["result"]
        time.sleep(15)


song_id = submit({
    "model": "suno-music",
    "custom": True,
    "title": "Night Drive",
    "style": "synthwave, female vocal",
    "prompt": "[Verse]\nCity lights are calling me home\n[Chorus]\nWe keep on driving through the night",
})
track = wait(song_id)["tracks"][0]

extension_id = submit({
    "model": "suno-extend",
    "source_task_id": song_id,
    "source_index": track["index"],
    "continue_at": 60,  # a second inside the source track
    "custom": True,
    "prompt": "[Bridge]\nHeadlights fading into dawn",
})
wait(extension_id)

full_id = submit({"model": "suno-concat", "source_task_id": extension_id})
print("Full song:", wait(full_id)["urls"][0])

const API = "https://api.snapgen.org/v1";
const auth = { Authorization: `Bearer ${process.env.SNAPGEN_API_KEY}` };
const FINAL = ["succeeded", "failed", "expired", "reconciliation_required"];

async function submit(body) {
  const response = await fetch(`${API}/music/generations`, {
    method: "POST",
    headers: { ...auth, "Content-Type": "application/json" },
    body: JSON.stringify(body),
  });
  if (!response.ok) throw new Error(await response.text());
  return (await response.json()).id;
}

async function wait(taskId) {
  for (;;) {
    const task = await (await fetch(`${API}/tasks/${taskId}`, { headers: auth })).json();
    if (FINAL.includes(task.status)) {
      if (task.status !== "succeeded") {
        throw new Error(`${taskId} ${task.status}: ${JSON.stringify(task.error)}`);
      }
      return task.result;
    }
    await new Promise((resolve) => setTimeout(resolve, 15_000));
  }
}

const songId = await submit({
  model: "suno-music",
  custom: true,
  title: "Night Drive",
  style: "synthwave, female vocal",
  prompt: "[Verse]\nCity lights are calling me home\n[Chorus]\nWe keep on driving through the night",
});
const [track] = (await wait(songId)).tracks;

const extensionId = await submit({
  model: "suno-extend",
  source_task_id: songId,
  source_index: track.index,
  continue_at: 60, // a second inside the source track
  custom: true,
  prompt: "[Bridge]\nHeadlights fading into dawn",
});
await wait(extensionId);

const fullId = await submit({ model: "suno-concat", source_task_id: extensionId });
console.log("Full song:", (await wait(fullId)).urls[0]);