324 lines
11 KiB
TypeScript
324 lines
11 KiB
TypeScript
import { copyFileSync, existsSync, readFileSync, writeFileSync } from 'node:fs'
|
|
import { join } from 'node:path'
|
|
import type { Job } from '~/server/utils/jobs'
|
|
import { pendingFromJob, remainingAfterCurrentShot, writePendingJob } from '~/server/utils/pending'
|
|
import { emitChainJob, waitForComfySocket, watchComfyJob } from '~/server/utils/watch'
|
|
import { comfyFilenamePrefix, queuePrompt, uploadImage } from '~/server/utils/comfy'
|
|
import { buildWorkflow } from '~/server/utils/workflow'
|
|
import { clipVideoPath, deleteRetryDraft, extendTempDir, stillPath } from '~/server/utils/library'
|
|
import { extractLastFrame, probeHasAudio } from '~/server/utils/ffmpeg'
|
|
import { ensureComfyReady } from '~/server/utils/comfyLifecycle'
|
|
|
|
export type ChainImage = { filename: string; data: Buffer; type?: string }
|
|
|
|
type VideoChainParams = {
|
|
prompt: string
|
|
image: ChainImage
|
|
width: number
|
|
height: number
|
|
steps: number
|
|
seed: number
|
|
turbo: boolean
|
|
length: number
|
|
sound: boolean
|
|
cfg: number
|
|
fps: number
|
|
samplerName: string
|
|
scheduler: string
|
|
extensions: { prompt: string; duration: number }[]
|
|
workflow: 'v1' | 'v2'
|
|
duration: number
|
|
useIdentityRefs: boolean
|
|
referenceImages: Array<ChainImage | null>
|
|
}
|
|
|
|
function sleep(ms: number) {
|
|
return new Promise(resolve => setTimeout(resolve, ms))
|
|
}
|
|
|
|
export function frameLength(seconds: number, fps: number) {
|
|
return Math.max(5, Math.floor(seconds * fps))
|
|
}
|
|
|
|
function assertJobActive(job: Job) {
|
|
if (job.status === 'cancelled') {
|
|
throw new Error('Job interrupted.')
|
|
}
|
|
if (job.status === 'error') {
|
|
throw new Error(job.error || 'Generation failed')
|
|
}
|
|
}
|
|
|
|
function loadStillImage(owner: string, stillId?: string, filename?: string): ChainImage | null {
|
|
if (!stillId) return null
|
|
const path = stillPath(owner, stillId)
|
|
if (!existsSync(path)) return null
|
|
return {
|
|
filename: filename || 'still.png',
|
|
data: readFileSync(path),
|
|
type: 'image/png'
|
|
}
|
|
}
|
|
|
|
function paramsFromJob(job: Job): VideoChainParams {
|
|
const library = job.library
|
|
if (!library) {
|
|
throw new Error('Cannot continue the shot chain: job metadata is missing')
|
|
}
|
|
const image = loadStillImage(library.ownerKey, library.stillId, library.stillFilename)
|
|
if (!image?.data.length) {
|
|
throw new Error('Cannot continue the shot chain: the start still is missing from the library')
|
|
}
|
|
const referenceImages: Array<ChainImage | null> = [null, null, null, null]
|
|
for (const [index, stillId] of (library.referenceStillIds || []).entries()) {
|
|
if (!stillId || index >= 4) continue
|
|
referenceImages[index] = loadStillImage(library.ownerKey, stillId, `identity-ref-${index + 2}.png`)
|
|
}
|
|
const fps = library.fps || 24
|
|
const duration = library.duration || 5
|
|
return {
|
|
prompt: library.prompt,
|
|
image,
|
|
width: library.width,
|
|
height: library.height,
|
|
steps: library.steps,
|
|
seed: library.seed,
|
|
turbo: library.turbo,
|
|
length: frameLength(duration, fps),
|
|
sound: library.sound !== false,
|
|
cfg: library.cfg || (library.turbo ? 1.5 : 4),
|
|
fps,
|
|
samplerName: library.samplerName || 'res_multistep',
|
|
scheduler: library.scheduler || 'simple',
|
|
extensions: library.extensions || [],
|
|
workflow: library.workflow || 'v1',
|
|
duration,
|
|
useIdentityRefs: library.useIdentityRefs === true,
|
|
referenceImages
|
|
}
|
|
}
|
|
|
|
export async function queueMiniMax(
|
|
job: Job,
|
|
params: Omit<VideoChainParams, 'extensions'> & { persist: boolean }
|
|
) {
|
|
assertJobActive(job)
|
|
job.socketReady = false
|
|
job.promptId = undefined
|
|
job.video = undefined
|
|
job.segmentBuffer = undefined
|
|
|
|
const done = watchComfyJob(job, { persist: params.persist })
|
|
job.status = 'uploading'
|
|
const chainIndex = job.library?.chainIndex || 0
|
|
const uploading = params.useIdentityRefs
|
|
? (chainIndex > 0 ? 'Uploading identity stills for next shot...' : 'Uploading image to ComfyUI...')
|
|
: (chainIndex > 0 ? 'Uploading last frame to ComfyUI...' : 'Uploading image to ComfyUI...')
|
|
const queueing = (job.library?.chainIndex || 0) > 0
|
|
? 'Queueing extension on MiniMax H3...'
|
|
: 'Queueing MiniMax H3 job...'
|
|
emitChainJob(job, { type: 'status', message: uploading, progress: 4 })
|
|
const uploaded = await uploadImage(params.image, job.id)
|
|
const referenceNames: string[] = ['', '', '', '']
|
|
if (params.useIdentityRefs) {
|
|
for (const [index, ref] of (params.referenceImages || []).entries()) {
|
|
if (!ref?.data?.length) continue
|
|
const next = await uploadImage({
|
|
...ref,
|
|
filename: `ref${index + 1}_${ref.filename || 'identity.png'}`
|
|
}, job.id)
|
|
referenceNames[index] = next.name
|
|
}
|
|
}
|
|
if (job.library) {
|
|
job.library.imageName = uploaded.name
|
|
job.library.imageSubfolder = uploaded.subfolder
|
|
job.library.referenceImageNames = referenceNames.filter(Boolean)
|
|
}
|
|
emitChainJob(job, { type: 'status', message: queueing, 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: params.sound,
|
|
workflow: params.workflow,
|
|
duration: params.duration,
|
|
useIdentityRefs: params.useIdentityRefs,
|
|
referenceImageNames: params.useIdentityRefs ? referenceNames : []
|
|
})
|
|
|
|
const queued = await queuePrompt(graph, job.clientId)
|
|
job.promptId = queued.prompt_id
|
|
job.status = 'running'
|
|
if (job.library?.draftId) {
|
|
await deleteRetryDraft(job.library.ownerKey, job.library.draftId).catch(() => null)
|
|
job.library.draftId = undefined
|
|
}
|
|
if (job.library && job.promptId) {
|
|
writePendingJob(pendingFromJob(job))
|
|
}
|
|
emitChainJob(job, { type: 'status', message: 'Job queued on ComfyUI', progress: 8 })
|
|
await done
|
|
assertJobActive(job)
|
|
if (!params.persist && !job.segmentBuffer?.length) {
|
|
throw new Error('Segment finished without a video')
|
|
}
|
|
}
|
|
|
|
function seedCurrentVideo(job: Job) {
|
|
const library = job.library
|
|
if (!library) {
|
|
throw new Error('Cannot continue the shot chain: job metadata is missing')
|
|
}
|
|
const tmpDir = library.extendTmpDir || extendTempDir(library.ownerKey, job.id)
|
|
library.extendTmpDir = tmpDir
|
|
const currentPath = join(tmpDir, 'current.mp4')
|
|
if (job.segmentBuffer?.length) {
|
|
writeFileSync(currentPath, job.segmentBuffer)
|
|
job.segmentBuffer = undefined
|
|
} else if (!existsSync(currentPath)) {
|
|
const clipId = job.clipId
|
|
if (!clipId) {
|
|
throw new Error('Cannot continue the shot chain: the previous clip is missing')
|
|
}
|
|
const source = clipVideoPath(library.ownerKey, clipId)
|
|
if (!existsSync(source)) {
|
|
throw new Error('Cannot continue the shot chain: the previous clip file is missing')
|
|
}
|
|
copyFileSync(source, currentPath)
|
|
}
|
|
return { tmpDir, currentPath, part1Path: join(tmpDir, 'part1.mp4'), framePath: join(tmpDir, 'last_frame.png') }
|
|
}
|
|
|
|
export async function continueQueuedExtensions(job: Job, params: VideoChainParams) {
|
|
const extensions = params.extensions || []
|
|
if (!extensions.length || !job.library) return
|
|
if (job.library.chainContinuing) return
|
|
job.library.chainContinuing = true
|
|
|
|
const ready = (status: { state: string; message: string; queueRunning?: number; queuePending?: number }) => {
|
|
emitChainJob(job, {
|
|
type: status.state === 'busy' ? 'busy' : 'status',
|
|
message: status.message,
|
|
progress: status.state === 'online' ? 3 : 1,
|
|
busy: status.state === 'busy',
|
|
queueRunning: status.queueRunning,
|
|
queuePending: status.queuePending
|
|
})
|
|
}
|
|
|
|
try {
|
|
const { currentPath, part1Path, framePath } = seedCurrentVideo(job)
|
|
const startFrom = job.library.chainIndex || 0
|
|
|
|
for (let i = startFrom; i < extensions.length; i++) {
|
|
assertJobActive(job)
|
|
const ext = extensions[i]
|
|
const isLast = i === extensions.length - 1
|
|
job.library.chainIndex = i + 1
|
|
job.library.chainStep = i + 2
|
|
job.library.chainLabel = `Extension ${i + 1}`
|
|
job.library.extendPart1Path = undefined
|
|
job.library.thumb = undefined
|
|
|
|
emitChainJob(job, { type: 'status', message: 'Waiting 3s buffer...', progress: 1 })
|
|
await sleep(3000)
|
|
assertJobActive(job)
|
|
|
|
copyFileSync(currentPath, part1Path)
|
|
try {
|
|
await extractLastFrame(part1Path, framePath)
|
|
} catch (error) {
|
|
const detail = error instanceof Error ? error.message : String(error)
|
|
throw new Error(detail.includes('last frame')
|
|
? detail
|
|
: `Could not extract the last frame for the next extension: ${detail}`)
|
|
}
|
|
const frame = readFileSync(framePath)
|
|
job.library.extendPart1Path = part1Path
|
|
job.library.prompt = ext.prompt
|
|
const sound = await probeHasAudio(currentPath)
|
|
const seed = Math.floor(Math.random() * 2_147_483_647)
|
|
|
|
await ensureComfyReady(ready)
|
|
await queueMiniMax(job, {
|
|
prompt: ext.prompt,
|
|
image: params.useIdentityRefs
|
|
? params.image
|
|
: { filename: 'last_frame.png', data: frame, type: 'image/png' },
|
|
width: params.width,
|
|
height: params.height,
|
|
steps: params.steps,
|
|
seed,
|
|
turbo: params.turbo,
|
|
length: frameLength(ext.duration, params.fps),
|
|
sound,
|
|
cfg: params.cfg,
|
|
fps: params.fps,
|
|
samplerName: params.samplerName,
|
|
scheduler: params.scheduler,
|
|
persist: isLast,
|
|
workflow: params.workflow,
|
|
duration: ext.duration,
|
|
useIdentityRefs: params.useIdentityRefs,
|
|
referenceImages: params.useIdentityRefs ? params.referenceImages : []
|
|
})
|
|
|
|
if (!isLast) {
|
|
if (!job.segmentBuffer?.length) {
|
|
throw new Error('Extension finished without a stitched video')
|
|
}
|
|
writeFileSync(currentPath, job.segmentBuffer)
|
|
job.segmentBuffer = undefined
|
|
job.library.extendPart1Path = undefined
|
|
}
|
|
}
|
|
} finally {
|
|
if (job.library) job.library.chainContinuing = false
|
|
}
|
|
}
|
|
|
|
export async function continueQueuedExtensionsIfNeeded(job: Job) {
|
|
if (!job.library) return
|
|
if (job.status === 'error' || job.status === 'cancelled' || job.status === 'complete') return
|
|
const remaining = remainingAfterCurrentShot(job.library)
|
|
if (!remaining.length) return
|
|
await continueQueuedExtensions(job, paramsFromJob(job))
|
|
}
|
|
|
|
export async function runGeneration(job: Job, params: VideoChainParams) {
|
|
const ready = (status: { state: string; message: string; queueRunning?: number; queuePending?: number }) => {
|
|
emitChainJob(job, {
|
|
type: status.state === 'busy' ? 'busy' : 'status',
|
|
message: status.message,
|
|
progress: status.state === 'online' ? 3 : 1,
|
|
busy: status.state === 'busy',
|
|
queueRunning: status.queueRunning,
|
|
queuePending: status.queuePending
|
|
})
|
|
}
|
|
|
|
await ensureComfyReady(ready)
|
|
|
|
const { extensions, ...base } = params
|
|
await queueMiniMax(job, {
|
|
...base,
|
|
persist: extensions.length === 0
|
|
})
|
|
|
|
if (!extensions.length || !job.library) return
|
|
await continueQueuedExtensions(job, params)
|
|
}
|