Queue durable Studio 2 jobs with sticky identity and saved metadata
This commit is contained in:
Towsty
@@ -0,0 +1,274 @@
|
||||
import { readFileSync, mkdirSync, existsSync, unlinkSync } from 'node:fs';
|
||||
import { join } from 'node:path';
|
||||
import template from '../../assets/studio2_minimax_native.json';
|
||||
import { nativeVideoGraph, attachHeroReference } from '~/shared/studio2/graphs.mjs';
|
||||
import { compilePrompt, scopedFile } from '~/shared/studio2/contracts.mjs';
|
||||
import { createJob, restoreJob, getJob, emitJob, type Job } from '../jobs';
|
||||
import { markStudioLive, onLiveVideoSettled, type StudioJob } from '../studioQueue';
|
||||
import { readRecord, saveRecord, studio2Root, records } from './store';
|
||||
import { ensureComfyReady } from '../comfyLifecycle';
|
||||
import { comfyFetch, queuePrompt, fetchHistory, extractVideo } from '../comfy';
|
||||
import { extractEditedImage } from '../imageComfy';
|
||||
import { buildImageV2Workflow, resolveKreaGenerateAssets } from '../imageWorkflowV2';
|
||||
import { resolveGraphLoraNames, applyUserLoraToGraph, ensureComfyLoraNames } from '../loras';
|
||||
import { getClip, getStill, stillPath, clipVideoPath, saveStill, saveClip, downloadComfyImage, downloadComfyVideo, attachStudio2Metadata } from '../library';
|
||||
import { imageDimensions } from '../resolution';
|
||||
import { stitchExtension } from '../ffmpeg';
|
||||
import { videoSourcePaths } from '../videoSources';
|
||||
import { resolveExtensionHandoffFrame, persistClipAnchorFrame } from '../extensionFrame';
|
||||
import { acquireSharedGpu, sharedGpuHeaders } from '../sharedGpu';
|
||||
type DiskFile = {
|
||||
filename: string;
|
||||
subfolder: string;
|
||||
type: string;
|
||||
};
|
||||
function update(r: any, state: string) { r.state = state; r.updatedAt = Date.now(); saveRecord(r); }
|
||||
async function upload(r: any, name: string, data: Buffer): Promise<string> {
|
||||
const prefix = String(useRuntimeConfig().comfyFilenamePrefix).replace(/\/$/, '') + `/studio2/${r.id}/${r.index}`;
|
||||
const body = new FormData();
|
||||
body.append('image', new Blob([new Uint8Array(data)]), name + '.png');
|
||||
body.append('subfolder', prefix);
|
||||
body.append('type', 'input');
|
||||
body.append('overwrite', 'true');
|
||||
const response = await comfyFetch('/upload/image', { method: 'POST', body });
|
||||
if (!response.ok)
|
||||
throw new Error('Reference image upload failed');
|
||||
const file = await response.json() as any;
|
||||
const disk = { filename: file.name, subfolder: file.subfolder || '', type: 'input' };
|
||||
if (!scopedFile(disk, String(useRuntimeConfig().comfyFilenamePrefix)))
|
||||
throw new Error('Comfy returned an input outside this instance prefix');
|
||||
r.files.push(disk);
|
||||
saveRecord(r);
|
||||
return `${disk.subfolder}/${disk.filename}`;
|
||||
}
|
||||
async function purge(r: any) {
|
||||
const config = useRuntimeConfig(), prefix = String(config.comfyFilenamePrefix);
|
||||
const files = (r.files as DiskFile[]).filter(file => scopedFile(file, prefix));
|
||||
r.purgeResult = 'Left on host';
|
||||
if (files.length !== r.files.length || !files.length)
|
||||
return;
|
||||
try {
|
||||
const response = await fetch(String(config.comfyControlUrl).replace(/\/$/, '') + '/studio2/purge', {
|
||||
method: 'POST', headers: { 'Content-Type': 'application/json', ...sharedGpuHeaders() },
|
||||
body: JSON.stringify({ prefix, files }), signal: AbortSignal.timeout(12000)
|
||||
});
|
||||
const result = await response.json() as any;
|
||||
if (response.ok && result.ok && result.cleared === files.length)
|
||||
r.purgeResult = 'Cleared input + output';
|
||||
}
|
||||
catch { /* Keep the durable library copy and report unconfirmed host cleanup. */ }
|
||||
}
|
||||
async function prepareGraph(r: any) {
|
||||
const q = r.request, s = q.settings, prefix = String(useRuntimeConfig().comfyFilenamePrefix).replace(/\/$/, '') + `/studio2/${r.id}/${r.index}`;
|
||||
const load = async (id: string, label: string) => id ? upload(r, label, readFileSync(stillPath(r.owner, getStill(r.owner, id).id))) : '';
|
||||
const hero = await load(q.identityStillId, 'hero'), a = await load(q.imageAId, 'source'), b = await load(q.imageBId, 'compose'), mask = await load(q.maskId, 'mask');
|
||||
let graph: any;
|
||||
if (['video', 'extend'].includes(q.mode)) {
|
||||
if (q.engine !== 'minimax')
|
||||
throw new Error('Studio 2 native identity video currently requires MiniMax. Use the existing xAIGen studio for LTX.');
|
||||
let start = a;
|
||||
if (q.startClipId) {
|
||||
const dir = join(studio2Root(), r.id);
|
||||
mkdirSync(dir, { recursive: true });
|
||||
const dest = join(dir, 'handoff.png');
|
||||
await resolveExtensionHandoffFrame({ ownerKey: r.owner, sourceClipId: q.startClipId, sourceVideoPath: clipVideoPath(r.owner, q.startClipId), destPath: dest });
|
||||
start = await upload(r, 'start', readFileSync(dest));
|
||||
unlinkSync(dest);
|
||||
}
|
||||
const end = await load(q.endStillId, 'end'), guides = [];
|
||||
for (const [i, g] of q.guides.entries())
|
||||
guides.push({ image: await load(g.stillId, `guide${i}`), frame: g.frame });
|
||||
graph = nativeVideoGraph(template, q, { hero, start, end, guides }, prefix);
|
||||
await ensureComfyLoraNames('video');
|
||||
resolveGraphLoraNames(graph, 'video');
|
||||
applyUserLoraToGraph(graph, s.loraStack);
|
||||
r.graphId = 'studio2_minimax_native.json';
|
||||
}
|
||||
else {
|
||||
const mode = q.mode === 'iterate' ? (a ? 'edit' : 'generate') : q.mode;
|
||||
const found = q.engine === 'krea' ? await resolveKreaGenerateAssets() : null;
|
||||
const assets = found ? { kreaUnetName: found.unet, kreaClipName: found.clip, kreaVaeName: found.vae, kreaConceptLora: found.conceptLora } : {};
|
||||
graph = buildImageV2Workflow({ ...s, ...assets, engine: q.engine, mode, task: 'scene', prompt: q.compiledPrompt, imageAName: a, imageBName: mode === 'compose' ? b : undefined, maskName: mode === 'refine' ? mask : undefined, filenamePrefix: prefix + '/image' }).graph;
|
||||
if (q.engine === 'flux')
|
||||
attachHeroReference(graph, hero);
|
||||
else if (hero && ['edit', 'compose', 'iterate'].includes(q.mode))
|
||||
throw new Error('Krea native hero-reference binding is not supported by this graph. Select Klein explicitly; no fallback is performed.');
|
||||
await ensureComfyLoraNames('image');
|
||||
resolveGraphLoraNames(graph, 'image');
|
||||
r.graphId = `${q.engine === 'flux' ? 'klein' : 'krea'}_v2_${mode}.json`;
|
||||
}
|
||||
saveRecord(r);
|
||||
return graph;
|
||||
}
|
||||
async function run(r: any, job: Job) {
|
||||
try {
|
||||
const prompts = r.prompts || [r.request.promptSections, ...r.request.batch];
|
||||
r.prompts = prompts;
|
||||
saveRecord(r);
|
||||
for (r.index = r.index || 0; r.index < prompts.length; r.index++) {
|
||||
if (job.status === 'cancelled')
|
||||
throw new Error('Cancelled');
|
||||
if (!r.promptId) {
|
||||
while (!await acquireSharedGpu()) {
|
||||
if (job.status === 'cancelled')
|
||||
throw new Error('Cancelled');
|
||||
await new Promise(resolve => setTimeout(resolve, 2000));
|
||||
}
|
||||
r.files = [];
|
||||
r.request.promptSections = prompts[r.index];
|
||||
r.request.compiledPrompt = compilePrompt(prompts[r.index], ['video', 'extend'].includes(r.request.mode));
|
||||
update(r, 'waking');
|
||||
await ensureComfyReady(message => emitJob(job, { type: 'status', message }));
|
||||
if (job.status === 'cancelled')
|
||||
throw new Error('Cancelled');
|
||||
const graph = await prepareGraph(r);
|
||||
if (job.status === 'cancelled')
|
||||
throw new Error('Cancelled');
|
||||
update(r, 'submitting');
|
||||
const queued = await queuePrompt(graph, job.clientId);
|
||||
job.promptId = queued.prompt_id;
|
||||
r.promptId = queued.prompt_id;
|
||||
r.renderStartedAt = Date.now();
|
||||
update(r, 'rendering');
|
||||
}
|
||||
let history: any, missingSince = 0;
|
||||
for (;;) {
|
||||
try {
|
||||
history = await fetchHistory(r.promptId);
|
||||
}
|
||||
catch {
|
||||
history = null;
|
||||
}
|
||||
const e = history?.[r.promptId];
|
||||
if (job.status === 'cancelled')
|
||||
throw new Error('Cancelled');
|
||||
if (['error', 'interrupted'].includes(e?.status?.status_str))
|
||||
throw new Error('Comfy generation failed or was interrupted; inspect the host error log.');
|
||||
if (!e) {
|
||||
try {
|
||||
const response = await comfyFetch('/queue', { signal: AbortSignal.timeout(5000) });
|
||||
if (response.ok) {
|
||||
const q = await response.json() as any;
|
||||
const present = [...(q.queue_running || []), ...(q.queue_pending || [])].some((entry: any) => entry[1] === r.promptId);
|
||||
if (present)
|
||||
missingSince = 0;
|
||||
else if (!missingSince)
|
||||
missingSince = Date.now();
|
||||
else if (Date.now() - missingSince > 60000)
|
||||
throw new Error('PROMPT_MISSING');
|
||||
}
|
||||
}
|
||||
catch (error: any) {
|
||||
if (error.message === 'PROMPT_MISSING')
|
||||
throw new Error('The saved Comfy prompt is no longer in history or the queue. It was not resubmitted.');
|
||||
}
|
||||
}
|
||||
if (e?.status?.completed || e?.status?.status_str === 'success')
|
||||
break;
|
||||
if (job.status === 'cancelled')
|
||||
throw new Error('Cancelled');
|
||||
await new Promise(resolve => setTimeout(resolve, 1500));
|
||||
}
|
||||
const messages = history[r.promptId]?.status?.messages || [];
|
||||
const executionStart = messages.find((m: any) => m[0] === 'execution_start')?.[1]?.timestamp;
|
||||
const executionEnd = messages.find((m: any) => m[0] === 'execution_success')?.[1]?.timestamp;
|
||||
r.gpuSeconds = executionStart && executionEnd ? (executionEnd - executionStart) / 1000 : null;
|
||||
update(r, 'saving');
|
||||
job.saving = true;
|
||||
const q = r.request, s = q.settings, video = ['video', 'extend'].includes(q.mode);
|
||||
const file = video ? extractVideo(history, r.promptId) : extractEditedImage(history, r.promptId);
|
||||
if (!file)
|
||||
throw new Error('Comfy finished without a media output');
|
||||
if (!r.files.some((f: DiskFile) => f.filename === file.filename && f.subfolder === file.subfolder))
|
||||
r.files.push(file);
|
||||
let asset: any = r.savedAssetId ? (video ? getClip(r.owner, r.savedAssetId) : getStill(r.owner, r.savedAssetId)) : null;
|
||||
if (!asset && video) {
|
||||
const data = await downloadComfyVideo(file);
|
||||
const sourceSegments = q.startClipId ? videoSourcePaths(clipVideoPath(r.owner, q.startClipId)) : [];
|
||||
const tmp = join(studio2Root(), r.id);
|
||||
mkdirSync(tmp, { recursive: true });
|
||||
const assembled = q.startClipId ? await stitchExtension({ part1Path: clipVideoPath(r.owner, q.startClipId), part2: data, tmpDir: tmp, sourcePaths: sourceSegments }) : data;
|
||||
asset = await saveClip({ ownerKey: r.owner, folderId: q.folderId, prompt: q.compiledPrompt, ...s, aspect: `${s.width}:${s.height}`, hideThumbnail: false, video: assembled, originalSegment: data, sourceSegments, fps: 24, sound: true, familyId: r.familyId, parentClipId: q.startClipId || undefined, chainIndex: r.index, comfyFilename: file.filename });
|
||||
const anchor = history[r.promptId].outputs?.anchor_save?.images?.[0];
|
||||
if (anchor) {
|
||||
r.files.push(anchor);
|
||||
await persistClipAnchorFrame({ ownerKey: r.owner, clipId: asset.id, fromBuffer: await downloadComfyImage(anchor) });
|
||||
}
|
||||
}
|
||||
else if (!asset) {
|
||||
const data = await downloadComfyImage(file), size = imageDimensions(data) || { width: s.width, height: s.height };
|
||||
asset = await saveStill({ ownerKey: r.owner, folderId: q.folderId, filename: file.filename, data, ...size, role: 'output', prompt: q.compiledPrompt, familyId: r.familyId, chainIndex: r.index, settings: s });
|
||||
}
|
||||
if (!asset)
|
||||
throw new Error('Library did not save the output');
|
||||
r.savedAssetId = asset.id;
|
||||
saveRecord(r);
|
||||
r.savedAt = Date.now();
|
||||
r.wallTime = (r.savedAt - r.startedAt) / 1000;
|
||||
const metadata = { ...structuredClone(q), id: r.id, kind: video ? 'video' : 'image', graphId: r.graphId, promptId: r.promptId, queuedAt: r.queuedAt, startedAt: r.startedAt, savedAt: r.savedAt, gpuSeconds: r.gpuSeconds, wallTime: r.wallTime, outputWidth: asset.width || s.width, outputHeight: asset.height || s.height, purgeResult: 'Left on host' };
|
||||
await attachStudio2Metadata(r.owner, asset.id, metadata);
|
||||
await purge(r);
|
||||
metadata.purgeResult = r.purgeResult;
|
||||
await attachStudio2Metadata(r.owner, asset.id, metadata);
|
||||
if (!r.outputs.some((o: any) => o.id === asset.id))
|
||||
r.outputs.push({ id: asset.id, kind: video ? 'clip' : 'still', studio2: metadata });
|
||||
saveRecord(r);
|
||||
job.saving = false;
|
||||
if (video) {
|
||||
q.startClipId = asset.id;
|
||||
q.startFrameSource = { kind: 'previous-last-frame', clipId: asset.id };
|
||||
}
|
||||
else if (q.mode === 'iterate')
|
||||
q.imageAId = asset.id;
|
||||
r.promptId = '';
|
||||
r.savedAssetId = '';
|
||||
r.index++;
|
||||
saveRecord(r);
|
||||
r.index--;
|
||||
// Identity remains the original hero throughout the batch.
|
||||
}
|
||||
if (job.status === 'cancelled')
|
||||
throw new Error('Cancelled');
|
||||
job.status = 'complete';
|
||||
update(r, 'complete');
|
||||
emitJob(job, { type: 'complete', message: 'Studio 2 output saved' });
|
||||
}
|
||||
catch (e: any) {
|
||||
job.status = job.status === 'cancelled' ? 'cancelled' : 'error';
|
||||
job.error = e.message;
|
||||
r.error = e.message;
|
||||
update(r, job.status === 'cancelled' ? 'cancelled' : 'failed');
|
||||
emitJob(job, { type: 'error', message: e.message });
|
||||
}
|
||||
finally {
|
||||
job.saving = false;
|
||||
await onLiveVideoSettled(job);
|
||||
}
|
||||
}
|
||||
export async function startStudio2Job(item: StudioJob) {
|
||||
const r = readRecord(item.ownerKey, item.payload.studio2Id!), job = createJob(['video', 'extend'].includes(r.request.mode) ? 'video' : 'edit');
|
||||
job.studio2 = true;
|
||||
job.library = { ownerKey: item.ownerKey, folderId: r.request.folderId, extensions: [], queueAutoRun: false } as any;
|
||||
job.status = 'running';
|
||||
r.startedAt = Date.now();
|
||||
r.liveId = job.id;
|
||||
r.clientId = job.clientId;
|
||||
saveRecord(r);
|
||||
await markStudioLive(item.ownerKey, item.id, job.id);
|
||||
void run(r, job);
|
||||
}
|
||||
export function resumeStudio2Jobs() {
|
||||
for (const r of records()) {
|
||||
if (!r.liveId || getJob(r.liveId) || ['waiting', 'complete', 'failed', 'cancelled'].includes(r.state))
|
||||
continue;
|
||||
if (r.state === 'submitting' && !r.promptId) {
|
||||
r.state = 'failed';
|
||||
r.error = 'Submission outcome unknown after restart. Inspect the host before retrying.';
|
||||
saveRecord(r);
|
||||
continue;
|
||||
}
|
||||
const job = restoreJob({ id: r.liveId, clientId: r.clientId, startedAt: r.startedAt, promptId: r.promptId || '', library: { ownerKey: r.owner, folderId: r.request.folderId, extensions: [], queueAutoRun: false } as any });
|
||||
job.studio2 = true;
|
||||
void run(r, job);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user