Show Krea in Image v2 sampler progress instead of Flux.2 Klein.

Co-authored-by: Cursor <cursoragent@cursor.com>

server/utils/imageChainV2.ts

@@ -2,7 +2,7 @@ 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, resolveKreaGenerateAssets } from '~/server/utils/imageWorkflowV2'
+import { buildImageV2Workflow, IMAGE_V2_NODE_LABELS, resolveKreaGenerateAssets } from '~/server/utils/imageWorkflowV2'
 import { ensureComfyLoraNames } from '~/server/utils/loras'
 import { imageDimensions } from '~/server/utils/resolution'
 import { emitChainJob } from '~/server/utils/watch'
@@ -148,6 +148,11 @@ export async function runEditV2(job: Job, params: EditV2RunParams) {
         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',
@@ -209,7 +214,7 @@ export async function runEditV2(job: Job, params: EditV2RunParams) {
       job.status = 'complete'
       emitChainJob(job, {
         type: 'complete',
-        message: 'Klein v2 finished on Beast',
+        message: `${engineLabel} v2 finished on Beast`,
         progress: 100,
         stillId: still?.id,
         filename: output.filename,

server/utils/imageComfy.ts

@@ -244,10 +244,13 @@ export async function waitForImageEdit(opts: {
   promptId: string
   clientId: string
   timeoutMs?: number
+  engineLabel?: string
+  nodeLabel?: (node: string) => string | undefined
   onProgress?: (event: { message: string; progress: number; step?: number; maxStep?: number; node?: string | null }) => void
   isCancelled?: () => boolean
 }) {
   const timeoutMs = opts.timeoutMs || 180_000
+  const engineLabel = String(opts.engineLabel || 'Flux.2 Klein').trim() || 'Flux.2 Klein'
   const started = Date.now()
   let settled = false
   let lastError: string | null = null
@@ -281,19 +284,21 @@ export async function waitForImageEdit(opts: {
       if (payload.type === 'progress' && payload.data) {
         const max = Math.max(1, Number(payload.data.max || 4))
         const value = Number(payload.data.value || 0)
+        const node = payload.data.node || null
         opts.onProgress?.({
-          message: 'Sampling Flux.2 Klein...',
+          message: `Sampling ${engineLabel}...`,
           progress: Math.min(90, 20 + Math.round((value / max) * 60)),
           step: value,
           maxStep: max,
-          node: payload.data.node || null
+          node
         })
       }
       if (payload.type === 'executing' && payload.data?.node) {
+        const node = payload.data.node
         opts.onProgress?.({
-          message: 'Running Flux.2 Klein...',
+          message: opts.nodeLabel?.(node) || `Running ${engineLabel}...`,
-          progress: payload.data.node === '9' || payload.data.node === '94' || payload.data.node === '21' ? 92 : 30,
+          progress: node === '9' || node === '94' || node === '21' ? 92 : 30,
-          node: payload.data.node
+          node
         })
       }
       if (payload.type === 'execution_error') {
@@ -327,7 +332,7 @@ export async function waitForImageEdit(opts: {
   if (result) return result
   throw createError({
     statusCode: 502,
-    statusMessage: lastError || 'Image edit timed out waiting for Flux.2 Klein to finish'
+    statusMessage: lastError || `Image edit timed out waiting for ${engineLabel} to finish`
   })
 }
 

server/utils/imageWorkflowV2.ts

@@ -654,7 +654,7 @@ export const IMAGE_V2_NODE_LABELS: Record<string, string> = {
   '33': 'Applying mask',
   '14': 'Building empty latent',
   '15': 'Sampling',
-  '19': 'Sampling Klein v2',
+  '19': 'Sampling',
   '20': 'Decoding still',
   '21': 'Saving still'
 }