Files
aigen/server/api/extend.post.ts
T

278 lines
9.0 KiB
TypeScript

import { copyFileSync, existsSync, readFileSync } from 'node:fs'
import { join } from 'node:path'
import { composeShotPrompt } from '~/utils/globalLocks'
import { parsePostedLoraStack, listStudioLoras, persistLoraFields } from '~/server/utils/loras'
import { readLoraStack } from '~/utils/loras'
import { clampVideoCfg } from '~/utils/generationPresets'
import { defaultVideoSteps, isLtxWorkflow, LTX_DISABLED_MESSAGE, ltxWorkflowEnabled, parseVideoWorkflow } from '~/utils/videoModels'
const SAMPLERS = new Set(['res_multistep', 'euler', 'dpmpp_2m'])
const SCHEDULERS = new Set(['simple', 'ddim_uniform', 'sgm_uniform'])
function parseExtendDuration(raw: unknown) {
const seconds = Number(raw)
if (!Number.isFinite(seconds)) return 5
return Math.min(120, Math.max(0.5, Math.round(seconds * 10) / 10))
}
function frameLength(seconds: number, fps: number) {
return Math.max(5, Math.floor(seconds * fps))
}
function parseFps(raw: unknown) {
const fps = Number(raw)
return fps === 12 || fps === 30 || fps === 24 ? fps : 24
}
function parseCfg(raw: unknown, turbo: boolean) {
return clampVideoCfg(raw, turbo ? 1.5 : 4)
}
function parseSampler(raw: unknown) {
const value = String(raw || '')
return SAMPLERS.has(value) ? value : 'res_multistep'
}
function parseScheduler(raw: unknown) {
const value = String(raw || '')
return SCHEDULERS.has(value) ? value : 'simple'
}
export default defineEventHandler(async (event) => {
const body = await readBody<{ clipId?: string; prompt?: string; duration?: number; lora?: unknown; loraStack?: unknown }>(event).catch(() => ({}))
const clipId = String(body?.clipId || '').trim()
const prompt = String(body?.prompt || '').trim()
if (!clipId) {
throw createError({ statusCode: 400, statusMessage: 'A clip is required to extend' })
}
if (!prompt) {
throw createError({ statusCode: 400, statusMessage: 'An extension prompt is required' })
}
const { owner } = assertLibraryOwner(event)
const source = getClip(owner, clipId)
assertFolderAccess(event, source.folderId)
const sourcePath = clipVideoPath(owner, source.id)
if (!existsSync(sourcePath)) {
throw createError({ statusCode: 404, statusMessage: 'Source video is missing' })
}
const destFolder = publicLibrary(event).folders.find(folder => folder.id === source.folderId)
const folderLocked = Boolean(destFolder?.protected && !destFolder.unlocked)
const durationSeconds = parseExtendDuration(body?.duration)
const composedPrompt = composeShotPrompt({
globalLocks: source.globalLocks,
prompt,
shotIndex: (source.chainIndex || 0) + 1,
familyRefs: source.permanenceRefs
})
const workflow = parseVideoWorkflow(source.workflow)
if (isLtxWorkflow(workflow) && !ltxWorkflowEnabled()) {
throw createError({ statusCode: 400, statusMessage: LTX_DISABLED_MESSAGE })
}
const turbo = source.turbo !== false
const fps = parseFps(source.fps)
const cfg = parseCfg(source.cfg, turbo)
const samplerName = parseSampler(source.samplerName)
const scheduler = parseScheduler(source.scheduler)
const steps = Number(source.steps || defaultVideoSteps(turbo, workflow))
const seed = Math.floor(Math.random() * 2_147_483_647)
await listStudioLoras()
const postedStack = parsePostedLoraStack(body?.loraStack ?? body?.lora, 'video')
const loraStack = postedStack.length ? postedStack : parsePostedLoraStack(readLoraStack(source), 'video')
const loraFields = persistLoraFields(loraStack)
const length = frameLength(durationSeconds, fps)
const width = source.width || 1344
const height = source.height || 768
const clipName = nextClipPartName(clipTitle(source))
const job = createJob()
job.kind = 'video'
job.maxStep = steps
job.hideThumbnail = source.hideThumbnail
const tmpDir = extendTempDir(owner, job.id)
const part1Path = join(tmpDir, 'part1.mp4')
const framePath = join(tmpDir, 'last_frame.png')
job.library = {
ownerKey: owner,
folderId: source.folderId,
hideThumbnail: source.hideThumbnail,
folderLocked,
name: clipName,
prompt: composedPrompt,
aspect: source.aspect || '16:9',
width,
height,
steps,
turbo,
seed,
cfg,
fps,
samplerName,
scheduler,
duration: durationSeconds,
extendTmpDir: tmpDir,
extendPart1Path: part1Path,
extendSourceClipId: source.id,
familyId: source.familyId || crypto.randomUUID(),
parentClipId: source.id,
chainIndex: (source.chainIndex || 0) + 1,
workflow,
globalLocks: source.globalLocks,
permanenceRefs: source.permanenceRefs,
...loraFields
}
emitJob(job, { type: 'status', message: 'Extracting last frame & stitching extension...', progress: 2 })
void runExtend(job, {
sourcePath,
part1Path,
framePath,
prompt: composedPrompt,
width,
height,
steps,
seed,
turbo,
length,
cfg,
fps,
samplerName,
scheduler,
...loraFields
}).catch((error) => {
removeExtendTemp(job.library?.extendTmpDir)
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,
clientId: job.clientId,
seed,
width,
height,
steps,
duration: durationSeconds,
hideThumbnail: source.hideThumbnail,
folderLocked
}
})
async function runExtend(
job: ReturnType<typeof createJob>,
params: {
sourcePath: string
part1Path: string
framePath: string
prompt: string
width: number
height: number
steps: number
seed: number
turbo: boolean
length: number
cfg: number
fps: number
samplerName: string
scheduler: string
loraName?: string
loraStack?: import('~/utils/loras').LoraStackItem[]
}
) {
copyFileSync(params.sourcePath, params.part1Path)
await extractLastFrame(params.sourcePath, params.framePath)
const frame = readFileSync(params.framePath)
if (!frame.length || frame.length < 64) {
throw new Error('Could not extract the last frame of the clip being extended')
}
if (job.library) job.library.thumb = undefined
const sound = await probeHasAudio(params.part1Path)
await ensureComfyReady((status) => {
emitJob(job, {
type: status.state === 'busy' ? 'busy' : 'status',
message: status.message,
progress: status.state === 'online' ? Math.max(job.progress, 3) : Math.max(job.progress, 2),
busy: status.state === 'busy',
queueRunning: status.queueRunning,
queuePending: status.queuePending
})
})
const done = watchComfyJob(job)
job.status = 'uploading'
emitJob(job, { type: 'status', message: 'Uploading last frame to ComfyUI...', progress: 4 })
const uploaded = await uploadImage({ filename: 'last_frame.png', data: frame, type: 'image/png' }, job.id)
if (job.library) {
job.library.imageName = uploaded.name
job.library.imageSubfolder = uploaded.subfolder
}
emitJob(job, { type: 'status', message: 'Queueing extension on MiniMax H3...', progress: 6 })
await waitForComfySocket(job, 4000)
const graph = buildWorkflow({
prompt: params.prompt,
imageName: uploaded.name,
width: params.width,
height: params.height,
steps: params.steps,
seed: params.seed,
turbo: params.turbo,
length: params.length,
cfg: params.cfg,
fps: params.fps,
samplerName: params.samplerName,
scheduler: params.scheduler,
filenamePrefix: comfyFilenamePrefix(),
sound,
workflow: job.library?.workflow,
duration: job.library?.duration,
loraStack: params.loraStack || job.library?.loraStack || params.loraName || job.library?.loraName
})
const queued = await queuePrompt(graph, job.clientId)
job.promptId = queued.prompt_id
job.status = 'running'
if (job.library && job.promptId) {
writePendingJob({
jobId: job.id,
promptId: job.promptId,
clientId: job.clientId,
ownerKey: job.library.ownerKey,
folderId: job.library.folderId,
hideThumbnail: job.library.hideThumbnail,
folderLocked: job.library.folderLocked,
name: job.library.name,
prompt: job.library.prompt,
aspect: job.library.aspect,
width: job.library.width,
height: job.library.height,
steps: job.library.steps,
turbo: job.library.turbo,
seed: job.library.seed,
startedAt: job.startedAt,
imageName: job.library.imageName,
imageSubfolder: job.library.imageSubfolder,
extendTmpDir: job.library.extendTmpDir,
extendPart1Path: job.library.extendPart1Path,
familyId: job.library.familyId,
parentClipId: job.library.parentClipId,
chainIndex: job.library.chainIndex,
stillId: job.library.stillId,
sound: job.library.sound,
workflow: job.library.workflow,
duration: job.library.duration,
globalLocks: job.library.globalLocks,
permanenceRefs: job.library.permanenceRefs,
loraName: job.library.loraName,
loraStack: job.library.loraStack
})
}
emitJob(job, { type: 'status', message: 'Job queued on ComfyUI', progress: 8 })
await done
}