aigen/server/api/v2/generate.post.ts

import { addStudioJob, kickStudioQueue, listStudioJobs, videoJobsBusy } from '~/server/utils/studioQueue'
import { comfyConfigured } from '~/server/utils/comfy'
import { imageDimensions } from '~/server/utils/resolution'
import { clampImageCfg, clampImageScaleMegapixels, clampImageSteps } from '~/utils/generationPresets'
import {
  IMAGE_V2_CFG_DEFAULT,
  IMAGE_V2_CONSISTENCY_CLIP,
  IMAGE_V2_CONSISTENCY_MODEL,
  IMAGE_V2_SNOFS_CLIP,
  IMAGE_V2_SNOFS_MODEL,
  IMAGE_V2_STEPS_DEFAULT,
  IMAGE_V2_TURBO_CFG,
  IMAGE_V2_TURBO_STEPS,
  IMAGE_V2_DENOISE_DEFAULT,
  clampImageV2Denoise,
  clampImageV2Strength,
  parseImageV2Mode,
  parseImageV2Task,
  type ImageV2Mode,
  type ImageV2Task
} from '~/utils/imageV2'
import { getStill, rememberInputStill, stillPath } from '~/server/utils/library'
import { existsSync, readFileSync } from 'node:fs'

type ImageFile = { filename: string; data: Buffer; type?: string }

function parseBool(raw: unknown) {
  return raw === true || raw === 'true' || raw === '1' || raw === 1
}

function isHttpUrl(value: string) {
  return /^https?:\/\//i.test(value)
}

async function fileFromUrl(url: string): Promise<ImageFile> {
  const res = await fetch(url, { signal: AbortSignal.timeout(20_000) })
  if (!res.ok) {
    throw createError({ statusCode: 400, statusMessage: `Could not fetch image (${res.status})` })
  }
  const mime = String(res.headers.get('content-type') || 'image/png').split(';')[0]
  if (!/^image\//i.test(mime)) {
    throw createError({ statusCode: 400, statusMessage: 'image_a / image_b / mask URL must be an image' })
  }
  const data = Buffer.from(await res.arrayBuffer())
  if (data.length > 40 * 1024 * 1024) {
    throw createError({ statusCode: 400, statusMessage: 'Image is too large' })
  }
  const ext = mime.includes('jpeg') || mime.includes('jpg') ? '.jpg' : mime.includes('webp') ? '.webp' : '.png'
  return { filename: `remote${ext}`, data, type: mime }
}

function fileFromStill(ownerKey: string, id: string): ImageFile {
  const still = getStill(ownerKey, id)
  const path = stillPath(ownerKey, still.id)
  if (!existsSync(path)) {
    throw createError({ statusCode: 400, statusMessage: 'That still is missing from the library' })
  }
  return {
    filename: still.filename || `${still.id}.png`,
    data: readFileSync(path),
    type: 'image/png'
  }
}

async function resolveImageRef(ownerKey: string, raw: unknown, uploaded: ImageFile | null) {
  if (uploaded) return uploaded
  const value = String(raw || '').trim()
  if (!value || value === 'null') return null
  if (isHttpUrl(value)) return fileFromUrl(value)
  return fileFromStill(ownerKey, value)
}

function readMultipart(parts: Array<{ name?: string; filename?: string; type?: string; data?: Buffer }> | null) {
  const fields: Record<string, string> = {}
  let imageA: ImageFile | null = null
  let imageB: ImageFile | null = null
  let mask: ImageFile | null = null
  for (const part of parts || []) {
    if ((part.name === 'image_a' || part.name === 'image') && part.filename && part.data?.length) {
      imageA = { filename: part.filename, data: part.data, type: part.type }
    } else if ((part.name === 'image_b' || part.name === 'image2') && part.filename && part.data?.length) {
      imageB = { filename: part.filename, data: part.data, type: part.type }
    } else if (part.name === 'mask' && part.filename && part.data?.length) {
      mask = { filename: part.filename, data: part.data, type: part.type }
    } else if (part.name && part.data) {
      fields[part.name] = part.data.toString('utf8')
    }
  }
  return { fields, imageA, imageB, mask }
}

