aigen/server/api/edit.post.ts

import { ensureComfyReady } from '~/server/utils/comfyLifecycle'
import { getComfyHost, comfyConfigured, uploadImage, queuePrompt, purgeComfyArtifacts } from '~/server/utils/comfy'
import { withImageComfyHost, waitForImageEdit, downloadEditedImage } from '~/server/utils/imageComfy'
import { buildEditWorkflow } from '~/server/utils/imageWorkflow'
import { imageDimensions } from '~/server/utils/resolution'

function parseSteps(raw: unknown) {
  const value = Number(raw)
  if (!Number.isFinite(value)) return 20
  return Math.min(40, Math.max(1, Math.round(value)))
}

function parseCfg(raw: unknown) {
  const value = Number(raw)
  if (!Number.isFinite(value)) return 1
  return Math.min(8, Math.max(0.5, Math.round(value * 10) / 10))
}

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 prompt = (fields.prompt || '').trim()
  if (!image) {
    throw createError({ statusCode: 400, statusMessage: 'An input image is required' })
  }
  if (!prompt) {
    throw createError({ statusCode: 400, statusMessage: 'An edit prompt is required' })
  }
  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 steps = parseSteps(fields.steps)
  const cfg = parseCfg(fields.cfg)
  const seed = fields.seed && fields.seed !== 'random'
    ? Number(fields.seed)
    : Math.floor(Math.random() * 2_147_483_647)
  const size = imageDimensions(image.data)

  const still = isPipelineFrameFilename(image.filename)
    ? null
    : await saveStill({
      ownerKey,
      folderId,
      filename: image.filename,
      data: image.data,
      width: size?.width || 0,
      height: size?.height || 0,
      hideInput
    })
  if (reference && !isPipelineFrameFilename(reference.filename)) {
    await saveStill({
      ownerKey,
      folderId,
      filename: reference.filename,
      data: reference.data,
      hideInput
    })
  }

  const job = createJob()
  job.kind = 'edit'
  job.maxStep = steps
  job.hideThumbnail = hideInput
  job.library = {
    ownerKey,
    folderId,
    hideThumbnail: hideInput,
    hideInput,
    folderLocked,
    name: (fields.name || '').trim().slice(0, 80),
    prompt,
    aspect: 'auto',
    width: size?.width || 0,
    height: size?.height || 0,
    steps,
    turbo: true,
    seed,
    cfg,
    stillId: still?.id
  }
  emitJob(job, { type: 'status', message: 'Checking Beast ComfyUI...', progress: 2 })

  void runEdit(job, {
    image,
    reference,
    prompt,
    negative: (fields.negative || '').trim(),
    steps,
    seed,
    cfg
  }).catch(async (error) => {
    if (job.status === 'error' || job.status === 'cancelled') 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,
    seed,
    steps,
    hideThumbnail: hideInput,
    folderLocked,
    dual: Boolean(reference)
  }
})

async function runEdit(
  job: ReturnType<typeof createJob>,
  params: {
    image: { filename: string; data: Buffer; type?: string }
    reference: { filename: string; data: Buffer; type?: string } | null
    prompt: string
    negative: string
    steps: number
    seed: number
    cfg: number
  }
) {
  await ensureComfyReady((status) => {
    emitJob(job, {
      type: status.state === 'busy' ? 'busy' : 'status',
      message: status.message,
      progress: status.state === 'online' ? Math.max(job.progress, 6) : Math.max(job.progress, 3),
      busy: status.state === 'busy'
    })
  })
  if (job.status === 'cancelled') throw new Error('Job interrupted.')
  job.imageComfyHost = getComfyHost()

  await withImageComfyHost(job.imageComfyHost, async () => {
    const dual = Boolean(params.reference)
    emitJob(job, {
      type: 'status',
      message: dual ? 'Uploading both stills to Beast...' : 'Uploading still to Beast...',
      progress: 8
    })
    const uploaded = await uploadImage(params.image, job.id)
    const uploadedRef = params.reference
      ? await uploadImage({
        ...params.reference,
        filename: `ref_${params.reference.filename || 'image2.png'}`
      }, job.id)
      : null
    if (job.status === 'cancelled') throw new Error('Job interrupted.')

    emitJob(job, {
      type: 'status',
      message: dual
        ? 'Queueing Flux.2 Klein two-image edit on Beast...'
        : 'Queueing Flux.2 Klein edit on Beast...',
      progress: 12
    })
    const graph = buildEditWorkflow({
      prompt: params.prompt,
      negative: params.negative,
      imageName: uploaded.name,
      referenceImageName: uploadedRef?.name,
      steps: params.steps,
      seed: params.seed,
      cfg: params.cfg,
      filenamePrefix: `aigen_edit_${job.id.slice(0, 8)}`
    })
    const queued = await queuePrompt(graph, job.clientId)
    job.promptId = queued.prompt_id
    job.status = 'running'
    emitJob(job, {
      type: 'status',
      message: dual ? 'Applying image 2 onto image 1...' : 'Editing still on Beast...',
      progress: 18,
      maxStep: params.steps
    })

    const output = await waitForImageEdit({
      promptId: queued.prompt_id,
      clientId: job.clientId,
      timeoutMs: 10 * 60 * 1000,
      onProgress: (event) => {
        emitJob(job, {
          type: 'status',
          message: event.message,
          progress: event.progress,
          step: event.step,
          maxStep: event.maxStep || params.steps,
          node: event.node
        })
      },
      isCancelled: () => job.status === 'cancelled'
    })

    emitJob(job, { type: 'status', message: 'Saving edited still...', progress: 94 })
    const buffer = await downloadEditedImage(output)
    const size = imageDimensions(buffer)
    const still = await saveStill({
      ownerKey: job.library!.ownerKey,
      folderId: job.library!.folderId,
      filename: job.library?.name ? `${job.library.name}.png` : output.filename,
      data: buffer,
      width: size?.width || 0,
      height: size?.height || 0,
      hideInput: job.library?.hideInput
    })
    job.stillId = still?.id
    await purgeComfyArtifacts({
      video: { filename: output.filename, subfolder: output.subfolder, type: output.type },
      imageName: uploaded.name,
      imageSubfolder: uploaded.subfolder,
      extraImageNames: uploadedRef?.name ? [uploadedRef.name] : [],
      promptId: job.promptId
    })
    job.status = 'complete'
    emitJob(job, {
      type: 'complete',
      message: 'Edit finished on Beast',
      progress: 100,
      stillId: still?.id,
      filename: output.filename,
      subfolder: output.subfolder,
      mediaType: 'image',
      hideThumbnail: job.hideThumbnail,
      folderLocked: job.library?.folderLocked
    })
  })
}