From 09c96c0ae7214b5f54f240f2578c228d9fe85a58 Mon Sep 17 00:00:00 2001 From: Towsty Date: Thu, 10 Sep 2026 23:48:25 -0500 Subject: [PATCH] Fit Studio 2 inputs and report live generation progress --- server/utils/studio2/media.ts | 13 +++++++++++++ server/utils/studio2/progress.ts | 28 ++++++++++++++++++++++++++++ server/utils/studio2/runner.ts | 30 ++++++++++++++++++++---------- tests/studio2-progress.test.mjs | 19 +++++++++++++++++++ tests/studio2-runner.test.mjs | 2 +- 5 files changed, 81 insertions(+), 11 deletions(-) create mode 100644 server/utils/studio2/media.ts create mode 100644 server/utils/studio2/progress.ts create mode 100644 tests/studio2-progress.test.mjs diff --git a/server/utils/studio2/media.ts b/server/utils/studio2/media.ts new file mode 100644 index 0000000..c414dca --- /dev/null +++ b/server/utils/studio2/media.ts @@ -0,0 +1,13 @@ +import { mkdtempSync, writeFileSync, readFileSync, rmSync } from 'node:fs' +import { join } from 'node:path' +import { tmpdir } from 'node:os' +import { runFfmpeg } from '../ffmpeg' +export async function fitStill(data: Buffer, width: number, height: number) { + const dir=mkdtempSync(join(tmpdir(),'studio2-fit-')) + try { + const input=join(dir,'source'),output=join(dir,'fitted.png') + writeFileSync(input,data) + await runFfmpeg(['-y','-i',input,'-vf',`scale=${width}:${height}:force_original_aspect_ratio=decrease,pad=${width}:${height}:(ow-iw)/2:(oh-ih)/2:color=black,setsar=1`,'-frames:v','1',output]) + return readFileSync(output) + } finally { rmSync(dir,{recursive:true,force:true}) } +} diff --git a/server/utils/studio2/progress.ts b/server/utils/studio2/progress.ts new file mode 100644 index 0000000..236ee0d --- /dev/null +++ b/server/utils/studio2/progress.ts @@ -0,0 +1,28 @@ +import { comfyWsUrl } from '../comfy' +import { saveRecord } from './store' + +/** Listen only to this client's prompt; history polling remains the completion authority. */ +export function watchProgress(record: any, job: any) { + let ws: WebSocket | undefined, timer: ReturnType, stopped = false + const connect = () => { + if (stopped) return + try { + ws = new WebSocket(comfyWsUrl(job.clientId)) + ws.addEventListener('message', event => { + if (typeof event.data !== 'string') return + try { + const {type,data} = JSON.parse(event.data) + if (!record.promptId || (data?.prompt_id && data.prompt_id !== record.promptId)) return + if (type === 'progress' && Number(data.max) > 0) { + record.progress = {step:Number(data.value),maxStep:Number(data.max),percent:Math.round(Number(data.value)/Number(data.max)*100),updatedAt:Date.now()} + saveRecord(record) + } + } catch { /* Binary previews and malformed events cannot interrupt a job. */ } + }) + ws.addEventListener('error', () => { ws?.close() }) + ws.addEventListener('close', () => { if (!stopped) timer=setTimeout(connect,3000) }) + } catch { if (!stopped) timer=setTimeout(connect,3000) } + } + connect() + return () => { stopped=true;clearTimeout(timer);ws?.close() } +} diff --git a/server/utils/studio2/runner.ts b/server/utils/studio2/runner.ts index 292f056..0a12d3c 100644 --- a/server/utils/studio2/runner.ts +++ b/server/utils/studio2/runner.ts @@ -1,3 +1,6 @@ +import { watchProgress } from './progress'; +import { fitStill } from './media'; +import { resolveSize } from '~/shared/studio2/input.mjs'; import { readFileSync, mkdirSync, existsSync, unlinkSync } from 'node:fs'; import { join } from 'node:path'; import template from '../../assets/studio2_minimax_native.json'; @@ -25,6 +28,7 @@ type DiskFile = { 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 { const prefix = String(useRuntimeConfig().comfyFilenamePrefix).replace(/\/$/, '') + `/studio2/${r.id}/${r.index}`; + if (name !== 'hero') data = await fitStill(data, r.request.settings.width, r.request.settings.height); const body = new FormData(); body.append('image', new Blob([new Uint8Array(data)]), name + '.png'); body.append('subfolder', prefix); @@ -61,7 +65,10 @@ async function purge(r: any) { 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'); + const source = q.startClipId ? getClip(r.owner,q.startClipId) : (q.imageAId || q.identityStillId) ? getStill(r.owner,q.imageAId || q.identityStillId) : null; + Object.assign(s,resolveSize(s.aspect,source,s)); + saveRecord(r); + const hero = q.lockFace === false && q.lockOutfit === false ? '' : 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') @@ -101,6 +108,7 @@ async function prepareGraph(r: any) { 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; @@ -115,8 +123,9 @@ async function run(r: any, job: Job) { await new Promise(resolve => setTimeout(resolve, 2000)); } 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.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') @@ -188,12 +197,7 @@ async function run(r: any, job: Job) { 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) }); - } + 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 }; @@ -207,11 +211,16 @@ async function run(r: any, job: Job) { 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); + 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); - 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) { @@ -241,6 +250,7 @@ async function run(r: any, job: Job) { emitJob(job, { type: 'error', message: e.message }); } finally { + stopProgress(); job.saving = false; await onLiveVideoSettled(job); } diff --git a/tests/studio2-progress.test.mjs b/tests/studio2-progress.test.mjs new file mode 100644 index 0000000..bd092e7 --- /dev/null +++ b/tests/studio2-progress.test.mjs @@ -0,0 +1,19 @@ +import test from 'node:test' +import assert from 'node:assert/strict' +import fs from 'node:fs' +import ts from 'typescript' +test('progress accepts this client, rejects other prompts and closes on completion',()=>{ + const source=fs.readFileSync(new URL('../server/utils/studio2/progress.ts',import.meta.url),'utf8') + const tree=ts.createSourceFile('progress.ts',source,ts.ScriptTarget.Latest,true) + const fn=tree.statements.find(n=>ts.isFunctionDeclaration(n)&&n.name.text==='watchProgress') + const code=ts.transpileModule(fn.getText(tree).replace(/^export /,''),{compilerOptions:{target:ts.ScriptTarget.ES2022}}).outputText + let socket,saves=0 + class Socket { listeners={};constructor(){socket=this}addEventListener(name,fn){this.listeners[name]=fn}close(){this.closed=true} } + const watch=new Function('WebSocket','comfyWsUrl','saveRecord',`${code};return watchProgress`)(Socket,()=>'/fake',()=>saves++) + const record={promptId:'mine'},stop=watch(record,{clientId:'client'}) + const send=data=>socket.listeners.message({data:JSON.stringify({type:'progress',data})}) + send({prompt_id:'another',value:1,max:8});assert.equal(saves,0) + send({prompt_id:'mine',value:3,max:8});assert.equal(record.progress.step,3);assert.equal(record.progress.maxStep,8) + send({value:4,max:8});assert.equal(record.progress.percent,50) + stop();assert.equal(socket.closed,true) +}) diff --git a/tests/studio2-runner.test.mjs b/tests/studio2-runner.test.mjs index e4e90b4..8956ef1 100644 --- a/tests/studio2-runner.test.mjs +++ b/tests/studio2-runner.test.mjs @@ -16,7 +16,7 @@ function fixture(mode='video') { const job={status:'running',clientId:'client',library:{}} let id=0 const scope={ - compilePrompt,saveRecord(){},update(r,state){r.state=state;phases.push(state)},acquireSharedGpu:async()=>true,ensureComfyReady:async()=>{},emitJob(){}, + watchProgress:()=>()=>{},compilePrompt,saveRecord(){},update(r,state){r.state=state;phases.push(state)},acquireSharedGpu:async()=>true,ensureComfyReady:async()=>{},emitJob(){}, prepareGraph:async r=>{queued.push({hero:r.request.identityStillId,start:r.request.startClipId,action:r.request.promptSections.action});return {}},queuePrompt:async()=>({prompt_id:`p${++id}`}), fetchHistory:async p=>({[p]:{status:{status_str:'success'},outputs:{}}}),comfyFetch:()=>{throw new Error('Unexpected network call')}, extractVideo:()=>({filename:'video.mp4',subfolder:'preview/studio2/r',type:'output'}),extractEditedImage:()=>({filename:'image.png',subfolder:'preview/studio2/r',type:'output'}),