Fit Studio 2 inputs and report live generation progress

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}) }
+}

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<typeof setTimeout>, 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() }
+}

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<string> {
     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 });
+                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 });
-                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 };
@@ -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);
     }

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)
+})

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'}),