aigen/server/api/edit.post.ts

import { imageComfyConfigured, uploadImageEdit, queueImagePrompt, waitForImageEdit, downloadEditedImage } from '~/server/utils/imageComfy'
import { ensureImageComfyReady } from '~/server/utils/imageComfyLifecycle'
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 4
  return Math.min(20, 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))
}

function parseDenoise(raw: unknown) {
  const value = Number(raw)
  if (!Number.isFinite(value)) return 0.75
  return Math.min(1, Math.max(0.05, Math.round(value * 100) / 100))
}

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
  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')
    }
  }

  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 (!imageComfyConfigured()) {
    throw createError({
      statusCode: 503,
      statusMessage: 'Image edit ComfyUI is not configured yet. Set IMAGE_COMFY_HOST on this instance.'
    })
  }

  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 denoise = parseDenoise(fields.denoise)
  const seed = fields.seed && fields.seed !== 'random'
    ? Number(fields.seed)
    : Math.floor(Math.random() * 2_147_483_647)
  const size = imageDimensions(image.data)

  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
  }
  emitJob(job, { type: 'status', message: 'Checking sidecar ComfyUI...', progress: 2 })

  void runEdit(job, {
    image,
    prompt,
    negative: (fields.negative || '').trim(),
    steps,
    seed,
    cfg,
    denoise
  }).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
  }
})

async function runEdit(
  job: ReturnType<typeof createJob>,
  params: {
    image: { filename: string; data: Buffer; type?: string }
    prompt: string
    negative: string
    steps: number
    seed: number
    cfg: number
    denoise: number
  }
) {
  await ensureImageComfyReady((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.')

  emitJob(job, { type: 'status', message: 'Uploading still to Qwen Edit...', progress: 8 })
  const uploaded = await uploadImageEdit(params.image, job.id)
  if (job.status === 'cancelled') throw new Error('Job interrupted.')

  emitJob(job, { type: 'status', message: 'Queueing Qwen Image Edit Lightning...', progress: 12 })
  const graph = buildEditWorkflow({
    prompt: params.prompt,
    negative: params.negative,
    imageName: uploaded.name,
    steps: params.steps,
    seed: params.seed,
    cfg: params.cfg,
    denoise: params.denoise,
    filenamePrefix: `aigen_edit_${job.id.slice(0, 8)}`
  })
  const queued = await queueImagePrompt(graph, job.clientId)
  job.promptId = queued.prompt_id
  job.status = 'running'
  emitJob(job, { type: 'status', message: 'Editing still...', progress: 18, maxStep: params.steps })

  const output = await waitForImageEdit({
    promptId: queued.prompt_id,
    clientId: job.clientId,
    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
  job.status = 'complete'
  emitJob(job, {
    type: 'complete',
    message: 'Edit finished',
    progress: 100,
    stillId: still.id,
    filename: output.filename,
    subfolder: output.subfolder,
    mediaType: 'image',
    hideThumbnail: job.hideThumbnail,
    folderLocked: job.library?.folderLocked
  })
}