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, IMAGE_V2_NODE_LABELS, resolveKreaGenerateAssets } from '~/server/utils/imageWorkflowV2'
import { ensureComfyLoraNames, persistLoraFields } from '~/server/utils/loras'
import { imageDimensions } from '~/server/utils/resolution'
import { emitChainJob } from '~/server/utils/watch'
import { saveStill, stillChainName } from '~/server/utils/library'
import { comfyJobPrefix } from '~/utils/outputNames'
import { parseImageV2Engine, type ImageV2Engine, type ImageV2Mode, type ImageV2Task } from '~/utils/imageV2'
import type { EditImageFile } from '~/server/utils/imageChain'

export type EditV2RunParams = {
  mode: ImageV2Mode
  engine?: ImageV2Engine
  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
  aspect?: string
  sourceStillId?: string
  referenceStillId?: string
  loraStack?: import('~/utils/loras').LoraStackItem[]
  passes?: { prompt: string }[]
}

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

  const canChain = params.mode === 'edit' || params.mode === 'compose'
  const extraPasses = canChain
    ? (params.passes || []).map(item => ({ prompt: String(item.prompt || '').trim() })).filter(item => item.prompt)
    : []
  const prompts = [params.prompt, ...extraPasses.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 () => {
      const generate = params.mode === 'generate'
      const engine = parseImageV2Engine(params.engine, 'flux')
      const engineLabel = engine === 'krea' ? 'Krea' : 'Flux'
      const { applyImageFill } = await import('~/server/utils/imagePad')
      const { composeFillPrompt } = await import('~/utils/imageAspect')
      let current = params.image || null
      let mask = params.mask
      let parentStillId: string | undefined
      let uploadedRef: { name: string } | null = null
      let uploadedMask: { name: string } | null = null

      await ensureComfyLoraNames('image')
      if (!generate) await assertImageScaleToTotalPixelsNode()
      const krea = engine === 'krea' ? await resolveKreaGenerateAssets() : null

      if (!generate && params.mode === 'compose' && params.reference) {
        uploadedRef = await uploadImage({
          ...params.reference,
          filename: `ref_${params.reference.filename || 'image_b.png'}`
        }, job.id)
      }

      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 last = index === prompts.length - 1
        const seed = index === 0 ? params.seed : Math.floor(Math.random() * 2_147_483_647)
        library.prompt = prompts[index]
        library.seed = seed
        library.chainIndex = index
        library.chainStep = index + 1
        library.chainLabel = chainTotal > 1 ? `Pass ${index + 1}` : undefined
        const passName = stillChainName(library.name || '', index)
        let prompt = prompts[index]
        let image = current
        let padded = false

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

        emitChainJob(job, {
          type: 'status',
          message: generate
            ? `Queueing ${engineLabel} 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
        })

        if (index === 0 && !generate && image) {
          const filled = await applyImageFill({
            image,
            mask: params.mode === 'refine' ? mask : null,
            aspect: params.aspect || 'auto'
          })
          image = filled.image
          current = filled.image
          if (filled.mask) {
            mask = filled.mask
            uploadedMask = await uploadImage({
              ...filled.mask,
              filename: `mask_${filled.mask.filename || 'refine-mask.png'}`
            }, job.id)
          }
          padded = filled.padded
          prompt = composeFillPrompt(prompt, padded)
          if (padded) {
            emitChainJob(job, {
              type: 'status',
              message: `Filling to ${filled.aspect} without stretching...`,
              progress: 7
            })
          }
        }

        if (index === 0 && params.mode === 'refine' && mask && !uploadedMask) {
          uploadedMask = await uploadImage({
            ...mask,
            filename: `mask_${mask.filename || 'refine-mask.png'}`
          }, job.id)
        }

        const uploaded = generate || !image ? null : await uploadImage(image, job.id)
        if (job.status === 'cancelled') throw new Error('Job interrupted.')

        if (!generate) {
          emitChainJob(job, {
            type: 'status',
            message: params.mode === 'refine'
              ? `Queueing ${engineLabel} v2 refine on Beast...`
              : params.mode === 'compose'
                ? `Queueing ${engineLabel} v2 compose (${params.task}) on Beast...`
                : `Queueing ${engineLabel} v2 edit on Beast...`,
            progress: 12
          })
        }

        const built = buildImageV2Workflow({
          mode: params.mode,
          engine,
          task: params.task,
          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,
          megapixels: params.megapixels,
          width: params.width,
          height: params.height,
          turbo: params.turbo === true,
          filenamePrefix: chainTotal > 1 ? `${comfyJobPrefix(job.id)}_p${index + 1}` : comfyJobPrefix(job.id),
          kreaUnetName: krea?.unet,
          kreaClipName: krea?.clip,
          kreaVaeName: krea?.vae,
          kreaConceptLora: krea?.conceptLora,
          loraStack: params.loraStack
        })
        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,
          engineLabel,
          nodeLabel: (node) => {
            if (node === '15' || node === '19') return `Sampling ${engineLabel}`
            return IMAGE_V2_NODE_LABELS[node]
          },
          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 ${engineLabel} v2 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: 'imagev2',
            mode: params.mode,
            engine,
            task: params.task,
            negative: params.negative,
            snofsModel: params.snofsModel,
            snofsClip: params.snofsClip,
            consistencyModel: params.consistencyModel,
            consistencyClip: params.consistencyClip,
            steps: params.steps,
            cfg: params.cfg,
            megapixels: params.mode === 'generate' ? undefined : params.megapixels,
            turbo: params.turbo === true,
            strength: params.mode === 'generate' ? undefined : params.strength,
            width: params.mode === 'generate' ? params.width : undefined,
            height: params.mode === 'generate' ? params.height : undefined,
            aspect: params.aspect || 'auto',
            seed,
            sourceStillId: params.sourceStillId || library.stillId,
            referenceStillId: params.referenceStillId,
            ...persistLoraFields(params.loraStack)
          }
        })
        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: last
            ? [uploadedRef?.name, uploadedMask?.name].filter((name): name is string => Boolean(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
            ? `${engineLabel} v2 chain finished on Beast`
            : `${engineLabel} 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)
  }
}