Add Image v2 as a sibling of Image edit, with separate Klein graphs.

Keep the v1 Image tab on its dual-branch template. The host agent follows Desktop's live port so Coolify stays on :8198. Co-authored-by: Cursor <cursoragent@cursor.com>
2026-08-28 21:22:37 -05:00
parent 55f8cb1def
commit 21fe2412f9
18 changed files with 1678 additions and 67 deletions
@@ -0,0 +1,168 @@
+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
+  reference: EditImageFile | null
+  prompt: string
+  negative: string
+  steps: number
+  seed: number
+  cfg: number
+  snofsModel: number
+  snofsClip: number
+  consistencyModel: number
+  consistencyClip: number
+  megapixels: number
+}
+
+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 === '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.')
+  }
+
+  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 () => {
+      emitChainJob(job, {
+        type: 'status',
+        message: params.mode === 'compose' ? 'Uploading stills A and B to Beast...' : 'Uploading still A to Beast...',
+        progress: 8
+      })
+      const uploaded = await uploadImage(params.image, job.id)
+      const uploadedRef = params.reference
+        ? await uploadImage({
+          ...params.reference,
+          filename: `ref_${params.reference.filename || 'image_b.png'}`
+        }, job.id)
+        : null
+      if (job.status === 'cancelled') throw new Error('Job interrupted.')
+
+      emitChainJob(job, {
+        type: 'status',
+        message: params.mode === 'compose'
+          ? `Queueing Klein v2 compose (${params.task}) on Beast...`
+          : 'Queueing Klein v2 edit on Beast...',
+        progress: 12
+      })
+      await ensureComfyLoraNames('image')
+      await assertImageScaleToTotalPixelsNode()
+      const built = buildImageV2Workflow({
+        mode: params.mode,
+        task: params.task,
+        prompt: params.prompt,
+        negative: params.negative,
+        imageAName: uploaded.name,
+        imageBName: uploadedRef?.name,
+        snofsModel: params.snofsModel,
+        snofsClip: params.snofsClip,
+        consistencyModel: params.consistencyModel,
+        consistencyClip: params.consistencyClip,
+        steps: params.steps,
+        cfg: params.cfg,
+        seed: params.seed,
+        megapixels: params.megapixels,
+        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 ? [uploadedRef.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)
+  }
+}