aigen/server/api/generate.post.ts

import { existsSync, readFileSync } from 'node:fs'
import { addStudioJob, kickStudioQueue, listStudioJobs, videoJobsBusy } from '~/server/utils/studioQueue'
import { listStudioLoras, parsePostedLoraStack, parseShotLoraStacks, persistLoraFields } from '~/server/utils/loras'
import { clampVideoCfg } from '~/utils/generationPresets'
import { defaultVideoSteps, isLtxWorkflow, isTextToVideo, LTX_DISABLED_MESSAGE, ltxWorkflowEnabled, parseVideoWorkflow, videoEngineOf } from '~/utils/videoModels'
import { allowIdentityRefs, normalizePermanenceRefs, resolveGlobalLocks, type PermanenceRef } from '~/utils/globalLocks'
import { composePromptParts, ingestPromptParts, persistPromptWrappers } from '~/utils/promptParts'
import type { QueuedExtension } from '~/server/utils/library'
import {
  clampRefinementDenoise,
  parseRefineExtensionFrame,
  parseSaveLosslessAnchor,
  REFINEMENT_DENOISE_DEFAULT
} from '~/utils/extensionFrame'

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 QueuedExtension[]
  try {
    const parsed = JSON.parse(raw)
    if (!Array.isArray(parsed)) return []
    return parsed
      .map((item: { prompt?: unknown; promptPre?: unknown; promptPost?: unknown; duration?: unknown; permanenceRefs?: unknown; loraName?: unknown; loraStack?: unknown }) => {
        const loraStack = parsePostedLoraStack(item?.loraStack ?? item?.loraName, 'video')
        const wrappers = persistPromptWrappers({
          ...(item && Object.prototype.hasOwnProperty.call(item, 'promptPre') ? { promptPre: String(item.promptPre ?? '') } : {}),
          ...(item && Object.prototype.hasOwnProperty.call(item, 'promptPost') ? { promptPost: String(item.promptPost ?? '') } : {})
        })
        return {
          prompt: String(item?.prompt || '').trim(),
          duration: parseExtendDuration(item?.duration),
          permanenceRefs: normalizePermanenceRefs(item?.permanenceRefs),
          ...persistLoraFields(loraStack),
          ...wrappers
        }
      })
      .filter(item => composePromptParts(item.promptPre || '', item.prompt, item.promptPost || ''))
  } catch (error) {
    if (error && typeof error === 'object' && 'statusCode' in error) throw error
    return []
  }
}

function parseFamilyPermanenceRefs(raw: string | undefined) {
  if (!raw) return [] as PermanenceRef[]
  try {
    return normalizePermanenceRefs(JSON.parse(raw))
  } catch {
    return []
  }
}

function parseShotPermanenceRefs(raw: string | undefined, shotCount: number) {
  if (!raw) return Array.from({ length: shotCount }, () => [] as PermanenceRef[])
  try {
    const parsed = JSON.parse(raw)
    if (!Array.isArray(parsed)) return Array.from({ length: shotCount }, () => [] as PermanenceRef[])
    return Array.from({ length: shotCount }, (_, index) => normalizePermanenceRefs(parsed[index]))
  } catch {
    return Array.from({ length: shotCount }, () => [] as PermanenceRef[])
  }
}

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

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, ltx = false) {
  const fallback = ltx ? 3.5 : (turbo ? 1.5 : 4)
  return clampVideoCfg(raw, fallback)
}

