aigen/server/api/generate.post.ts

import { copyFileSync, readFileSync, writeFileSync } from 'node:fs'
import { join } from 'node:path'

function isComfyBusyTimeout(error: unknown) {
  const err = error as { statusCode?: number; status?: number; data?: { code?: string }; message?: string; statusMessage?: string }
  if (err?.statusCode === 409 || err?.status === 409 || err?.data?.code === 'COMFY_BUSY') return true
  const text = `${err?.statusMessage || ''} ${err?.message || ''}`
  return text.includes('still busy after 3 minutes')
}

function parseDuration(raw: unknown) {
  const seconds = Number(raw)
  if (!Number.isFinite(seconds)) return 5
  return Math.min(120, Math.max(0.5, Math.round(seconds * 10) / 10))
}

function parseExtendDuration(raw: unknown) {
  return parseDuration(raw)
}

function parseExtensions(raw: string | undefined) {
  if (!raw) return [] as { prompt: string; duration: number }[]
  try {
    const parsed = JSON.parse(raw)
    if (!Array.isArray(parsed)) return []
    return parsed
      .slice(0, 8)
      .map((item: { prompt?: unknown; duration?: unknown }) => ({
        prompt: String(item?.prompt || '').trim(),
        duration: parseExtendDuration(item?.duration)
      }))
      .filter(item => item.prompt)
  } catch {
    return []
  }
}

const SAMPLERS = new Set(['res_multistep', 'euler', 'dpmpp_2m'])
const SCHEDULERS = new Set(['simple', 'ddim_uniform', 'sgm_uniform'])

function frameLength(seconds: number, fps: number) {
  return Math.max(5, Math.floor(seconds * fps))
}

function parseFps(raw: string | undefined) {
  const fps = Number(raw)
  return fps === 12 || fps === 30 || fps === 24 ? fps : 24
}

function parseCfg(raw: string | undefined, turbo: boolean) {
  const fallback = turbo ? 1.5 : 4
  const value = Number(raw)
  if (!Number.isFinite(value)) return fallback
  const snapped = Math.round(value * 2) / 2
  return Math.min(10, Math.max(1, snapped))
}

function parseSampler(raw: string | undefined) {
  return SAMPLERS.has(raw || '') ? raw! : 'res_multistep'
}

function parseScheduler(raw: string | undefined) {
  return SCHEDULERS.has(raw || '') ? raw! : 'simple'
}

function sleep(ms: number) {
  return new Promise(resolve => setTimeout(resolve, ms))
}

function assertJobActive(job: ReturnType<typeof createJob>) {
  if (job.status === 'cancelled') {
    throw new Error('Job interrupted.')
  }
  if (job.status === 'error') {
    throw new Error(job.error || 'Generation failed')
  }
}

