aigen/server/utils/imageChainV2.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 { buildImageV2Workflow } from '~/server/utils/imageWorkflowV2'
import { ensureComfyLoraNames } from '~/server/utils/loras'
import { imageDimensions } from '~/server/utils/resolution'
import { emitChainJob } from '~/server/utils/watch'
import { saveStill } from '~/server/utils/library'
import type { ImageV2Mode, ImageV2Task } from '~/utils/imageV2'
import type { EditImageFile } from '~/server/utils/imageChain'

export type EditV2RunParams = {
  mode: ImageV2Mode
  task: ImageV2Task
  image?: EditImageFile | null
  reference: EditImageFile | null
  mask?: EditImageFile | null
  prompt: string
  negative: string
  steps: number
  seed: number
  cfg: number
  snofsModel: number
  snofsClip: number
  consistencyModel: number
  consistencyClip: number
  megapixels: number
  strength?: number
  width?: number
  height?: number
  turbo?: boolean
}

export async function runEditV2(job: Job, params: EditV2RunParams) {
  const library = job.library
  if (!library) throw new Error('Edit job is missing library metadata')
  if (params.mode === 'refine' && !params.mask) {
    throw new Error('Refine requires a mask. Refusing to fall back to Edit.')
  }
  if (params.mode === 'compose' && !params.reference) {
    throw new Error('Compose requires image B. Refusing to fall back to one-image generation.')
  }
  if (params.mode === 'edit' && params.reference) {
    throw new Error('Edit mode takes one image. Use Compose for two stills.')
  }
  if (params.mode !== 'generate' && !params.image) {
    throw new Error('This v2 mode requires still A.')
  }

  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 () => {
      const generate = params.mode === 'generate'
      emitChainJob(job, {
        type: 'status',
        message: generate
          ? 'Queueing Klein v2 generate on Beast...'
          : params.mode === 'refine'
            ? 'Uploading canvas and mask to Beast...'
            : params.mode === 'compose' ? 'Uploading stills A and B to Beast...' : 'Uploading still A to Beast...',
        progress: generate ? 12 : 8
      })
      const uploaded = generate || !params.image ? null : await uploadImage(params.image, job.id)
      const uploadedRef = generate || params.mode === 'refine'
        ? null
        : params.reference
          ? await uploadImage({
            ...params.reference,
            filename: `ref_${params.reference.filename || 'image_b.png'}`
          }, job.id)
          : null
      const uploadedMask = !generate && params.mode === 'refine' && params.mask
        ? await uploadImage({
          ...params.mask,
          filename: `mask_${params.mask.filename || 'refine-mask.png'}`
        }, job.id)
        : null
      if (job.status === 'cancelled') throw new Error('Job interrupted.')

      if (!generate) {
        emitChainJob(job, {
          type: 'status',
          message: params.mode === 'refine'
            ? 'Queueing Klein v2 refine on Beast...'
            : params.mode === 'compose'
              ? `Queueing Klein v2 compose (${params.task}) on Beast...`
              : 'Queueing Klein v2 edit on Beast...',
          progress: 12
        })
      }
      await ensureComfyLoraNames('image')
      if (!generate) await assertImageScaleToTotalPixelsNode()
      const built = buildImageV2Workflow({
        mode: params.mode,
        task: params.task,
        prompt: params.prompt,
        negative: params.negative,
        imageAName: uploaded?.name,
        imageBName: uploadedRef?.name,
        maskName: uploadedMask?.name,
        strength: params.strength,
        snofsModel: params.snofsModel,
        snofsClip: params.snofsClip,
        consistencyModel: params.consistencyModel,
        consistencyClip: params.consistencyClip,
        steps: params.steps,
        cfg: params.cfg,
        seed: params.seed,
        megapixels: params.megapixels,
        width: params.width,
        height: params.height,
        turbo: params.turbo === true,
        filenamePrefix: `aigen_v2_${job.id.slice(0, 8)}`
      })
      const queued = await queuePrompt(built.graph, job.clientId)
      job.promptId = queued.prompt_id
      job.status = 'running'
      emitChainJob(job, {
        type: 'status',
        message: `Running ${built.workflowFile}...`,
        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 Klein v2 still...', progress: 94 })
      const buffer = await downloadEditedImage(output)
      const size = imageDimensions(buffer)
      const still = await saveStill({
        ownerKey: library.ownerKey,
        folderId: library.folderId,
        filename: output.filename,
        data: buffer,
        width: size?.width || 0,
        height: size?.height || 0,
        hideInput: job.hideThumbnail === true,
        role: 'output',
        name: library.name || undefined,
        prompt: built.prompt
      })
      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, uploadedMask?.name].filter((name): name is string => Boolean(name)),
        promptId: job.promptId
      })

      job.status = 'complete'
      emitChainJob(job, {
        type: 'complete',
        message: 'Klein v2 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)
  }
}