function parseSampler(raw: string | undefined, ltx = false) {
  return SAMPLERS.has(raw || '') ? raw! : (ltx ? 'euler' : 'res_multistep')
}

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

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
  const referenceImages: Array<{ filename: string; data: Buffer; type?: string } | null> = [null, null, null, null]
  let compactRefSlot = 0
  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 && /^refImage(\d)$/.test(part.name) && part.filename && part.data?.length) {
      const slot = Number(part.name.replace('refImage', ''))
      if (slot >= 0 && slot < 4) {
        referenceImages[slot] = { filename: part.filename, data: part.data, type: part.type }
      }
    } else if (part.name === 'refImage' && part.filename && part.data?.length) {
      if (compactRefSlot < 4) {
        referenceImages[compactRefSlot] = { filename: part.filename, data: part.data, type: part.type }
        compactRefSlot += 1
      }
    } else if (part.name && part.data) {
      fields[part.name] = part.data.toString('utf8')
    }
  }

  await listStudioLoras()
  const workflow = parseVideoWorkflow(fields.workflow)
  if (isLtxWorkflow(workflow) && !ltxWorkflowEnabled()) {
    throw createError({ statusCode: 400, statusMessage: LTX_DISABLED_MESSAGE })
  }
  const textToVideo = isTextToVideo(workflow)
  const ownerKey = libraryOwnerKey(event)
  const postedStillId = (fields.stillId || '').trim()
  let libraryStill: { id: string; filename: string; name?: string } | null = null
  if (postedStillId) {
    const path = stillPath(ownerKey, postedStillId)
    if (existsSync(path)) {
      try {
        const existing = getStill(ownerKey, postedStillId)
        libraryStill = { id: existing.id, filename: existing.filename, name: existing.name }
        if (!image) {
          image = {
            filename: existing.name ? `${existing.name}.png` : (existing.filename || 'still.png'),
            data: readFileSync(path),
            type: 'image/png'
          }
        }
      } catch {
        libraryStill = null
      }
    }
  }
  if (!image && !textToVideo) {
    throw createError({ statusCode: 400, statusMessage: 'An initial image is required' })
  }
  const posted = ingestPromptParts(fields)
  const promptPre = posted.promptPre
  const promptPost = posted.promptPost
  let promptMid = posted.promptMid
  if (!composePromptParts(promptPre, promptMid, promptPost)) {
    throw createError({ statusCode: 400, statusMessage: 'A motion prompt is required' })
  }
  const resolvedLocks = resolveGlobalLocks(fields.globalLocks || '', promptMid)
  promptMid = resolvedLocks.prompt
  const prompt = (promptPre || promptPost)
    ? composePromptParts(promptPre, promptMid, promptPost)
    : (posted.prompt || promptMid)
  const globalLocks = resolvedLocks.locks
  const extensions = parseExtensions(fields.extensions)
  const permanenceRefs = parseFamilyPermanenceRefs(fields.permanenceRefs)
  const shotPermanenceRefs = parseShotPermanenceRefs(fields.shotPermanenceRefs, 1 + extensions.length)
  const defaultStack = parsePostedLoraStack(fields.loraStack || fields.lora, 'video')
  const shotLoraStacks = parseShotLoraStacks(fields.shotLoraStacks || fields.shotLoras, 1 + extensions.length)
  if (defaultStack.length && !shotLoraStacks[0].length) shotLoraStacks[0] = defaultStack
  for (const [index, item] of extensions.entries()) {
    if (!shotPermanenceRefs[index + 1]?.length && item.permanenceRefs?.length) {
      shotPermanenceRefs[index + 1] = item.permanenceRefs
    }
    const shotStack = shotLoraStacks[index + 1]
    if (!shotStack.length && item.loraStack?.length) shotLoraStacks[index + 1] = item.loraStack
    else if (shotStack.length) Object.assign(item, persistLoraFields(shotStack))
  }
  const requestedIdentity = workflow === 'v2' && !textToVideo && !isLtxWorkflow(workflow) && fields.useIdentityRefs === 'true'
  const useIdentityRefs = allowIdentityRefs(requestedIdentity, permanenceRefs, globalLocks, shotPermanenceRefs)
  const { width, height } = resolveOutputSize(
    fields.aspect,
    image?.data || Buffer.alloc(0),
    Number(fields.width),
    Number(fields.height)
  )
  const turbo = fields.turbo !== 'false'
  const steps = Number(fields.steps || defaultVideoSteps(turbo, workflow))
  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, isLtxWorkflow(workflow))
  const samplerName = parseSampler(fields.sampler_name, isLtxWorkflow(workflow))
  const scheduler = parseScheduler(fields.scheduler)
  const hideThumbnail = fields.hideThumbnail === 'true'
  const hideInput = fields.hideInput === 'true'
  const sound = !isLtxWorkflow(workflow) && fields.sound !== 'false'
  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 = libraryStill || (image
    ? await rememberInputStill({
      ownerKey,
      folderId,
      filename: image.filename,
      data: image.data,
      width,
      height,
      hideInput
    })
    : null)
  if (!textToVideo && !still?.id) {
    throw createError({ statusCode: 400, statusMessage: 'An initial image is required' })
  }
  const clipName = allocateOutputName(ownerKey, folderId, (fields.name || '').trim(), {
    inputFilename: libraryStill?.name || image?.filename || still?.filename || '',
    engine: videoEngineOf(workflow),
    loraStack: defaultStack
  })
  console.info('[aigen generate]', {
    hasImage: Boolean(image?.data?.length),
    imageBytes: image?.data?.length || 0,
    imageName: image?.filename || '',
    stillId: still?.id || '',
    postedStillId,
    workflow,
    promptChars: promptMid.length,
    textToVideo
  })
  const referenceStillIds: Array<string | null> = [null, null, null, null]
  if (useIdentityRefs) {
    for (const [index, ref] of referenceImages.entries()) {
      if (!ref?.data?.length || isPipelineFrameFilename(ref.filename)) continue
      const saved = await rememberInputStill({
        ownerKey,
        folderId,
        filename: ref.filename || `identity-ref-${index + 2}.png`,
        data: ref.data,
        width,
        height,
        hideInput
      })
      if (saved?.id) referenceStillIds[index] = saved.id
    }
  }

  const chainTotal = 1 + extensions.length
  const familyId = crypto.randomUUID()
  const queueAutoRun = fields.queueAutoRun === 'true'
  const refineExtensionFrame = parseRefineExtensionFrame(fields.refineExtensionFrame, true)
  const saveLosslessAnchor = parseSaveLosslessAnchor(fields.saveLosslessAnchor, false)
  const refinementDenoise = clampRefinementDenoise(fields.refinementDenoise, REFINEMENT_DENOISE_DEFAULT)
  const studio = await addStudioJob({
    ownerKey,
    familyId,
    payload: {
      prompt,
      promptMid: promptMid || undefined,
      promptPre: promptPre || undefined,
      promptPost: promptPost || undefined,
      name: clipName,
      folderId,
      aspect: fields.aspect || 'auto',
      width,
      height,
      steps,
      turbo,
      seed,
      cfg,
      fps,
      samplerName,
      scheduler,
      duration: durationSeconds,
      sound,
      workflow,
      useIdentityRefs,
      stillId: still?.id,
      stillFilename: still?.filename,
      hideThumbnail,
      hideInput,
      folderLocked,
      referenceStillIds,
      extensions,
      queueAutoRun,
      refineExtensionFrame,
      saveLosslessAnchor,
      refinementDenoise,
      globalLocks: globalLocks || undefined,
      permanenceRefs: permanenceRefs.length ? permanenceRefs : undefined,
      shotPermanenceRefs: shotPermanenceRefs.some(item => item.length) ? shotPermanenceRefs : undefined,
      ...persistLoraFields(defaultStack),
      shotLoras: shotLoraStacks.some(stack => stack.length) ? shotLoraStacks.map(stack => stack[0]?.name || '') : undefined,
      shotLoraStacks: shotLoraStacks.some(stack => stack.length) ? shotLoraStacks : undefined
    }
  })
  await kickStudioQueue()
  let latest = listStudioJobs(ownerKey).find(item => item.id === studio.id)
  if (!latest?.liveJobId && !(await videoJobsBusy())) {
    await kickStudioQueue()
    latest = listStudioJobs(ownerKey).find(item => item.id === studio.id)
  }
  const liveJobId = latest?.liveJobId || ''
  const queueId = latest?.shotQueueId

  return {
    jobId: liveJobId || studio.id,
    studioJobId: studio.id,
    queued: !liveJobId,
    clientId: '',
    seed,
    width,
    height,
    aspect: fields.aspect || 'auto',
    cfg,
    fps,
    sampler_name: samplerName,
    scheduler,
    steps,
    turbo,
    folderId,
    hideThumbnail,
    folderLocked,
    chainTotal,
    queueId,
    queuedRemaining: extensions.length
  }
})