export default defineEventHandler(async (event) => {
  const form = await readMultipartFormData(event)
  if (!form?.length) {
    throw createError({ statusCode: 400, statusMessage: 'Expected multipart form data' })
  }

  const fields: Record<string, string> = {}
  let image: { filename: string; data: Buffer; type?: string } | null = null
  for (const part of form) {
    if (part.name === 'image' && part.filename && part.data?.length) {
      image = { filename: part.filename, data: part.data, type: part.type }
    } else if (part.name && part.data) {
      fields[part.name] = part.data.toString('utf8')
    }
  }

  if (!image) {
    throw createError({ statusCode: 400, statusMessage: 'An initial image is required' })
  }
  const prompt = (fields.prompt || '').trim()
  if (!prompt) {
    throw createError({ statusCode: 400, statusMessage: 'A motion prompt is required' })
  }
  const extensions = parseExtensions(fields.extensions)

  const { width, height } = resolveOutputSize(fields.aspect, image.data, Number(fields.width), Number(fields.height))
  const turbo = fields.turbo !== 'false'
  const steps = Number(fields.steps || (turbo ? 8 : 20))
  const seed = fields.seed && fields.seed !== 'random'
    ? Number(fields.seed)
    : Math.floor(Math.random() * 2_147_483_647)
  const durationSeconds = parseDuration(fields.duration || '5')
  const fps = parseFps(fields.fps)
  const cfg = parseCfg(fields.cfg, turbo)
  const samplerName = parseSampler(fields.sampler_name)
  const scheduler = parseScheduler(fields.scheduler)
  const length = frameLength(durationSeconds, fps)
  const hideThumbnail = fields.hideThumbnail === 'true'
  const hideInput = fields.hideInput === 'true'
  const sound = fields.sound !== 'false'
  const clipName = (fields.name || '').trim().slice(0, 80)
  const ownerKey = libraryOwnerKey(event)
  const library = publicLibrary(event)
  const folderId = library.folders.some(folder => folder.id === fields.folderId)
    ? fields.folderId
    : library.folders[0]?.id
  if (!folderId) {
    throw createError({ statusCode: 400, statusMessage: 'Create a library folder before generating' })
  }
  const destFolder = library.folders.find(folder => folder.id === folderId)
  const folderLocked = Boolean(destFolder?.protected && !destFolder.unlocked)
  assertFolderExists(event, folderId)
  const still = await saveStill({
    ownerKey,
    folderId,
    filename: image.filename,
    data: image.data,
    width,
    height,
    hideInput
  })

  const chainTotal = 1 + extensions.length
  const job = createJob()
  job.maxStep = steps
  job.hideThumbnail = hideThumbnail
  job.library = {
    ownerKey,
    folderId,
    hideThumbnail,
    hideInput,
    folderLocked,
    name: clipName,
    prompt,
    aspect: fields.aspect || 'auto',
    width,
    height,
    steps,
    turbo,
    seed,
    cfg,
    fps,
    samplerName,
    scheduler,
    thumb: image.data,
    stillId: still.id,
    stillFilename: still.filename,
    duration: durationSeconds,
    sound,
    draftId: (fields.draftId || '').trim() || undefined,
    extensions,
    chainIndex: 0,
    chainStep: 1,
    chainTotal,
    chainLabel: extensions.length ? 'Initial' : undefined
  }
  emitChainJob(job, { type: 'status', message: 'Checking ComfyUI...', progress: 1 })

  void runGeneration(job, {
    prompt,
    image,
    width,
    height,
    steps,
    seed,
    turbo,
    length,
    sound,
    cfg,
    fps,
    samplerName,
    scheduler,
    extensions
  }).catch(async (error) => {
    removeExtendTemp(job.library?.extendTmpDir)
    if (isComfyBusyTimeout(error) && job.library) {
      try {
        const draft = await saveRetryDraft({
          ownerKey: job.library.ownerKey,
          id: job.library.draftId,
          folderId: job.library.folderId,
          stillId: job.library.stillId,
          stillFilename: job.library.stillFilename,
          name: job.library.name,
          prompt: job.library.prompt,
          aspect: job.library.aspect,
          width: job.library.width,
          height: job.library.height,
          steps: job.library.steps,
          turbo: job.library.turbo,
          seed: job.library.seed,
          duration: job.library.duration || durationSeconds,
          hideThumbnail: job.library.hideThumbnail,
          hideInput: job.library.hideInput === true,
          sound: job.library.sound !== false,
          cfg: job.library.cfg,
          fps: job.library.fps,
          samplerName: job.library.samplerName,
          scheduler: job.library.scheduler,
          extensions: job.library.extensions
        })
        job.library.draftId = draft.id
        job.status = 'deferred'
        const message = 'ComfyUI was still busy after 3 minutes. Saved this job — click it under Held jobs to restore the form.'
        emitJob(job, { type: 'deferred', draftId: draft.id, message, progress: job.progress })
        return
      } catch {
        // fall through to a normal error if the held job could not be saved
      }
    }
    if (job.status === 'error' || job.status === 'cancelled' || job.status === 'deferred') return
    const message = error instanceof Error ? error.message : String(error)
    job.status = 'error'
    job.error = message
    emitJob(job, { type: 'error', error: message, message })
  })

  return {
    jobId: job.id,
    clientId: job.clientId,
    seed,
    width,
    height,
    aspect: fields.aspect || 'auto',
    cfg,
    fps,
    sampler_name: samplerName,
    scheduler,
    steps,
    turbo,
    folderId,
    hideThumbnail,
    folderLocked,
    chainTotal
  }
})

type GenerateParams = {
  prompt: string
  image: { filename: string; data: Buffer; type?: string }
  width: number
  height: number
  steps: number
  seed: number
  turbo: boolean
  length: number
  sound: boolean
  cfg: number
  fps: number
  samplerName: string
  scheduler: string
  extensions: { prompt: string; duration: number }[]
}

