Music generation

import { writeFileSync } from 'node:fs'
import {
  AUDIOGEN_ACESTEP_5HZ_LM_0_6B_Q8_0,
  AUDIOGEN_ACESTEP_V15_TURBO_Q4_K_M,
  AUDIOGEN_QWEN3_EMBEDDING_0_6B_Q8_0,
  AUDIOGEN_VAE_BF16,
  audioGen,
  loadModel,
  unloadModel
} from '@qvac/sdk'

// Re-render an existing song with a new caption (ACE-Step "cover-nofsq"),
// optionally conditioning the timbre on a second reference clip.
//
// Usage:
//   bun examples/audiogen/generate-cover.ts <source.wav|mp3|...> "orchestral arrangement" [reference.wav] [output.wav]
//
// Both audio inputs are file paths: the SDK decodes them (any FFmpeg-decodable
// format) to the 48 kHz stereo float PCM the engine expects. Pass raw
// interleaved stereo 48 kHz Float32 LE PCM as a Buffer instead when the audio
// is already in memory.
const sourcePath = process.argv[2]
const caption = process.argv[3] ?? 'Orchestral arrangement with dramatic strings'
const referencePath = process.argv[4]
const outputPath = process.argv[5] ?? 'audiogen-cover.wav'

if (!sourcePath) {
  console.error(
    'Usage: bun examples/audiogen/generate-cover.ts <source-audio> "<caption>" [reference-audio] [output.wav]'
  )
  process.exit(1)
}

let modelId: string | undefined

try {
  console.log('▸ Loading ACE-Step AudioGen models...')
  modelId = await loadModel({
    modelType: 'audiogen',
    modelConfig: {
      textEncModelSrc: AUDIOGEN_QWEN3_EMBEDDING_0_6B_Q8_0,
      lmModelSrc: AUDIOGEN_ACESTEP_5HZ_LM_0_6B_Q8_0,
      ditModelSrc: AUDIOGEN_ACESTEP_V15_TURBO_Q4_K_M,
      vaeModelSrc: AUDIOGEN_VAE_BF16,
      useGPU: true
    }
  })
  console.log(`▸ Model loaded: ${modelId}`)
  console.log(`▸ Covering ${sourcePath} as: ${caption}`)

  const run = audioGen({
    modelId,
    caption,
    lyrics: '[Instrumental]',
    taskType: 'cover-nofsq',
    sourceAudio: sourcePath,
    ...(referencePath !== undefined && { referenceAudio: referencePath }),
    // cover-nofsq keeps the full source context; coverNoiseStrength blends the
    // initial noise toward the source latent (0 = pure noise, 1 ≈ source).
    audioCoverStrength: 1,
    coverNoiseStrength: 0.75,
    seed: 22886
  })
  console.log(`▸ requestId: ${run.requestId}`)

  for await (const progress of run.progressStream) {
    const value =
      progress.total > 0 ? `${progress.step}/${progress.total}` : `${progress.step} (indeterminate)`
    console.log(`▸ ${progress.stage}: ${value}`)
  }

  const [audio, stats] = await Promise.all([run.audio, run.stats])
  writeFileSync(
    outputPath,
    createWav(audio.pcm, audio.sampleRate, audio.channels, audio.bitsPerSample)
  )
  if (stats) console.log(`▸ Stats: ${JSON.stringify(stats)}`)
  console.log(`▸ Saved ${outputPath}`)

  await unloadModel({ modelId })
  modelId = undefined
  process.exit(0)
} catch (error) {
  if (modelId !== undefined) {
    try {
      await unloadModel({ modelId })
    } catch {
      // Preserve the generation error as the primary failure.
    }
  }
  console.error('✖', error)
  process.exit(1)
}

function createWav(pcm: Uint8Array, sampleRate: number, channels: number, bitsPerSample: number) {
  const header = new ArrayBuffer(44)
  const view = new DataView(header)
  const blockAlign = channels * (bitsPerSample / 8)

  writeAscii(view, 0, 'RIFF')
  view.setUint32(4, 36 + pcm.byteLength, true)
  writeAscii(view, 8, 'WAVE')
  writeAscii(view, 12, 'fmt ')
  view.setUint32(16, 16, true)
  view.setUint16(20, 1, true)
  view.setUint16(22, channels, true)
  view.setUint32(24, sampleRate, true)
  view.setUint32(28, sampleRate * blockAlign, true)
  view.setUint16(32, blockAlign, true)
  view.setUint16(34, bitsPerSample, true)
  writeAscii(view, 36, 'data')
  view.setUint32(40, pcm.byteLength, true)

  const wav = new Uint8Array(44 + pcm.byteLength)
  wav.set(new Uint8Array(header))
  wav.set(pcm, 44)
  return wav
}

function writeAscii(view: DataView, offset: number, value: string) {
  for (let index = 0; index < value.length; index++) {
    view.setUint8(offset + index, value.charCodeAt(index))
  }
}

import { writeFileSync } from 'node:fs'
import { audioGen, loadModel, unloadModel, type ModelProgressUpdate } from '@qvac/sdk'

