aigen/server/utils/imageChain.ts

import { createJob, emitJob, type Job } from '~/server/utils/jobs'
import { ensureComfyReady } from '~/server/utils/comfyLifecycle'
import { assertImageScaleToTotalPixelsNode, getComfyHost, uploadImage, queuePrompt, purgeComfyArtifacts } from '~/server/utils/comfy'
import { withImageComfyHost, waitForImageEdit, downloadEditedImage } from '~/server/utils/imageComfy'
import { buildEditWorkflow } from '~/server/utils/imageWorkflow'
import { ensureComfyLoraNames } from '~/server/utils/loras'
import { imageDimensions } from '~/server/utils/resolution'
import { emitChainJob } from '~/server/utils/watch'
import { saveStill, stillChainName } from '~/server/utils/library'

export type EditImageFile = { filename: string; data: Buffer; type?: string }

export type EditRunParams = {
  image: EditImageFile
  reference: EditImageFile | null
  prompt: string
  passes: { prompt: string }[]
  negative: string
  steps: number
  seed: number
  cfg: number
  loraName?: string
  loraStack?: import('~/utils/loras').LoraStackItem[]
  scaleToTotalPixels?: boolean
  scaleMegapixels?: number
}

export async function runEdit(job: Job, params: EditRunParams) {
  const library = job.library
  if (!library) throw new Error('Edit job is missing library metadata')
  const prompts = [params.prompt, ...params.passes.map(item => item.prompt)]
  const chainTotal = prompts.length
  library.chainTotal = chainTotal
  library.familyId = chainTotal > 1 ? (library.familyId || crypto.randomUUID()) : library.familyId

  try {
    await ensureComfyReady((status) => {
      emitChainJob(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'
      })
    }, { skipBusyWait: true })
    if (job.status === 'cancelled') throw new Error('Job interrupted.')
    job.imageComfyHost = getComfyHost()

    await withImageComfyHost(job.imageComfyHost, async () => {
      let current = params.image
      let parentStillId: string | undefined
      for (let index = 0; index < prompts.length; index++) {
        if (job.status === 'cancelled') throw new Error('Job interrupted.')
        if (index > 0 && library.stopAfterCurrent === true) break
        const prompt = prompts[index]
        const last = index === prompts.length - 1
        if (index > 0) {
          await ensureComfyReady((status) => {
            emitChainJob(job, {
              type: status.state === 'busy' ? 'busy' : 'status',
              message: status.message,
              progress: status.state === 'online' ? 4 : 2,
              busy: status.state === 'busy'
            })
          }, { skipBusyWait: true })
        }
        const seed = index === 0 ? params.seed : Math.floor(Math.random() * 2_147_483_647)
        library.prompt = prompt
        library.seed = seed
        library.chainIndex = index
        library.chainStep = index + 1
        library.chainLabel = chainTotal > 1 ? `Pass ${index + 1}` : undefined
        const reference = index === 0 ? params.reference : null
        const dual = Boolean(reference)
        const passName = stillChainName(library.name || '', index)

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

        emitChainJob(job, {
          type: 'status',
          message: dual
            ? 'Queueing Flux.2 Klein two-image edit on Beast...'
            : 'Queueing Flux.2 Klein edit on Beast...',
          progress: 12
        })
        await ensureComfyLoraNames('image')
        if (params.scaleToTotalPixels) await assertImageScaleToTotalPixelsNode()
        const graph = buildEditWorkflow({
          prompt,
          negative: params.negative,
          imageName: uploaded.name,
          referenceImageName: uploadedRef?.name,
          steps: params.steps,
          seed,
          cfg: params.cfg,
          filenamePrefix: `aigen_edit_${job.id.slice(0, 8)}_p${index + 1}`,
          loraStack: params.loraStack || params.loraName,
          scaleToTotalPixels: params.scaleToTotalPixels,
          scaleMegapixels: params.scaleMegapixels
        })
        const queued = await queuePrompt(graph, job.clientId)
        job.promptId = queued.prompt_id
        job.status = 'running'
        emitChainJob(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) => {
            emitChainJob(job, {
              type: 'status',
              message: event.message,
              progress: event.progress,
              step: event.step,
              maxStep: event.maxStep || params.steps,
              node: event.node
            })
          },
          isCancelled: () => job.status === 'cancelled'
        })

        emitChainJob(job, { type: 'status', message: 'Saving edited still...', progress: 94 })
        const buffer = await downloadEditedImage(output)
        const size = imageDimensions(buffer)
        const still = await saveStill({
          ownerKey: library.ownerKey,
          folderId: library.folderId,
          filename: passName ? `${passName}.png` : output.filename,
          data: buffer,
          width: size?.width || 0,
          height: size?.height || 0,
          hideInput: job.hideThumbnail === true,
          role: 'output',
          name: passName || undefined,
          prompt,
          familyId: library.familyId,
          parentStillId,
          chainIndex: index,
          settings: {
            kind: 'image',
            steps: params.steps,
            cfg: params.cfg,
            negative: params.negative,
            seed,
            scaleToTotalPixels: params.scaleToTotalPixels === true,
            scaleMegapixels: params.scaleMegapixels,
            loraName: params.loraName,
            loraStack: params.loraStack,
            sourceStillId: library.stillId
          }
        })
        job.stillId = still?.id
        parentStillId = 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
        })

        if (!last) {
          emitChainJob(job, {
            type: 'checkpoint',
            message: `Pass ${index + 1} saved`,
            progress: 100,
            stillId: still?.id,
            hideThumbnail: job.hideThumbnail,
            folderLocked: library.folderLocked
          })
          current = {
            filename: still?.filename || `pass_${index + 1}.png`,
            data: buffer,
            type: 'image/png'
          }
          if (library.stopAfterCurrent === true) break
          continue
        }

        job.status = 'complete'
        emitChainJob(job, {
          type: 'complete',
          message: chainTotal > 1 ? 'Edit chain finished on Beast' : 'Edit finished on Beast',
          progress: 100,
          stillId: still?.id,
          filename: output.filename,
          subfolder: output.subfolder,
          mediaType: 'image',
          hideThumbnail: job.hideThumbnail,
          folderLocked: library.folderLocked
        })
      }
    })
  } catch (error) {
    if (job.status !== 'error' && job.status !== 'cancelled') {
      const message = error instanceof Error ? error.message : String(error)
      job.status = 'error'
      job.error = message
      emitJob(job, { type: 'error', error: message, message })
    }
  } finally {
    const { onLiveVideoSettled } = await import('~/server/utils/studioQueue')
    await onLiveVideoSettled(job)
  }
}

export function createEditLiveJob(params: {
  steps: number
  hideThumbnail: boolean
  library: NonNullable<Job['library']>
}) {
  const job = createJob()
  job.kind = 'edit'
  job.maxStep = params.steps
  job.hideThumbnail = params.hideThumbnail
  job.library = params.library
  return job
}