340 lines
20 KiB
TypeScript
340 lines
20 KiB
TypeScript
import { kleinIdentityPlan, applyKleinIdentity } from '~/shared/studio2/klein-identity.mjs';
|
|
import { stylePrompt } from '~/shared/studio2/styles.mjs';
|
|
import { purge } from './cleanup';
|
|
import { watchProgress } from './progress';
|
|
import { fitStill, needsFit } from './media';
|
|
import { resolveRequestSize } from './size';
|
|
import { queueSeeds } from '~/shared/studio2/seed.mjs';
|
|
import { readFileSync, writeFileSync, mkdirSync, existsSync, unlinkSync } from 'node:fs';
|
|
import { join } from 'node:path';
|
|
import template from '../../assets/studio2_minimax_native.json';
|
|
import { nativeVideoGraph, attachHeroReference, applyResolvedImageSize } 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, probeDuration } 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(); if (['complete','failed','cancelled'].includes(state)) r.finishedAt ??= r.updatedAt; saveRecord(r); }
|
|
// Labels whose pixels must never be resampled: the hero identity still, and the
|
|
// extend/video start frame (which is either the exact anchor PNG or an already-canvas-sized
|
|
// extracted last frame). Fitting these every hop is what destroys identity across chained extends.
|
|
const NEVER_FIT = new Set(['hero', 'start']);
|
|
async function upload(r: any, name: string, data: Buffer): Promise<string> {
|
|
const prefix = String(useRuntimeConfig().comfyFilenamePrefix).replace(/\/$/, '') + `/studio2/${r.id}/${r.index}`;
|
|
let fitted = false;
|
|
if (!NEVER_FIT.has(name) && needsFit(data, r.request.settings.width, r.request.settings.height)) {
|
|
data = await fitStill(data, r.request.settings.width, r.request.settings.height);
|
|
fitted = true;
|
|
}
|
|
if (name === 'start') { r.handoffFitted = fitted; saveRecord(r); }
|
|
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 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 video=['video','extend'].includes(q.mode);
|
|
const identityPlan=kleinIdentityPlan(q);
|
|
const sourceId=identityPlan ? identityPlan.sourceId : q.imageAId;
|
|
await resolveRequestSize(r.owner,identityPlan ? {...q,imageAId:sourceId} : q);
|
|
saveRecord(r);
|
|
const a=await load(sourceId,'source'),b=await load(q.imageBId,'compose'),mask=await load(q.maskId,'mask');
|
|
const attachHero=(video || q.engine==='flux') && (q.lockFace!==false || q.lockOutfit!==false) && !(!video && sourceId && sourceId===q.identityStillId);
|
|
const hero=attachHero ? await load(q.identityStillId,'hero') : '';
|
|
r.sourceStillId=sourceId || null;
|
|
r.heroReferenceAttached=!!hero || (!video && q.engine==='flux' && !!sourceId && sourceId===q.identityStillId);
|
|
|
|
let graph: any;
|
|
if (['video', 'extend'].includes(q.mode)) {
|
|
let start = a;
|
|
if (q.startClipId) {
|
|
const dir = join(studio2Root(), r.id);
|
|
mkdirSync(dir, { recursive: true });
|
|
const dest = join(dir, 'handoff.png');
|
|
const handoff = await resolveExtensionHandoffFrame({ ownerKey: r.owner, sourceClipId: q.startClipId, sourceVideoPath: clipVideoPath(r.owner, q.startClipId), destPath: dest });
|
|
r.handoffSource = handoff.source;
|
|
start = await upload(r, 'start', readFileSync(dest));
|
|
unlinkSync(dest);
|
|
}
|
|
if (q.engine === 'ltx') {
|
|
const { buildWorkflow } = await import('../workflow');
|
|
const { frameLength } = await import('../videoChain');
|
|
const { ltxWorkflowEnabled, LTX_DISABLED_MESSAGE } = await import('~/utils/videoModels');
|
|
if (!ltxWorkflowEnabled()) throw new Error(LTX_DISABLED_MESSAGE);
|
|
const textToVideo = q.mode === 'video' && !start;
|
|
if (!start && !textToVideo) throw new Error('Choose a start still for LTX video.');
|
|
const fps = s.fps || 24;
|
|
graph = buildWorkflow({
|
|
workflow: textToVideo ? 'ltx-t2v' : 'ltx',
|
|
imageName: start || '',
|
|
prompt: q.compiledPrompt,
|
|
width: s.width,
|
|
height: s.height,
|
|
length: frameLength(s.duration || 5, fps),
|
|
fps,
|
|
steps: s.steps,
|
|
cfg: s.cfg,
|
|
seed: s.seed,
|
|
samplerName: 'euler',
|
|
scheduler: 'simple',
|
|
turbo: false,
|
|
duration: s.duration || 5,
|
|
filenamePrefix: prefix + '/video',
|
|
loraStack: s.loraStack,
|
|
useIdentityRefs: false
|
|
});
|
|
await ensureComfyLoraNames('video');
|
|
r.graphId = textToVideo ? 'workflow_ltx_video.json#t2v' : 'workflow_ltx_video.json';
|
|
} else {
|
|
if (q.engine !== 'minimax')
|
|
throw new Error('Studio 2 video supports MiniMax H3 and LTX (xAIGen).');
|
|
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 = identityPlan?.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, negative: stylePrompt(q.imageStyles, true), imageAName: a, imageBName: mode === 'compose' ? b : undefined, maskName: mode === 'refine' ? mask : undefined, filenamePrefix: prefix + '/image' }).graph;
|
|
applyResolvedImageSize(graph,s);
|
|
if (identityPlan) Object.assign(r,applyKleinIdentity(graph,identityPlan));
|
|
if (q.engine === 'flux')
|
|
attachHeroReference(graph, hero);
|
|
if (!identityPlan) {
|
|
r.sampleLatent=q.mode==='refine'?'masked source latent':'empty latent';
|
|
r.sampleDenoise=null;
|
|
}
|
|
if (q.engine !== 'flux' && 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) {
|
|
const stopProgress = watchProgress(r,job);
|
|
try {
|
|
const prompts = r.prompts || [r.request.promptSections, ...r.request.batch];
|
|
r.prompts = prompts;
|
|
r.request.shotSeeds ||= queueSeeds(r.request.settings,prompts.length);
|
|
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.request.settings.seed = r.request.shotSeeds[r.index];
|
|
r.files = [];
|
|
r.progress = null;
|
|
r.request.promptSections = prompts[r.index];
|
|
r.request.compiledPrompt = compilePrompt(prompts[r.index], ['video', 'extend'].includes(r.request.mode), r.request);
|
|
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);
|
|
if (!data.length || data.length < 64)
|
|
throw new Error('Comfy wrote an empty video file (<64 bytes); the job did not render. Nothing was saved.');
|
|
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;
|
|
if (!assembled.length || assembled.length < 64)
|
|
throw new Error('The assembled video file is empty (<64 bytes); the job did not render. Nothing was saved.');
|
|
const checkPath = join(tmp, 'duration-check.mp4');
|
|
writeFileSync(checkPath, assembled);
|
|
let duration = 0;
|
|
try { duration = await probeDuration(checkPath); }
|
|
finally { try { unlinkSync(checkPath); } catch { /* ignore */ } }
|
|
if (!Number.isFinite(duration) || duration <= 0)
|
|
throw new Error('The rendered video has no readable duration; the job did not render. Nothing was saved.');
|
|
asset = await saveClip({ ownerKey: r.owner, folderId: q.folderId, prompt: q.compiledPrompt, ...s, aspect: s.aspect || 'auto', hideThumbnail: false, video: assembled, originalSegment: data, sourceSegments, fps: s.fps || 24, sound: true, familyId: r.familyId, parentClipId: q.startClipId || undefined, chainIndex: r.index, comfyFilename: file.filename });
|
|
}
|
|
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, sourceStillId:r.sourceStillId,heroReferenceAttached:r.heroReferenceAttached,sampleLatent:r.sampleLatent,sampleDenoise:r.sampleDenoise,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', handoffSource: r.handoffSource || null, fitted: video && q.startClipId ? !!r.handoffFitted : null };
|
|
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);
|
|
const anchor = video && history[r.promptId].outputs?.anchor_save?.images?.[0];
|
|
if (anchor) {
|
|
if (!r.files.some((f: DiskFile)=>f.filename===anchor.filename && f.subfolder===anchor.subfolder)) r.files.push(anchor);
|
|
await persistClipAnchorFrame({ownerKey:r.owner,clipId:asset.id,fromBuffer:await downloadComfyImage(anchor)});
|
|
}
|
|
await purge(r);
|
|
metadata.purgeResult = r.purgeResult;
|
|
await attachStudio2Metadata(r.owner, asset.id, 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 {
|
|
stopProgress();
|
|
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();
|
|
delete r.finishedAt;
|
|
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);
|
|
}
|
|
}
|