async function queueMiniMax(
  job: ReturnType<typeof createJob>,
  params: Omit<GenerateParams, 'extensions'> & { persist: boolean }
) {
  assertJobActive(job)
  job.socketReady = false
  job.promptId = undefined
  job.video = undefined
  job.segmentBuffer = undefined

  const done = watchComfyJob(job, { persist: params.persist })
  job.status = 'uploading'
  const uploading = (job.library?.chainIndex || 0) > 0
    ? 'Uploading last frame to ComfyUI...'
    : 'Uploading image to ComfyUI...'
  const queueing = (job.library?.chainIndex || 0) > 0
    ? 'Queueing extension on MiniMax H3...'
    : 'Queueing MiniMax H3 job...'
  emitChainJob(job, { type: 'status', message: uploading, progress: 4 })
  const uploaded = await uploadImage(params.image, job.id)
  if (job.library) {
    job.library.imageName = uploaded.name
    job.library.imageSubfolder = uploaded.subfolder
  }
  emitChainJob(job, { type: 'status', message: queueing, progress: 6 })
  await waitForComfySocket(job, 4000)

  const graph = buildWorkflow({
    prompt: params.prompt,
    imageName: uploaded.name,
    width: params.width,
    height: params.height,
    steps: params.steps,
    seed: params.seed,
    turbo: params.turbo,
    length: params.length,
    cfg: params.cfg,
    fps: params.fps,
    samplerName: params.samplerName,
    scheduler: params.scheduler,
    filenamePrefix: comfyFilenamePrefix(),
    sound: params.sound
  })

  const queued = await queuePrompt(graph, job.clientId)
  job.promptId = queued.prompt_id
  job.status = 'running'
  if (job.library?.draftId) {
    await deleteRetryDraft(job.library.ownerKey, job.library.draftId).catch(() => null)
    job.library.draftId = undefined
  }
  if (job.library && job.promptId) {
    writePendingJob({
      jobId: job.id,
      promptId: job.promptId,
      clientId: job.clientId,
      ownerKey: job.library.ownerKey,
      folderId: job.library.folderId,
      hideThumbnail: job.library.hideThumbnail,
      folderLocked: job.library.folderLocked,
      name: job.library.name,
      prompt: job.library.prompt,
      aspect: job.library.aspect,
      width: job.library.width,
      height: job.library.height,
      steps: job.library.steps,
      turbo: job.library.turbo,
      seed: job.library.seed,
      startedAt: job.startedAt,
      imageName: job.library.imageName,
      imageSubfolder: job.library.imageSubfolder,
      extendTmpDir: job.library.extendTmpDir,
      extendPart1Path: job.library.extendPart1Path
    })
  }
  emitChainJob(job, { type: 'status', message: 'Job queued on ComfyUI', progress: 8 })
  await done
  assertJobActive(job)
  if (!params.persist && !job.segmentBuffer?.length) {
    throw new Error('Segment finished without a video')
  }
}

async function runGeneration(
  job: ReturnType<typeof createJob>,
  params: GenerateParams
) {
  const ready = (status: { state: string; message: string; queueRunning?: number; queuePending?: number }) => {
    emitChainJob(job, {
      type: status.state === 'busy' ? 'busy' : 'status',
      message: status.message,
      progress: status.state === 'online' ? 3 : 1,
      busy: status.state === 'busy',
      queueRunning: status.queueRunning,
      queuePending: status.queuePending
    })
  }

  await ensureComfyReady(ready)

  const { extensions, ...base } = params
  await queueMiniMax(job, {
    ...base,
    persist: extensions.length === 0
  })

  if (!extensions.length || !job.library) return

  const tmpDir = extendTempDir(job.library.ownerKey, job.id)
  job.library.extendTmpDir = tmpDir
  const currentPath = join(tmpDir, 'current.mp4')
  const part1Path = join(tmpDir, 'part1.mp4')
  const framePath = join(tmpDir, 'last_frame.png')
  writeFileSync(currentPath, job.segmentBuffer!)
  job.segmentBuffer = undefined

  for (let i = 0; i < extensions.length; i++) {
    assertJobActive(job)
    const ext = extensions[i]
    const isLast = i === extensions.length - 1
    job.library.chainIndex = i + 1
    job.library.chainStep = i + 2
    job.library.chainLabel = `Extension ${i + 1}`
    job.library.extendPart1Path = undefined

    emitChainJob(job, { type: 'status', message: 'Waiting 3s buffer...', progress: 1 })
    await sleep(3000)
    assertJobActive(job)

    copyFileSync(currentPath, part1Path)
    await extractLastFrame(part1Path, framePath)
    const frame = readFileSync(framePath)
    job.library.extendPart1Path = part1Path
    const sound = await probeHasAudio(currentPath)
    const seed = Math.floor(Math.random() * 2_147_483_647)

    await ensureComfyReady(ready)
    await queueMiniMax(job, {
      prompt: ext.prompt,
      image: { filename: 'last_frame.png', data: frame, type: 'image/png' },
      width: params.width,
      height: params.height,
      steps: params.steps,
      seed,
      turbo: params.turbo,
      length: frameLength(ext.duration, params.fps),
      sound,
      cfg: params.cfg,
      fps: params.fps,
      samplerName: params.samplerName,
      scheduler: params.scheduler,
      persist: isLast
    })

    if (!isLast) {
      if (!job.segmentBuffer?.length) {
        throw new Error('Extension finished without a stitched video')
      }
      writeFileSync(currentPath, job.segmentBuffer)
      job.segmentBuffer = undefined
      job.library.extendPart1Path = undefined
    }
  }
}