diff --git a/server/utils/imageChainV2.ts b/server/utils/imageChainV2.ts index 2afcb04..f705c64 100644 --- a/server/utils/imageChainV2.ts +++ b/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, diff --git a/server/utils/imageComfy.ts b/server/utils/imageComfy.ts index 9334d26..ba3073d 100644 --- a/server/utils/imageComfy.ts +++ b/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...', - progress: payload.data.node === '9' || payload.data.node === '94' || payload.data.node === '21' ? 92 : 30, - node: payload.data.node + message: opts.nodeLabel?.(node) || `Running ${engineLabel}...`, + progress: node === '9' || node === '94' || node === '21' ? 92 : 30, + 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` }) } diff --git a/server/utils/imageWorkflowV2.ts b/server/utils/imageWorkflowV2.ts index bf81e09..351ac30 100644 --- a/server/utils/imageWorkflowV2.ts +++ b/server/utils/imageWorkflowV2.ts @@ -654,7 +654,7 @@ export const IMAGE_V2_NODE_LABELS: Record = { '33': 'Applying mask', '14': 'Building empty latent', '15': 'Sampling', - '19': 'Sampling Klein v2', + '19': 'Sampling', '20': 'Decoding still', '21': 'Saving still' }