409 lines
16 KiB
TypeScript
409 lines
16 KiB
TypeScript
import { normalizeImageIterations, resolveImageIteration } from '~/utils/imageIterations'
|
|
import { createJob, emitJob, type Job } from '~/server/utils/jobs'
|
|
import { ensureComfyReady } from '~/server/utils/comfyLifecycle'
|
|
import { assertImageScaleToTotalPixelsNode, comfyImagePrefix, getComfyHost, uploadImage, queuePrompt } from '~/server/utils/comfy'
|
|
import { withImageComfyHost, waitForImageEdit, downloadEditedImage, purgeImageComfyArtifacts } 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 { restorePromptParts, wrappedPromptForComfy } from '~/utils/promptParts'
|
|
import { isXaigenStudio } from '~/utils/videoModels'
|
|
import {
|
|
imageV2GraphMode,
|
|
imageV2StackSpecials,
|
|
mergeLegacyImageV2Loras,
|
|
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?: import('~/utils/imageIterations').ImageIteration[]
|
|
passMode?: 'batch' | 'chain'
|
|
}
|
|
|
|
export async function runEditV2(job: Job, params: EditV2RunParams) {
|
|
const library = job.library
|
|
if (!library) throw new Error('Edit job is missing library metadata')
|
|
const graphMode = imageV2GraphMode(params.mode, {
|
|
hasImage: Boolean(params.image),
|
|
hasReference: Boolean(params.reference)
|
|
})
|
|
if (graphMode === 'refine' && !params.mask) {
|
|
throw new Error('Refine requires a mask. Refusing to fall back to Edit.')
|
|
}
|
|
if (graphMode === 'compose' && !params.reference) {
|
|
throw new Error('Compose requires image B. Refusing to fall back to one-image generation.')
|
|
}
|
|
if (graphMode === 'edit' && params.reference && params.mode !== 'iterate') {
|
|
throw new Error('Edit mode takes one image. Use Compose for two stills.')
|
|
}
|
|
if (graphMode !== 'generate' && !params.image) {
|
|
throw new Error('This v2 mode requires still A.')
|
|
}
|
|
|
|
const canChain = params.mode === 'iterate' || params.mode === 'edit' || params.mode === 'compose'
|
|
const extraPasses = canChain
|
|
? normalizeImageIterations(params.passes)
|
|
: []
|
|
const prompts = [
|
|
library.promptMid || restorePromptParts(params.prompt, {
|
|
pre: library.promptPre,
|
|
prompt: params.prompt,
|
|
post: library.promptPost
|
|
}).prompt,
|
|
...extraPasses.map(item => item.prompt || library.promptMid || library.prompt || params.prompt)
|
|
]
|
|
const chainTotal = prompts.length
|
|
// Text-only iterate is always independent generations. With a still, honor batch vs chain.
|
|
const batch = graphMode === 'generate' || params.passMode !== 'chain'
|
|
const unit = params.mode === 'iterate' ? 'Iteration' : (batch ? 'Batch' : 'Pass')
|
|
library.chainTotal = chainTotal
|
|
library.familyId = chainTotal > 1 ? (library.familyId || crypto.randomUUID()) : library.familyId
|
|
|
|
try {
|
|
// Keep the studio slot across all iterations, including upload/save gaps.
|
|
library.chainContinuing = true
|
|
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 = graphMode === '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 startPadded = false
|
|
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 && params.megapixels > 0) await assertImageScaleToTotalPixelsNode()
|
|
const krea = engine === 'krea' ? await resolveKreaGenerateAssets() : null
|
|
if (!generate && graphMode === '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 pass = index > 0 ? extraPasses[index - 1] : undefined
|
|
const { steps, cfg, seed, strength } = resolveImageIteration(params, pass)
|
|
const turbo = pass?.steps !== undefined || pass?.cfg !== undefined ? false : params.turbo === true
|
|
const appliedStack = mergeLegacyImageV2Loras(pass?.loraStack ?? params.loraStack ?? library.loraStack, {
|
|
engine,
|
|
mode: graphMode,
|
|
xaigen: isXaigenStudio(),
|
|
snofsModel: pass?.loraStack !== undefined ? 0 : params.snofsModel,
|
|
snofsClip: pass?.loraStack !== undefined ? 0 : params.snofsClip,
|
|
consistencyModel: pass?.loraStack !== undefined ? 0 : params.consistencyModel,
|
|
consistencyClip: pass?.loraStack !== undefined ? 0 : params.consistencyClip,
|
|
kreaConcept: krea?.conceptLora
|
|
})
|
|
const appliedSpecials = imageV2StackSpecials(appliedStack, engine)
|
|
const snofsModel = appliedSpecials.snofs?.strengthModel ?? (pass?.loraStack !== undefined ? 0 : params.snofsModel)
|
|
const snofsClip = appliedSpecials.snofs?.strengthClip ?? (pass?.loraStack !== undefined ? 0 : params.snofsClip)
|
|
const consistencyModel = appliedSpecials.consistency?.strengthModel ?? (pass?.loraStack !== undefined ? 0 : params.consistencyModel)
|
|
const consistencyClip = appliedSpecials.consistency?.strengthClip ?? (pass?.loraStack !== undefined ? 0 : params.consistencyClip)
|
|
|
|
library.steps = steps
|
|
library.cfg = cfg
|
|
job.maxStep = steps
|
|
library.prompt = prompts[index]
|
|
library.seed = seed
|
|
library.chainIndex = index
|
|
library.chainStep = index + 1
|
|
library.chainLabel = chainTotal > 1 ? `${unit} ${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} iterate generate on Beast...`
|
|
: graphMode === 'refine'
|
|
? 'Uploading canvas and mask to Beast...'
|
|
: graphMode === '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: graphMode === '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
|
|
startPadded = filled.padded
|
|
prompt = composeFillPrompt(prompt, padded)
|
|
if (padded) {
|
|
emitChainJob(job, {
|
|
type: 'status',
|
|
message: `Filling to ${filled.aspect} without stretching...`,
|
|
progress: 7
|
|
})
|
|
}
|
|
} else if (batch && startPadded) {
|
|
prompt = composeFillPrompt(prompt, true)
|
|
}
|
|
|
|
if (index === 0 && graphMode === '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 === 'iterate'
|
|
? `Queueing ${engineLabel} iterate (${graphMode}) on Beast...`
|
|
: graphMode === 'refine'
|
|
? `Queueing ${engineLabel} v2 refine on Beast...`
|
|
: graphMode === 'compose'
|
|
? `Queueing ${engineLabel} v2 compose (${params.task}) on Beast...`
|
|
: `Queueing ${engineLabel} v2 edit on Beast...`,
|
|
progress: 12
|
|
})
|
|
}
|
|
|
|
const built = buildImageV2Workflow({
|
|
mode: graphMode,
|
|
engine,
|
|
task: generate ? 't2i' : params.task,
|
|
prompt: wrappedPromptForComfy(library, prompt, index === 0 ? params.prompt : prompt),
|
|
negative: params.negative,
|
|
imageAName: uploaded?.name,
|
|
imageBName: uploadedRef?.name,
|
|
maskName: uploadedMask?.name,
|
|
strength,
|
|
snofsModel,
|
|
snofsClip,
|
|
consistencyModel,
|
|
consistencyClip,
|
|
steps,
|
|
cfg,
|
|
seed,
|
|
megapixels: params.megapixels,
|
|
width: params.width,
|
|
height: params.height,
|
|
aspect: params.aspect,
|
|
turbo,
|
|
filenamePrefix: chainTotal > 1 ? `${comfyImagePrefix()}_${comfyJobPrefix(job.id)}_p${index + 1}` : `${comfyImagePrefix()}_${comfyJobPrefix(job.id)}`,
|
|
kreaUnetName: krea?.unet,
|
|
kreaClipName: krea?.clip,
|
|
kreaVaeName: krea?.vae,
|
|
kreaConceptLora: krea?.conceptLora,
|
|
loraStack: appliedStack
|
|
})
|
|
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: 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 || steps,
|
|
node: event.node
|
|
})
|
|
},
|
|
isCancelled: () => job.status === 'cancelled'
|
|
})
|
|
|
|
emitChainJob(job, { type: 'status', message: `Saving ${engineLabel} v2 still...`, progress: 94 })
|
|
job.saving = true
|
|
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,
|
|
promptPre: library.promptPre,
|
|
promptPost: library.promptPost,
|
|
familyId: library.familyId,
|
|
parentStillId: batch ? (params.sourceStillId || library.stillId) : parentStillId,
|
|
chainIndex: index,
|
|
settings: {
|
|
kind: 'imagev2',
|
|
mode: params.mode,
|
|
engine,
|
|
task: generate ? 't2i' : params.task,
|
|
negative: params.negative,
|
|
snofsModel,
|
|
snofsClip,
|
|
consistencyModel,
|
|
consistencyClip,
|
|
steps,
|
|
cfg,
|
|
megapixels: generate || !(params.megapixels > 0) ? undefined : params.megapixels,
|
|
scaleToTotalPixels: generate ? undefined : params.megapixels > 0,
|
|
turbo,
|
|
strength: generate ? undefined : strength,
|
|
width: generate ? params.width : undefined,
|
|
height: generate ? params.height : undefined,
|
|
aspect: params.aspect || 'auto',
|
|
seed,
|
|
sourceStillId: params.sourceStillId || library.stillId,
|
|
referenceStillId: params.referenceStillId,
|
|
...persistLoraFields(appliedStack)
|
|
}
|
|
})
|
|
job.stillId = still?.id
|
|
if (!batch) parentStillId = still?.id
|
|
await purgeImageComfyArtifacts({
|
|
output: { filename: output.filename, subfolder: output.subfolder, type: output.type },
|
|
inputName: uploaded?.name,
|
|
inputSubfolder: uploaded?.subfolder,
|
|
extraInputNames: last
|
|
? [uploadedRef?.name, uploadedMask?.name].filter((name): name is string => Boolean(name))
|
|
: [],
|
|
promptId: job.promptId
|
|
})
|
|
|
|
job.saving = false
|
|
if (!last) {
|
|
emitChainJob(job, {
|
|
type: 'checkpoint',
|
|
message: `${unit} ${index + 1} of ${chainTotal} saved`,
|
|
progress: 100,
|
|
stillId: still?.id,
|
|
hideThumbnail: job.hideThumbnail,
|
|
folderLocked: library.folderLocked
|
|
})
|
|
if (!batch) {
|
|
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} ${params.mode === 'iterate' ? 'iterate' : `v2 ${batch ? 'batch' : '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 {
|
|
job.saving = false
|
|
library.chainContinuing = false
|
|
const { onLiveVideoSettled } = await import('~/server/utils/studioQueue')
|
|
await onLiveVideoSettled(job)
|
|
}
|
|
}
|