// Usage:
//   AUDIOGEN_MINIMAX_LM_MODEL=/models/mm3-lm-q8.gguf \
//   AUDIOGEN_MINIMAX_SYNTH_MODEL=/models/mm3-synth-q8.gguf \
//   bun run examples/audiogen/generate-music-minimax.ts "warm cinematic piano" output.wav
const caption = process.argv[2] ?? 'Warm cinematic piano with gentle strings'
const outputPath = process.argv[3] ?? 'minimax-music3-output.wav'

let modelId: string | undefined

try {
  const threads = numberFromEnv('AUDIOGEN_THREADS')
  const maxFrames = numberFromEnv('AUDIOGEN_MAX_FRAMES')
  const inferenceSteps = numberFromEnv('AUDIOGEN_STEPS')

  console.log('▸ Loading MiniMax-Music3 models...')
  modelId = await loadModel({
    modelType: 'audiogen',
    modelConfig: {
      engine: 'minimax',
      lmModelSrc: requiredModelPath('AUDIOGEN_MINIMAX_LM_MODEL'),
      synthModelSrc: requiredModelPath('AUDIOGEN_MINIMAX_SYNTH_MODEL'),
      useGPU: process.env['AUDIOGEN_USE_GPU'] === '1',
      ...(threads !== undefined && { threads })
    },
    onProgress: function (progress: ModelProgressUpdate) {
      logDownloadProgress(progress)
    }
  })

  console.log(`▸ Model loaded: ${modelId}`)
  console.log(`▸ Generating: ${caption}`)
  const run = audioGen({
    modelId,
    caption,
    lyrics: process.env['AUDIOGEN_LYRICS'] ?? '[Instrumental]',
    seed: numberFromEnv('AUDIOGEN_SEED') ?? 7,
    cfgScale: numberFromEnv('AUDIOGEN_CFG_SCALE') ?? 1.7,
    ...(maxFrames !== undefined && { maxFrames }),
    ...(inferenceSteps !== undefined && { inferenceSteps })
  })

  console.log(`▸ requestId: ${run.requestId}`)
  for await (const progress of run.progressStream) {
    const value =
      progress.total > 0 ? `${progress.step}/${progress.total}` : `${progress.step} (indeterminate)`
    console.log(`▸ ${progress.stage}: ${value}`)
  }

  const [audio, stats] = await Promise.all([run.audio, run.stats])
  writeFileSync(
    outputPath,
    createWav(audio.pcm, audio.sampleRate, audio.channels, audio.bitsPerSample)
  )
  if (stats) console.log(`▸ Stats: ${JSON.stringify(stats)}`)
  console.log(`▸ Saved ${outputPath}`)

  await unloadModel({ modelId })
  modelId = undefined
  console.log('▸ Model unloaded')
  process.exit(0)
} catch (error) {
  if (modelId !== undefined) {
    try {
      await unloadModel({ modelId })
    } catch {
      // Preserve the generation error as the primary failure.
    }
  }
  console.error('✖', error)
  process.exit(1)
}

function requiredModelPath(name: string) {
  const value = process.env[name]
  if (!value) throw new Error(`${name} is required`)
  return value
}

function numberFromEnv(name: string) {
  const value = process.env[name]
  if (value === undefined) return undefined
  if (value.trim() === '') throw new Error(`${name} must not be empty`)

  const number = Number(value)
  if (!Number.isFinite(number)) throw new Error(`${name} must be a finite number`)
  return number
}

function logDownloadProgress(progress: ModelProgressUpdate) {
  const downloadedMb = (progress.downloaded / 1e6).toFixed(1)
  const totalMb = (progress.total / 1e6).toFixed(1)
  const line = `▸ Downloading ${progress.percentage.toFixed(0)}% (${downloadedMb}/${totalMb} MB)`
  process.stderr.write(process.stderr.isTTY ? `\r${line}` : `${line}\n`)
  if (process.stderr.isTTY && progress.percentage >= 100) process.stderr.write('\n')
}

function createWav(pcm: Uint8Array, sampleRate: number, channels: number, bitsPerSample: number) {
  const header = new ArrayBuffer(44)
  const view = new DataView(header)
  const blockAlign = channels * (bitsPerSample / 8)

  writeAscii(view, 0, 'RIFF')
  view.setUint32(4, 36 + pcm.byteLength, true)
  writeAscii(view, 8, 'WAVE')
  writeAscii(view, 12, 'fmt ')
  view.setUint32(16, 16, true)
  view.setUint16(20, 1, true)
  view.setUint16(22, channels, true)
  view.setUint32(24, sampleRate, true)
  view.setUint32(28, sampleRate * blockAlign, true)
  view.setUint16(32, blockAlign, true)
  view.setUint16(34, bitsPerSample, true)
  writeAscii(view, 36, 'data')
  view.setUint32(40, pcm.byteLength, true)

  const wav = new Uint8Array(44 + pcm.byteLength)
  wav.set(new Uint8Array(header))
  wav.set(pcm, 44)
  return wav
}

function writeAscii(view: DataView, offset: number, value: string) {
  for (let index = 0; index < value.length; index++) {
    view.setUint8(offset + index, value.charCodeAt(index))
  }
}