aigen/server/api/edit.post.ts

import { addStudioJob, kickStudioQueue, listStudioJobs, videoJobsBusy } from '~/server/utils/studioQueue'
import { parsePostedLoraStack, listStudioLoras, persistLoraFields } from '~/server/utils/loras'
import { assertImageScaleToTotalPixelsNode, comfyConfigured } from '~/server/utils/comfy'
import { imageDimensions } from '~/server/utils/resolution'
import { parseImageAspect } from '~/utils/imageAspect'
import { clampImageCfg, clampImageScaleMegapixels, clampImageSteps, parseImagePassMode, parseImageScaleToTotalPixels } from '~/utils/generationPresets'

function parsePasses(raw: string | undefined) {
  if (!raw) return [] as { prompt: string }[]
  try {
    const parsed = JSON.parse(raw)
    if (!Array.isArray(parsed)) return []
    return parsed
      .map((item: { prompt?: unknown }) => ({ prompt: String(item?.prompt || '').trim() }))
      .filter(item => item.prompt)
  } catch {
    return []
  }
}

export default defineEventHandler(async (event) => {
  const form = await readMultipartFormData(event).catch(() => null)
  const fields: Record<string, string> = {}
  let image: { filename: string; data: Buffer; type?: string } | null = null
  let reference: { 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 === 'image2' || part.name === 'refImage') && part.filename && part.data?.length) {
      reference = { filename: part.filename, data: part.data, type: part.type }
    } else if (part.name && part.data) {
      fields[part.name] = part.data.toString('utf8')
    }
  }

  const promptMid = (fields.prompt || '').trim()
  const promptPre = (fields.promptPre || '').trim()
  const promptPost = (fields.promptPost || '').trim()
  const prompt = promptMid
  const extraPasses = parsePasses(fields.passes)
  const passMode = parseImagePassMode(fields.passMode)
  const aspect = parseImageAspect(fields.aspect)
  if (!image) {
    throw createError({ statusCode: 400, statusMessage: 'An input image is required' })
  }
  if (!promptMid && aspect === 'auto') {
    throw createError({ statusCode: 400, statusMessage: 'An edit prompt is required' })
  }
  if (extraPasses.some(item => !item.prompt)) {
    throw createError({ statusCode: 400, statusMessage: 'Each extra pass needs a prompt' })
  }
  if (!comfyConfigured()) {
    throw createError({
      statusCode: 503,
      statusMessage: 'Beast ComfyUI is not configured. Set COMFY_HOST.'
    })
  }

  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 editing' })
  }
  assertFolderExists(event, folderId)

  const destFolder = library.folders.find(folder => folder.id === folderId)
  const folderLocked = Boolean(destFolder?.protected && !destFolder.unlocked)
  const hideInput = fields.hideInput === 'true'
  const hideThumbnail = fields.hideThumbnail === 'true'
  const steps = clampImageSteps(fields.steps)
  const cfg = clampImageCfg(fields.cfg)
  const scaleToTotalPixels = parseImageScaleToTotalPixels(fields.scaleToTotalPixels)
  const scaleMegapixels = clampImageScaleMegapixels(fields.scaleMegapixels)
  if (scaleToTotalPixels) {
    try {
      await assertImageScaleToTotalPixelsNode()
    } catch (error) {
      const status = Number((error as { statusCode?: number })?.statusCode || 0)
      if (status === 400) throw error
    }
  }
  const seed = fields.seed && fields.seed !== 'random'
    ? Number(fields.seed)
    : Math.floor(Math.random() * 2_147_483_647)
  const size = imageDimensions(image.data)
  await listStudioLoras()
  const loraStack = parsePostedLoraStack(fields.loraStack || fields.lora, 'image')
  let inputFilename = image.filename
  const postedStillId = (fields.stillId || '').trim()
  if (postedStillId) {
    try {
      const existing = getStill(ownerKey, postedStillId)
      inputFilename = existing.name || existing.filename || inputFilename
    } catch { /* use the uploaded filename */ }
  }
  const clipName = allocateOutputName(ownerKey, folderId, (fields.name || '').trim(), {
    inputFilename,
    engine: 'klein',
    loraStack
  })

  const still = await rememberInputStill({
    ownerKey,
    folderId,
    filename: image.filename,
    data: image.data,
    width: size?.width || 0,
    height: size?.height || 0,
    hideInput
  })
  const savedRef = reference
    ? await rememberInputStill({
      ownerKey,
      folderId,
      filename: reference.filename,
      data: reference.data,
      hideInput
    })
    : null

  const passes = extraPasses
  const chainTotal = 1 + passes.length
  const familyId = crypto.randomUUID()
  const studio = await addStudioJob({
    ownerKey,
    familyId,
    kind: 'edit',
    payload: {
      prompt,
      promptMid: promptMid || undefined,
      promptPre: promptPre || undefined,
      promptPost: promptPost || undefined,
      name: clipName,
      folderId,
      aspect,
      width: size?.width || 0,
      height: size?.height || 0,
      steps,
      turbo: true,
      seed,
      cfg,
      fps: 24,
      samplerName: 'res_multistep',
      scheduler: 'simple',
      duration: 0,
      sound: false,
      workflow: 'v1',
      useIdentityRefs: false,
      stillId: still?.id,
      stillFilename: still?.filename,
      hideThumbnail,
      hideInput,
      folderLocked,
      referenceStillIds: [null, null, null, null],
      extensions: [],
      queueAutoRun: false,
      negative: (fields.negative || '').trim(),
      passes,
      passMode,
      referenceStillId: savedRef?.id,
      referenceStillFilename: savedRef?.filename,
      scaleToTotalPixels,
      scaleMegapixels,
      ...persistLoraFields(loraStack)
    }
  })
  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 || ''

  return {
    jobId: liveJobId || studio.id,
    studioJobId: studio.id,
    queued: !liveJobId,
    seed,
    steps,
    hideThumbnail,
    folderLocked,
    dual: Boolean(reference),
    chainTotal
  }
})