449 lines
15 KiB
TypeScript
449 lines
15 KiB
TypeScript
import { existsSync, mkdirSync, readdirSync, readFileSync, renameSync, writeFileSync } from 'node:fs'
|
|
import { join } from 'node:path'
|
|
import { getJob, listJobs, type Job } from '~/server/utils/jobs'
|
|
|
|
export type ShotQueueStatus = 'idle' | 'running' | 'paused' | 'complete' | 'error'
|
|
export type ShotSegmentStatus = 'pending' | 'running' | 'complete' | 'error'
|
|
|
|
export interface ShotQueueSegment {
|
|
index: number
|
|
prompt: string
|
|
duration: number
|
|
status: ShotSegmentStatus
|
|
clipId?: string
|
|
error?: string
|
|
permanenceRefs?: import('~/utils/globalLocks').PermanenceRef[]
|
|
}
|
|
|
|
export interface ShotQueue {
|
|
id: string
|
|
name: string
|
|
familyId: string
|
|
folderId: string
|
|
createdAt: number
|
|
updatedAt: number
|
|
status: ShotQueueStatus
|
|
autoRun: boolean
|
|
stopAfterCurrent?: boolean
|
|
currentClipId?: string
|
|
currentJobId?: string
|
|
stillId?: string
|
|
stillFilename?: string
|
|
hideThumbnail: boolean
|
|
hideInput?: boolean
|
|
aspect: string
|
|
width: number
|
|
height: number
|
|
steps: number
|
|
turbo: boolean
|
|
cfg?: number
|
|
fps?: number
|
|
samplerName?: string
|
|
scheduler?: string
|
|
workflow?: import('~/utils/videoModels').VideoWorkflowId
|
|
sound?: boolean
|
|
useIdentityRefs?: boolean
|
|
referenceStillIds?: Array<string | null>
|
|
globalLocks?: string
|
|
permanenceRefs?: import('~/utils/globalLocks').PermanenceRef[]
|
|
segments: ShotQueueSegment[]
|
|
lastError?: string
|
|
}
|
|
|
|
const writeChains = new Map<string, Promise<unknown>>()
|
|
const activeJobs = new Map<string, string>()
|
|
|
|
function libraryRoot() {
|
|
const config = useRuntimeConfig()
|
|
return (config.libraryDir || process.env.LIBRARY_DIR || '/data/library').replace(/\/$/, '')
|
|
}
|
|
|
|
function queuesPath(owner: string) {
|
|
return join(libraryRoot(), 'users', owner, 'shot-queues.json')
|
|
}
|
|
|
|
function ensureOwner(owner: string) {
|
|
mkdirSync(join(libraryRoot(), 'users', owner), { recursive: true })
|
|
}
|
|
|
|
function readQueues(owner: string): ShotQueue[] {
|
|
ensureOwner(owner)
|
|
const path = queuesPath(owner)
|
|
if (!existsSync(path)) return []
|
|
try {
|
|
const parsed = JSON.parse(readFileSync(path, 'utf8'))
|
|
return Array.isArray(parsed) ? parsed : []
|
|
} catch {
|
|
return []
|
|
}
|
|
}
|
|
|
|
function writeQueues(owner: string, queues: ShotQueue[]) {
|
|
ensureOwner(owner)
|
|
const path = queuesPath(owner)
|
|
const tmp = `${path}.tmp`
|
|
writeFileSync(tmp, JSON.stringify(queues, null, 2))
|
|
renameSync(tmp, path)
|
|
}
|
|
|
|
function mutate<T>(owner: string, fn: (queues: ShotQueue[]) => T): Promise<T> {
|
|
const prev = writeChains.get(owner) || Promise.resolve()
|
|
const run = prev.then(() => {
|
|
const queues = readQueues(owner)
|
|
const result = fn(queues)
|
|
writeQueues(owner, queues)
|
|
return result
|
|
})
|
|
writeChains.set(owner, run.then(() => undefined, () => undefined))
|
|
return run
|
|
}
|
|
|
|
export function pendingSegmentCount(queue: ShotQueue) {
|
|
return queue.segments.filter(segment => segment.status === 'pending' || segment.status === 'error').length
|
|
}
|
|
|
|
export function completedSegmentCount(queue: ShotQueue) {
|
|
return queue.segments.filter(segment => segment.status === 'complete').length
|
|
}
|
|
|
|
export function summarizeQueue(queue: ShotQueue) {
|
|
return {
|
|
...queue,
|
|
pendingCount: pendingSegmentCount(queue),
|
|
completedCount: completedSegmentCount(queue),
|
|
totalCount: queue.segments.length
|
|
}
|
|
}
|
|
|
|
export function listShotQueues(owner: string) {
|
|
return readQueues(owner)
|
|
.slice()
|
|
.sort((a, b) => b.updatedAt - a.updatedAt)
|
|
}
|
|
|
|
export function getShotQueue(owner: string, id: string) {
|
|
return readQueues(owner).find(queue => queue.id === id) || null
|
|
}
|
|
|
|
export function queueIsProcessing(id: string) {
|
|
return activeJobs.has(id)
|
|
}
|
|
|
|
export function setQueueJob(id: string, jobId: string | null) {
|
|
if (jobId) activeJobs.set(id, jobId)
|
|
else activeJobs.delete(id)
|
|
}
|
|
|
|
export function getQueueJobId(id: string) {
|
|
return activeJobs.get(id)
|
|
}
|
|
|
|
export async function createShotQueue(params: {
|
|
ownerKey: string
|
|
name?: string
|
|
familyId: string
|
|
folderId: string
|
|
autoRun: boolean
|
|
stillId?: string
|
|
stillFilename?: string
|
|
hideThumbnail: boolean
|
|
hideInput?: boolean
|
|
aspect: string
|
|
width: number
|
|
height: number
|
|
steps: number
|
|
turbo: boolean
|
|
cfg?: number
|
|
fps?: number
|
|
samplerName?: string
|
|
scheduler?: string
|
|
workflow?: import('~/utils/videoModels').VideoWorkflowId
|
|
sound?: boolean
|
|
useIdentityRefs?: boolean
|
|
referenceStillIds?: Array<string | null>
|
|
globalLocks?: string
|
|
permanenceRefs?: import('~/utils/globalLocks').PermanenceRef[]
|
|
initial: { prompt: string; duration: number; permanenceRefs?: import('~/utils/globalLocks').PermanenceRef[] }
|
|
extensions: { prompt: string; duration: number; permanenceRefs?: import('~/utils/globalLocks').PermanenceRef[] }[]
|
|
jobId?: string
|
|
}): Promise<ShotQueue> {
|
|
const now = Date.now()
|
|
const queue: ShotQueue = {
|
|
id: crypto.randomUUID(),
|
|
name: (params.name || '').trim() || params.initial.prompt.slice(0, 80),
|
|
familyId: params.familyId,
|
|
folderId: params.folderId,
|
|
createdAt: now,
|
|
updatedAt: now,
|
|
status: 'running',
|
|
autoRun: params.autoRun,
|
|
stopAfterCurrent: false,
|
|
currentJobId: params.jobId,
|
|
stillId: params.stillId,
|
|
stillFilename: params.stillFilename,
|
|
hideThumbnail: params.hideThumbnail,
|
|
hideInput: params.hideInput,
|
|
aspect: params.aspect,
|
|
width: params.width,
|
|
height: params.height,
|
|
steps: params.steps,
|
|
turbo: params.turbo,
|
|
cfg: params.cfg,
|
|
fps: params.fps,
|
|
samplerName: params.samplerName,
|
|
scheduler: params.scheduler,
|
|
workflow: params.workflow,
|
|
sound: params.sound,
|
|
useIdentityRefs: params.useIdentityRefs,
|
|
referenceStillIds: params.referenceStillIds,
|
|
globalLocks: params.globalLocks,
|
|
permanenceRefs: params.permanenceRefs,
|
|
segments: [
|
|
{ index: 0, prompt: params.initial.prompt, duration: params.initial.duration, status: 'running', permanenceRefs: params.initial.permanenceRefs },
|
|
...params.extensions.map((item, index) => ({
|
|
index: index + 1,
|
|
prompt: item.prompt,
|
|
duration: item.duration,
|
|
status: 'pending' as const,
|
|
permanenceRefs: item.permanenceRefs
|
|
}))
|
|
]
|
|
}
|
|
await mutate(params.ownerKey, (queues) => {
|
|
queues.unshift(queue)
|
|
return queue
|
|
})
|
|
if (params.jobId) setQueueJob(queue.id, params.jobId)
|
|
return queue
|
|
}
|
|
|
|
export async function updateShotQueue(owner: string, id: string, patch: (queue: ShotQueue) => void) {
|
|
return mutate(owner, (queues) => {
|
|
const queue = queues.find(item => item.id === id)
|
|
if (!queue) throw createError({ statusCode: 404, statusMessage: 'Queue not found' })
|
|
patch(queue)
|
|
queue.updatedAt = Date.now()
|
|
return structuredClone(queue)
|
|
})
|
|
}
|
|
|
|
export async function patchShotQueueSegments(
|
|
owner: string,
|
|
id: string,
|
|
segments: Array<{ index: number; prompt?: string; duration?: number }>
|
|
) {
|
|
return updateShotQueue(owner, id, (queue) => {
|
|
if (queue.status === 'complete') {
|
|
throw createError({ statusCode: 409, statusMessage: 'This batch is already finished' })
|
|
}
|
|
for (const next of segments) {
|
|
const segment = queue.segments.find(item => item.index === next.index)
|
|
if (!segment) continue
|
|
if (segment.status === 'complete' || segment.status === 'running') continue
|
|
if (typeof next.prompt === 'string') {
|
|
const prompt = next.prompt.trim()
|
|
if (!prompt) throw createError({ statusCode: 400, statusMessage: `Shot ${next.index + 1} needs a prompt` })
|
|
segment.prompt = prompt
|
|
}
|
|
if (typeof next.duration === 'number' && Number.isFinite(next.duration)) {
|
|
segment.duration = Math.min(120, Math.max(0.5, Math.round(next.duration * 10) / 10))
|
|
}
|
|
}
|
|
})
|
|
}
|
|
|
|
export async function markQueueSegment(owner: string, id: string, index: number, status: ShotSegmentStatus, extra: Partial<ShotQueueSegment> = {}) {
|
|
return updateShotQueue(owner, id, (queue) => {
|
|
const segment = queue.segments.find(item => item.index === index)
|
|
if (!segment) return
|
|
segment.status = status
|
|
if (extra.clipId) segment.clipId = extra.clipId
|
|
if (extra.error) segment.error = extra.error
|
|
else if (status !== 'error') delete segment.error
|
|
if (extra.clipId) queue.currentClipId = extra.clipId
|
|
if (status === 'complete' && pendingSegmentCount(queue) === 0) {
|
|
queue.status = 'complete'
|
|
queue.autoRun = false
|
|
queue.currentJobId = undefined
|
|
}
|
|
})
|
|
}
|
|
|
|
export async function syncQueueFromJob(job: Job, clipId?: string) {
|
|
const library = job.library
|
|
if (!library?.queueId) return null
|
|
const index = library.chainIndex || 0
|
|
const failed = job.status === 'error' || job.status === 'cancelled'
|
|
if (failed) {
|
|
return updateShotQueue(library.ownerKey, library.queueId, (queue) => {
|
|
const segment = queue.segments.find(item => item.index === index)
|
|
if (segment && segment.status !== 'complete') {
|
|
segment.status = 'error'
|
|
segment.error = job.error || 'Generation failed'
|
|
}
|
|
queue.status = job.status === 'cancelled' ? 'paused' : 'error'
|
|
queue.lastError = job.error
|
|
queue.currentJobId = undefined
|
|
queue.autoRun = job.status === 'cancelled' ? false : queue.autoRun
|
|
}).catch(() => null)
|
|
}
|
|
if (!clipId) return getShotQueue(library.ownerKey, library.queueId)
|
|
return markQueueSegment(library.ownerKey, library.queueId, index, 'complete', { clipId })
|
|
}
|
|
|
|
export async function prepareQueueBurst(owner: string, id: string, count: number | 'all') {
|
|
return mutate(owner, (queues) => {
|
|
const queue = queues.find(item => item.id === id)
|
|
if (!queue) throw createError({ statusCode: 404, statusMessage: 'Queue not found' })
|
|
if (activeJobs.has(queue.id)) {
|
|
throw createError({ statusCode: 409, statusMessage: 'This batch is already processing' })
|
|
}
|
|
const remaining = queue.segments.filter(segment => segment.status === 'pending' || segment.status === 'error')
|
|
if (!remaining.length) {
|
|
throw createError({ statusCode: 400, statusMessage: 'No remaining shots to process' })
|
|
}
|
|
const initial = queue.segments[0]
|
|
if (initial && initial.status !== 'complete') {
|
|
throw createError({ statusCode: 409, statusMessage: 'The first shot is still running or did not save. Wait for it, or generate again.' })
|
|
}
|
|
if (!queue.currentClipId) {
|
|
throw createError({ statusCode: 409, statusMessage: 'The previous clip is missing, so the next shot cannot start' })
|
|
}
|
|
const n = count === 'all' ? remaining.length : Math.min(remaining.length, Math.max(1, Math.floor(count)))
|
|
remaining.slice(0, n).forEach((segment) => {
|
|
segment.status = segment.index === remaining[0].index ? 'running' : 'pending'
|
|
delete segment.error
|
|
})
|
|
queue.autoRun = count === 'all'
|
|
queue.stopAfterCurrent = false
|
|
queue.status = 'running'
|
|
queue.lastError = undefined
|
|
queue.updatedAt = Date.now()
|
|
return { queue: structuredClone(queue), count: n }
|
|
})
|
|
}
|
|
|
|
export async function finishQueueBurst(owner: string, id: string, paused: boolean, error?: string) {
|
|
return updateShotQueue(owner, id, (queue) => {
|
|
queue.currentJobId = undefined
|
|
for (const segment of queue.segments) {
|
|
if (segment.status !== 'running') continue
|
|
if (error) {
|
|
segment.status = 'error'
|
|
segment.error = error
|
|
} else {
|
|
segment.status = 'pending'
|
|
delete segment.error
|
|
}
|
|
}
|
|
if (error) {
|
|
queue.status = 'error'
|
|
queue.lastError = error
|
|
queue.autoRun = false
|
|
return
|
|
}
|
|
if (pendingSegmentCount(queue) === 0) {
|
|
queue.status = 'complete'
|
|
queue.autoRun = false
|
|
return
|
|
}
|
|
queue.status = paused || !queue.autoRun ? 'paused' : 'idle'
|
|
if (paused) queue.autoRun = false
|
|
})
|
|
}
|
|
|
|
function liveJobForQueue(id: string): Job | undefined {
|
|
const mapped = activeJobs.get(id)
|
|
if (mapped) {
|
|
const job = getJob(mapped)
|
|
if (job) return job
|
|
}
|
|
return listJobs().find(job => (
|
|
job.library?.queueId === id
|
|
&& (job.status === 'queued' || job.status === 'uploading' || job.status === 'running')
|
|
))
|
|
}
|
|
|
|
export async function setShotQueuePause(owner: string, id: string, pause: boolean) {
|
|
const live = liveJobForQueue(id)
|
|
if (live?.library) {
|
|
live.library.stopAfterCurrent = pause
|
|
if (pause) live.library.queueAutoRun = false
|
|
}
|
|
return updateShotQueue(owner, id, (queue) => {
|
|
if (pause) {
|
|
queue.autoRun = false
|
|
queue.stopAfterCurrent = true
|
|
if (queue.status === 'running') queue.lastError = undefined
|
|
else if (queue.status !== 'complete') queue.status = 'paused'
|
|
} else {
|
|
queue.stopAfterCurrent = false
|
|
if (live?.library) live.library.queueAutoRun = queue.autoRun === true
|
|
if (queue.status === 'paused') queue.status = 'idle'
|
|
}
|
|
})
|
|
}
|
|
|
|
export async function pauseShotQueue(owner: string, id: string) {
|
|
return setShotQueuePause(owner, id, true)
|
|
}
|
|
|
|
export async function deleteShotQueue(owner: string, id: string, options: { force?: boolean } = {}) {
|
|
if (!options.force && activeJobs.has(id)) {
|
|
throw createError({ statusCode: 409, statusMessage: 'Stop this batch before deleting it' })
|
|
}
|
|
activeJobs.delete(id)
|
|
return mutate(owner, (queues) => {
|
|
const index = queues.findIndex(item => item.id === id)
|
|
if (index < 0) throw createError({ statusCode: 404, statusMessage: 'Queue not found' })
|
|
const [removed] = queues.splice(index, 1)
|
|
return removed
|
|
})
|
|
}
|
|
|
|
export async function clearShotQueues(owner: string) {
|
|
return mutate(owner, (queues) => {
|
|
const kept: ShotQueue[] = []
|
|
let removed = 0
|
|
for (const queue of queues) {
|
|
if (activeJobs.has(queue.id)) {
|
|
kept.push(queue)
|
|
continue
|
|
}
|
|
removed += 1
|
|
}
|
|
queues.splice(0, queues.length, ...kept)
|
|
return { removed, remaining: kept.length }
|
|
})
|
|
}
|
|
|
|
export function listOwnersWithQueues() {
|
|
const root = join(libraryRoot(), 'users')
|
|
if (!existsSync(root)) return []
|
|
return readdirSync(root, { withFileTypes: true })
|
|
.filter(entry => entry.isDirectory())
|
|
.map(entry => entry.name)
|
|
.filter(owner => existsSync(queuesPath(owner)))
|
|
}
|
|
|
|
export function remainingFromQueue(queue: ShotQueue): { prompt: string; duration: number; permanenceRefs?: import('~/utils/globalLocks').PermanenceRef[] }[] {
|
|
return queue.segments
|
|
.filter(segment => segment.index > 0)
|
|
.map(segment => ({ prompt: segment.prompt, duration: segment.duration, permanenceRefs: segment.permanenceRefs }))
|
|
}
|
|
|
|
export function lastCompletedIndex(queue: ShotQueue) {
|
|
const done = queue.segments.filter(segment => segment.status === 'complete')
|
|
if (!done.length) return -1
|
|
return Math.max(...done.map(segment => segment.index))
|
|
}
|
|
|
|
export function liveSegmentPrompt(queue: ShotQueue, extensionIndex: number) {
|
|
const segment = queue.segments.find(item => item.index === extensionIndex + 1)
|
|
return {
|
|
prompt: segment?.prompt || '',
|
|
duration: segment?.duration || 5,
|
|
permanenceRefs: segment?.permanenceRefs
|
|
}
|
|
}
|