towsty/aigen
1
0
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases Wiki Activity

Fit Studio 2 inputs and report live generation progress

Browse Source
This commit is contained in:
Towsty
2026-09-10 23:48:25 -05:00
parent 5c69650ffa
commit 09c96c0ae7
5 changed files with 81 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
server/utils/studio2/media.ts
+13
View File
@@ -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
+28
View File
@@ -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
+20 -10
View File
@@ -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 });
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);
}