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

Sample Klein hero identity from image latents

Browse Source
This commit is contained in:
Towsty
2026-09-12 11:46:56 -05:00
parent 3da3656140
commit 7bb4b16c1e
4 changed files with 79 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
server/utils/studio2/runner.ts
+14 -11
View File
@@ -1,3 +1,4 @@
import { kleinIdentityPlan, applyKleinIdentity } from '~/shared/studio2/klein-identity.mjs';
import { stylePrompt } from '~/shared/studio2/styles.mjs';
import { purge } from './cleanup';
import { watchProgress } from './progress';
@@ -7,7 +8,7 @@ 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, sampleKleinSource } from '~/shared/studio2/graphs.mjs';
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';
@@ -60,16 +61,16 @@ async function upload(r: any, name: string, data: Buffer): Promise<string> {
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))) : '';
await resolveRequestSize(r.owner,q);
saveRecord(r);
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 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 && q.mode!=='generate' && sourceId && sourceId===q.identityStillId);
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' && q.mode!=='generate' && !!sourceId && sourceId===q.identityStillId);
r.heroReferenceAttached=!!hero || (!video && q.engine==='flux' && !!sourceId && sourceId===q.identityStillId);
let graph: any;
if (['video', 'extend'].includes(q.mode)) {
@@ -95,16 +96,18 @@ async function prepareGraph(r: any) {
r.graphId = 'studio2_minimax_native.json';
}
else {
const mode = q.mode === 'iterate' ? (a ? 'edit' : 'generate') : q.mode;
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 (usesSourceLatent) sampleKleinSource(graph,.65);
if (identityPlan) Object.assign(r,applyKleinIdentity(graph,identityPlan));
if (q.engine === 'flux')
attachHeroReference(graph, hero);
r.sampleLatent=usesSourceLatent?'source image latent':q.mode==='refine'?'masked source latent':'empty latent';
r.sampleDenoise=usesSourceLatent ? .65 : null;
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');