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

Sample Klein edits from source and bind Studio 2 extensions correctly

Browse Source
This commit is contained in:
Towsty
2026-09-11 12:23:34 -05:00
parent 8fe69272c3
commit 28b919ecee
11 changed files with 118 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files
server/utils/studio2/runner.ts
+15 -4
View File
@@ -5,7 +5,7 @@ import { queueSeeds } from '~/shared/studio2/seed.mjs';
import { readFileSync, 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 { nativeVideoGraph, attachHeroReference, applyResolvedImageSize, sampleKleinSource } 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';
@@ -68,7 +68,15 @@ async function prepareGraph(r: any) {
const load = async (id: string, label: string) => id ? upload(r, label, readFileSync(stillPath(r.owner, getStill(r.owner, id).id))) : '';
await resolveRequestSize(r.owner,q);
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');
const video=['video','extend'].includes(q.mode);
const sourceId=q.imageAId || (!video && q.engine==='flux' && ['edit','iterate'].includes(q.mode) ? q.identityStillId : '');
const usesSourceLatent=!video && q.engine==='flux' && ['edit','iterate'].includes(q.mode) && !!sourceId;
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 && q.mode!=='generate' && sourceId && sourceId===q.identityStillId);
const hero=attachHero ? await load(q.identityStillId,'hero') : '';
r.sourceStillId=sourceId || null;
r.heroReferenceAttached=!!hero || (!video && q.engine==='flux' && q.mode!=='generate' && !!sourceId && sourceId===q.identityStillId);
let graph: any;
if (['video', 'extend'].includes(q.mode)) {
if (q.engine !== 'minimax')
@@ -97,9 +105,12 @@ async function prepareGraph(r: any) {
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;
applyResolvedImageSize(graph,s);
if (usesSourceLatent) sampleKleinSource(graph,.65);
if (q.engine === 'flux')
attachHeroReference(graph, hero);
else if (hero && ['edit', 'compose', 'iterate'].includes(q.mode))
r.sampleLatent=usesSourceLatent?'source image latent':q.mode==='refine'?'masked source latent':'empty latent';
r.sampleDenoise=usesSourceLatent ? .65 : 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');
@@ -212,7 +223,7 @@ async function run(r: any, job: Job) {
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' };
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' };
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);