Keep image jobs on the studio queue and confirm each job or library delete. Label stitch-frame slots when thumbs are hidden, and open output stills in a lightbox.
Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Towsty
@@ -0,0 +1,213 @@
|
||||
import { createJob, emitJob, type Job } from '~/server/utils/jobs'
|
||||
import { ensureComfyReady } from '~/server/utils/comfyLifecycle'
|
||||
import { getComfyHost, uploadImage, queuePrompt, purgeComfyArtifacts } from '~/server/utils/comfy'
|
||||
import { withImageComfyHost, waitForImageEdit, downloadEditedImage } from '~/server/utils/imageComfy'
|
||||
import { buildEditWorkflow } from '~/server/utils/imageWorkflow'
|
||||
import { imageDimensions } from '~/server/utils/resolution'
|
||||
import { emitChainJob } from '~/server/utils/watch'
|
||||
import { saveStill, stillChainName } from '~/server/utils/library'
|
||||
|
||||
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
|
||||
}
|
||||
|
||||
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
|
||||
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 () => {
|
||||
let current = params.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 prompt = prompts[index]
|
||||
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 = prompt
|
||||
library.seed = seed
|
||||
library.chainIndex = index
|
||||
library.chainStep = index + 1
|
||||
library.chainLabel = chainTotal > 1 ? `Pass ${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
|
||||
})
|
||||
const graph = buildEditWorkflow({
|
||||
prompt,
|
||||
negative: params.negative,
|
||||
imageName: uploaded.name,
|
||||
referenceImageName: uploadedRef?.name,
|
||||
steps: params.steps,
|
||||
seed,
|
||||
cfg: params.cfg,
|
||||
filenamePrefix: `aigen_edit_${job.id.slice(0, 8)}_p${index + 1}`
|
||||
})
|
||||
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,
|
||||
familyId: library.familyId,
|
||||
parentStillId,
|
||||
chainIndex: index
|
||||
})
|
||||
job.stillId = still?.id
|
||||
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: `Pass ${index + 1} saved`,
|
||||
progress: 100,
|
||||
stillId: still?.id,
|
||||
hideThumbnail: job.hideThumbnail,
|
||||
folderLocked: library.folderLocked
|
||||
})
|
||||
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 ? '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
|
||||
}
|
||||
Reference in New Issue
Block a user