export default defineEventHandler(async (event) => {
  const contentType = String(getHeader(event, 'content-type') || '')
  let fields: Record<string, unknown> = {}
  let uploadedA: ImageFile | null = null
  let uploadedB: ImageFile | null = null
  let uploadedMask: ImageFile | null = null

  if (contentType.includes('multipart/form-data')) {
    const form = await readMultipartFormData(event).catch(() => null)
    const parsed = readMultipart(form || [])
    fields = parsed.fields
    uploadedA = parsed.imageA
    uploadedB = parsed.imageB
    uploadedMask = parsed.mask
  } else {
    fields = await readBody<Record<string, unknown>>(event).catch(() => ({}))
  }

  const mode = parseImageV2Mode(fields.mode)
  if (!mode) {
    throw createError({ statusCode: 400, statusMessage: 'mode must be edit, compose, or refine' })
  }
  const task = parseImageV2Task(fields.task, mode === 'refine' ? 'refine' : 'scene')
  const prompt = String(fields.prompt || '').trim()
  if (!prompt) {
    throw createError({ statusCode: 400, statusMessage: 'A prompt is required' })
  }
  if (!comfyConfigured()) {
    throw createError({
      statusCode: 503,
      statusMessage: 'Beast ComfyUI is not configured. Set COMFY_HOST.'
    })
  }

  if (mode === 'edit' && (task === 'identity' || task === 'outfit' || task === 'face_lock')) {
    throw createError({
      statusCode: 400,
      statusMessage: `${task} requires Compose and image B. Edit is one image only.`
    })
  }

  const ownerKey = libraryOwnerKey(event)
  const imageA = await resolveImageRef(ownerKey, fields.image_a, uploadedA)
  const imageB = mode === 'refine' ? null : await resolveImageRef(ownerKey, fields.image_b, uploadedB)
  const mask = mode === 'refine' ? await resolveImageRef(ownerKey, fields.mask, uploadedMask) : null

  if (!imageA) {
    throw createError({ statusCode: 400, statusMessage: 'image_a is required' })
  }
  if (mode === 'refine' && !mask) {
    throw createError({
      statusCode: 400,
      statusMessage: 'Refine requires a mask. Refusing to fall back to Edit.'
    })
  }
  if (mode === 'edit' && imageB) {
    throw createError({
      statusCode: 400,
      statusMessage: 'Edit mode takes one image. Use Compose for two stills.'
    })
  }
  if (mode === 'compose' && !imageB) {
    throw createError({
      statusCode: 400,
      statusMessage: 'Compose requires image_b. Refusing to fall back to one-image generation.'
    })
  }

  const library = publicLibrary(event)
  const folderId = library.folders.some(folder => folder.id === String(fields.folderId || ''))
    ? String(fields.folderId)
    : library.folders[0]?.id
  if (!folderId) {
    throw createError({ statusCode: 400, statusMessage: 'Create a library folder before generating' })
  }
  assertFolderExists(event, folderId)

  const destFolder = library.folders.find(folder => folder.id === folderId)
  const folderLocked = Boolean(destFolder?.protected && !destFolder.unlocked)
  const hideInput = parseBool(fields.hideInput)
  const hideThumbnail = parseBool(fields.hideThumbnail)
  const turbo = parseBool(fields.turbo)
  const steps = turbo ? IMAGE_V2_TURBO_STEPS : clampImageSteps(fields.steps, IMAGE_V2_STEPS_DEFAULT)
  const cfg = turbo ? IMAGE_V2_TURBO_CFG : clampImageCfg(fields.cfg, IMAGE_V2_CFG_DEFAULT)
  const megapixels = clampImageScaleMegapixels(fields.megapixels ?? fields.scaleMegapixels, 1)
  const seed = fields.seed && String(fields.seed) !== 'random'
    ? Number(fields.seed)
    : Math.floor(Math.random() * 2_147_483_647)
  const size = imageDimensions(imageA.data)
  const clipName = String(fields.name || '').trim().slice(0, 80)
  const v2Mode = mode as ImageV2Mode
  const v2Task = (mode === 'refine' ? 'refine' : mode === 'edit' ? 'scene' : task) as ImageV2Task
  const refineStrength = mode === 'refine' ? clampImageV2Denoise(fields.strength, IMAGE_V2_DENOISE_DEFAULT) : undefined

  const still = await rememberInputStill({
    ownerKey,
    folderId,
    filename: imageA.filename,
    data: imageA.data,
    width: size?.width || 0,
    height: size?.height || 0,
    hideInput
  })
  const savedRef = imageB
    ? await rememberInputStill({
      ownerKey,
      folderId,
      filename: imageB.filename,
      data: imageB.data,
      hideInput
    })
    : null
  const savedMask = mask
    ? await rememberInputStill({
      ownerKey,
      folderId,
      filename: mask.filename || 'refine-mask.png',
      data: mask.data,
      hideInput
    })
    : null

  const studio = await addStudioJob({
    ownerKey,
    familyId: crypto.randomUUID(),
    kind: 'edit',
    payload: {
      prompt,
      name: clipName,
      folderId,
      aspect: 'auto',
      width: size?.width || 0,
      height: size?.height || 0,
      steps,
      turbo,
      seed,
      cfg,
      fps: 24,
      samplerName: 'euler',
      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: String(fields.negative || '').trim(),
      referenceStillId: savedRef?.id,
      referenceStillFilename: savedRef?.filename,
      scaleToTotalPixels: true,
      scaleMegapixels: megapixels,
      imagePipeline: 'v2',
      v2Mode,
      v2Task,
      snofsModel: clampImageV2Strength(fields.snofs_model, IMAGE_V2_SNOFS_MODEL),
      snofsClip: clampImageV2Strength(fields.snofs_clip, IMAGE_V2_SNOFS_CLIP),
      consistencyModel: clampImageV2Strength(fields.consistency_model, IMAGE_V2_CONSISTENCY_MODEL),
      consistencyClip: clampImageV2Strength(fields.consistency_clip, IMAGE_V2_CONSISTENCY_CLIP),
      maskStillId: savedMask?.id,
      maskStillFilename: savedMask?.filename,
      refineStrength
    }
  })
  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,
    cfg,
    mode: v2Mode,
    task: v2Task,
    workflow: v2Mode === 'refine' ? 'klein_v2_refine.json' : v2Mode === 'compose' ? 'klein_v2_compose.json' : 'klein_v2_edit.json',
    strength: refineStrength,
    hideThumbnail,
    folderLocked
  }
})