348 lines
11 KiB
TypeScript
348 lines
11 KiB
TypeScript
import { ensureComfyReady } from '~/server/utils/comfyLifecycle'
|
|
import { getComfyHost, comfyConfigured, 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'
|
|
|
|
function parseSteps(raw: unknown) {
|
|
const value = Number(raw)
|
|
if (!Number.isFinite(value)) return 20
|
|
return Math.min(40, Math.max(1, Math.round(value)))
|
|
}
|
|
|
|
function parseCfg(raw: unknown) {
|
|
const value = Number(raw)
|
|
if (!Number.isFinite(value)) return 1
|
|
return Math.min(8, Math.max(0.5, Math.round(value * 10) / 10))
|
|
}
|
|
|
|
function parsePasses(raw: string | undefined) {
|
|
if (!raw) return [] as { prompt: string }[]
|
|
try {
|
|
const parsed = JSON.parse(raw)
|
|
if (!Array.isArray(parsed)) return []
|
|
return parsed
|
|
.map((item: { prompt?: unknown }) => ({ prompt: String(item?.prompt || '').trim() }))
|
|
.filter(item => item.prompt)
|
|
} catch {
|
|
return []
|
|
}
|
|
}
|
|
|
|
export default defineEventHandler(async (event) => {
|
|
const form = await readMultipartFormData(event).catch(() => null)
|
|
const fields: Record<string, string> = {}
|
|
let image: { filename: string; data: Buffer; type?: string } | null = null
|
|
let reference: { filename: string; data: Buffer; type?: string } | null = null
|
|
for (const part of form || []) {
|
|
if (part.name === 'image' && part.filename && part.data?.length) {
|
|
image = { filename: part.filename, data: part.data, type: part.type }
|
|
} else if ((part.name === 'image2' || part.name === 'refImage') && part.filename && part.data?.length) {
|
|
reference = { filename: part.filename, data: part.data, type: part.type }
|
|
} else if (part.name && part.data) {
|
|
fields[part.name] = part.data.toString('utf8')
|
|
}
|
|
}
|
|
|
|
const prompt = (fields.prompt || '').trim()
|
|
const extraPasses = parsePasses(fields.passes)
|
|
if (!image) {
|
|
throw createError({ statusCode: 400, statusMessage: 'An input image is required' })
|
|
}
|
|
if (!prompt) {
|
|
throw createError({ statusCode: 400, statusMessage: 'An edit prompt is required' })
|
|
}
|
|
if (extraPasses.some(item => !item.prompt)) {
|
|
throw createError({ statusCode: 400, statusMessage: 'Each extra pass needs a prompt' })
|
|
}
|
|
if (!comfyConfigured()) {
|
|
throw createError({
|
|
statusCode: 503,
|
|
statusMessage: 'Beast ComfyUI is not configured. Set COMFY_HOST.'
|
|
})
|
|
}
|
|
|
|
const ownerKey = libraryOwnerKey(event)
|
|
const library = publicLibrary(event)
|
|
const folderId = library.folders.some(folder => folder.id === fields.folderId)
|
|
? fields.folderId
|
|
: library.folders[0]?.id
|
|
if (!folderId) {
|
|
throw createError({ statusCode: 400, statusMessage: 'Create a library folder before editing' })
|
|
}
|
|
assertFolderExists(event, folderId)
|
|
|
|
const destFolder = library.folders.find(folder => folder.id === folderId)
|
|
const folderLocked = Boolean(destFolder?.protected && !destFolder.unlocked)
|
|
const hideInput = fields.hideInput === 'true'
|
|
const hideThumbnail = fields.hideThumbnail === 'true'
|
|
const steps = parseSteps(fields.steps)
|
|
const cfg = parseCfg(fields.cfg)
|
|
const seed = fields.seed && fields.seed !== 'random'
|
|
? Number(fields.seed)
|
|
: Math.floor(Math.random() * 2_147_483_647)
|
|
const size = imageDimensions(image.data)
|
|
|
|
const still = await rememberInputStill({
|
|
ownerKey,
|
|
folderId,
|
|
filename: image.filename,
|
|
data: image.data,
|
|
width: size?.width || 0,
|
|
height: size?.height || 0,
|
|
hideInput
|
|
})
|
|
if (reference) {
|
|
await rememberInputStill({
|
|
ownerKey,
|
|
folderId,
|
|
filename: reference.filename,
|
|
data: reference.data,
|
|
hideInput
|
|
})
|
|
}
|
|
|
|
const passes = extraPasses
|
|
const chainTotal = 1 + passes.length
|
|
const familyId = chainTotal > 1 ? crypto.randomUUID() : undefined
|
|
const job = createJob()
|
|
job.kind = 'edit'
|
|
job.maxStep = steps
|
|
job.hideThumbnail = hideThumbnail
|
|
job.library = {
|
|
ownerKey,
|
|
folderId,
|
|
hideThumbnail,
|
|
hideInput,
|
|
folderLocked,
|
|
name: (fields.name || '').trim().slice(0, 80),
|
|
prompt,
|
|
aspect: 'auto',
|
|
width: size?.width || 0,
|
|
height: size?.height || 0,
|
|
steps,
|
|
turbo: true,
|
|
seed,
|
|
cfg,
|
|
stillId: still?.id,
|
|
familyId,
|
|
chainIndex: 0,
|
|
chainStep: 1,
|
|
chainTotal,
|
|
chainLabel: chainTotal > 1 ? 'Pass 1' : undefined,
|
|
passes
|
|
}
|
|
emitJob(job, {
|
|
type: 'status',
|
|
message: 'Checking Beast ComfyUI...',
|
|
progress: 2,
|
|
chainStep: 1,
|
|
chainTotal,
|
|
chainLabel: job.library.chainLabel
|
|
})
|
|
|
|
void runEdit(job, {
|
|
image,
|
|
reference,
|
|
prompt,
|
|
passes,
|
|
negative: (fields.negative || '').trim(),
|
|
steps,
|
|
seed,
|
|
cfg
|
|
}).catch(async (error) => {
|
|
if (job.status === 'error' || job.status === 'cancelled') return
|
|
const message = error instanceof Error ? error.message : String(error)
|
|
job.status = 'error'
|
|
job.error = message
|
|
emitJob(job, { type: 'error', error: message, message })
|
|
})
|
|
|
|
return {
|
|
jobId: job.id,
|
|
seed,
|
|
steps,
|
|
hideThumbnail,
|
|
folderLocked,
|
|
dual: Boolean(reference),
|
|
chainTotal
|
|
}
|
|
})
|
|
|
|
async function runEdit(
|
|
job: ReturnType<typeof createJob>,
|
|
params: {
|
|
image: { filename: string; data: Buffer; type?: string }
|
|
reference: { filename: string; data: Buffer; type?: string } | null
|
|
prompt: string
|
|
passes: { prompt: string }[]
|
|
negative: string
|
|
steps: number
|
|
seed: number
|
|
cfg: number
|
|
}
|
|
) {
|
|
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
|
|
|
|
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'
|
|
})
|
|
})
|
|
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.')
|
|
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'
|
|
})
|
|
})
|
|
}
|
|
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'
|
|
}
|
|
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
|
|
})
|
|
}
|
|
})
|
|
}
|