260 lines
9.7 KiB
TypeScript
260 lines
9.7 KiB
TypeScript
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 { buildEditWorkflow } from '~/server/utils/imageWorkflow'
|
|
import { ensureComfyLoraNames } 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 { joinPromptParts } from '~/utils/promptParts'
|
|
|
|
export type EditImageFile = { filename: string; data: Buffer; type?: string }
|
|
|
|
export type EditRunParams = {
|
|
image: EditImageFile
|
|
reference: EditImageFile | null
|
|
prompt: string
|
|
passes: { prompt: string }[]
|
|
negative: string
|
|
steps: number
|
|
seed: number
|
|
cfg: number
|
|
aspect?: string
|
|
loraName?: string
|
|
loraStack?: import('~/utils/loras').LoraStackItem[]
|
|
scaleToTotalPixels?: boolean
|
|
scaleMegapixels?: number
|
|
passMode?: 'batch' | 'chain'
|
|
}
|
|
|
|
export async function runEdit(job: Job, params: EditRunParams) {
|
|
const library = job.library
|
|
if (!library) throw new Error('Edit job is missing library metadata')
|
|
const prompts = [params.prompt, ...params.passes.map(item => item.prompt)]
|
|
const chainTotal = prompts.length
|
|
const batch = params.passMode !== 'chain'
|
|
const unit = batch ? 'Batch' : 'Pass'
|
|
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 { applyImageFill } = await import('~/server/utils/imagePad')
|
|
const { composeFillPrompt } = await import('~/utils/imageAspect')
|
|
const filled = await applyImageFill({ image: params.image, aspect: params.aspect || 'auto' })
|
|
if (filled.padded) {
|
|
emitChainJob(job, {
|
|
type: 'status',
|
|
message: `Filling to ${filled.aspect} without stretching...`,
|
|
progress: 7
|
|
})
|
|
}
|
|
let current = filled.image
|
|
let parentStillId: string | undefined
|
|
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 mid = (batch || index === 0) ? composeFillPrompt(prompts[index], filled.padded) : prompts[index]
|
|
const prompt = joinPromptParts(library.promptPre, mid, library.promptPost)
|
|
const last = index === prompts.length - 1
|
|
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 })
|
|
}
|
|
const seed = index === 0 ? params.seed : Math.floor(Math.random() * 2_147_483_647)
|
|
library.prompt = mid
|
|
library.seed = seed
|
|
library.chainIndex = index
|
|
library.chainStep = index + 1
|
|
library.chainLabel = chainTotal > 1 ? `${unit} ${index + 1}` : undefined
|
|
const reference = index === 0 ? params.reference : null
|
|
const dual = Boolean(reference)
|
|
const passName = stillChainName(library.name || '', index)
|
|
|
|
emitChainJob(job, {
|
|
type: 'status',
|
|
message: dual ? 'Uploading both stills to Beast...' : 'Uploading still to Beast...',
|
|
progress: 8
|
|
})
|
|
const uploaded = await uploadImage(current, job.id)
|
|
const uploadedRef = reference
|
|
? await uploadImage({
|
|
...reference,
|
|
filename: `ref_${reference.filename || 'image2.png'}`
|
|
}, job.id)
|
|
: null
|
|
if (job.status === 'cancelled') throw new Error('Job interrupted.')
|
|
|
|
emitChainJob(job, {
|
|
type: 'status',
|
|
message: dual
|
|
? 'Queueing Flux.2 Klein two-image edit on Beast...'
|
|
: 'Queueing Flux.2 Klein edit on Beast...',
|
|
progress: 12
|
|
})
|
|
await ensureComfyLoraNames('image')
|
|
if (params.scaleToTotalPixels) await assertImageScaleToTotalPixelsNode()
|
|
const graph = buildEditWorkflow({
|
|
prompt,
|
|
negative: params.negative,
|
|
imageName: uploaded.name,
|
|
referenceImageName: uploadedRef?.name,
|
|
steps: params.steps,
|
|
seed,
|
|
cfg: params.cfg,
|
|
filenamePrefix: `${comfyJobPrefix(job.id)}_p${index + 1}`,
|
|
loraStack: params.loraStack || params.loraName,
|
|
scaleToTotalPixels: params.scaleToTotalPixels,
|
|
scaleMegapixels: params.scaleMegapixels
|
|
})
|
|
const queued = await queuePrompt(graph, job.clientId)
|
|
job.promptId = queued.prompt_id
|
|
job.status = 'running'
|
|
emitChainJob(job, {
|
|
type: 'status',
|
|
message: dual ? 'Applying image 2 onto image 1...' : 'Editing still on Beast...',
|
|
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 edited 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: mid,
|
|
promptPre: library.promptPre,
|
|
promptPost: library.promptPost,
|
|
familyId: library.familyId,
|
|
parentStillId: batch ? library.stillId : parentStillId,
|
|
chainIndex: index,
|
|
settings: {
|
|
kind: 'image',
|
|
steps: params.steps,
|
|
cfg: params.cfg,
|
|
negative: params.negative,
|
|
seed,
|
|
scaleToTotalPixels: params.scaleToTotalPixels === true,
|
|
scaleMegapixels: params.scaleMegapixels,
|
|
aspect: params.aspect || 'auto',
|
|
loraName: params.loraName,
|
|
loraStack: params.loraStack,
|
|
sourceStillId: library.stillId
|
|
}
|
|
})
|
|
job.stillId = still?.id
|
|
if (!batch) parentStillId = 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
|
|
})
|
|
|
|
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
|
|
? (batch ? 'Edit batch finished on Beast' : 'Edit chain finished on Beast')
|
|
: 'Edit 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)
|
|
}
|
|
}
|
|
|
|
export function createEditLiveJob(params: {
|
|
steps: number
|
|
hideThumbnail: boolean
|
|
library: NonNullable<Job['library']>
|
|
}) {
|
|
const job = createJob()
|
|
job.kind = 'edit'
|
|
job.maxStep = params.steps
|
|
job.hideThumbnail = params.hideThumbnail
|
|
job.library = params.library
|
|
return job
|
|
}
|