From 7bb4b16c1ec49b95d9b71e62ba266e6235671c9f Mon Sep 17 00:00:00 2001 From: Towsty Date: Sat, 12 Sep 2026 11:46:56 -0500 Subject: [PATCH] Sample Klein hero identity from image latents --- components/studio2/Details.vue | 4 ++-- server/utils/studio2/runner.ts | 25 +++++++++++--------- shared/studio2/klein-identity.mjs | 33 +++++++++++++++++++++++++++ tests/studio2-identity-edit.test.mjs | 34 ++++++++++++++++++++++++---- 4 files changed, 79 insertions(+), 17 deletions(-) create mode 100644 shared/studio2/klein-identity.mjs diff --git a/components/studio2/Details.vue b/components/studio2/Details.vue index 09cdafa..314c387 100644 --- a/components/studio2/Details.vue +++ b/components/studio2/Details.vue @@ -8,13 +8,13 @@ const props=defineProps<{metadata?:any}>() const rows=computed(()=>{const m=props.metadata||{},s=m.settings||{};const time=(v:any)=>v?new Date(v).toLocaleString():null;return [ {label:'Kind / mode',value:[m.kind,m.mode].filter(Boolean).join(' / ')}, {label:'Engine',value:m.engine},{label:'Graph ID',value:m.graphId}, - {label:'Sample latent',value:m.sampleLatent},{label:'Edit denoise',value:m.sampleDenoise}, + {label:'Sample latent',value:m.sampleLatent},{label:'Sample denoise',value:m.sampleDenoise}, {label:'Source still',value:m.sourceStillId || m.imageAId || 'None'}, {label:'Hero reference attached',value:m.heroReferenceAttached==null?'Not recorded':m.heroReferenceAttached?'Yes':'No'}, {label:'Seed',value:s.seed},{label:'Steps',value:s.steps},{label:'CFG',value:s.cfg}, {label:'Size',value:s.width&&s.height?`${s.width} × ${s.height}`:null},{label:'Duration',value:s.duration!=null?`${s.duration}s`:null}, {label:'LoRAs',value:s.loraStack?.length?s.loraStack.map((l:any)=>`${l.name} · ${l.strengthModel}`).join('\n'):'None'}, - {label:'Identity still',value:m.identityStillId||'None'},{label:'Face lock',value:m.lockFace!==false?'On':'Off'},{label:'Outfit lock',value:m.lockOutfit!==false?'On':'Off'}, + {label:'Hero still',value:m.identityStillId||'None'},{label:'Face lock',value:m.lockFace!==false?'On':'Off'},{label:'Outfit lock',value:m.lockOutfit!==false?'On':'Off'}, {label:'Aspect',value:s.aspect||'auto'},{label:'FPS',value:s.fps||24},{label:'Hide input',value:m.hideInput?'On':'Off'}, {label:'Start frame',value:m.startFrameSource?JSON.stringify(m.startFrameSource):'None'}, {label:'Handoff source',value:m.handoffSource || 'None'},{label:'Handoff fitted (resampled)',value:m.fitted==null?'Not recorded':m.fitted?'Yes':'No'}, diff --git a/server/utils/studio2/runner.ts b/server/utils/studio2/runner.ts index 920c857..b7d645c 100644 --- a/server/utils/studio2/runner.ts +++ b/server/utils/studio2/runner.ts @@ -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 { 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'); diff --git a/shared/studio2/klein-identity.mjs b/shared/studio2/klein-identity.mjs new file mode 100644 index 0000000..219acb0 --- /dev/null +++ b/shared/studio2/klein-identity.mjs @@ -0,0 +1,33 @@ +import {sampleKleinSource} from './graphs.mjs' + +export function kleinIdentityPlan(request) { + if(request.engine!=='flux' || ['video','extend'].includes(request.mode)) return null + const sourceId=request.mode==='generate' ? request.identityStillId || request.imageAId || '' : request.imageAId || request.identityStillId || '' + const mode=['generate','iterate'].includes(request.mode) ? (sourceId?'edit':'generate') : request.mode + return {sourceId,mode,denoise:request.lockFace!==false?.45:.65} +} + +/** Validate the sampler's actual latent connection, not merely the presence of a reference. */ +export function assertKleinSample(graph, sourceId) { + let link=graph['19']?.inputs?.latent_image + const visited=new Set() + while(Array.isArray(link) && !visited.has(String(link[0]))) { + const id=String(link[0]);visited.add(id) + const node=graph[id] + if(node?.class_type==='VAEEncode' && id==='11') return 'source' + if(node?.class_type?.startsWith('Empty')) { + if(sourceId) throw new Error('Klein identity requires the source image latent; empty-latent sampling was rejected.') + return 'empty' + } + link=node?.class_type==='SetLatentNoiseMask'?node.inputs.samples:null + } + throw new Error('Klein sample latent could not be verified against the source still.') +} + +export function applyKleinIdentity(graph, plan) { + if(plan.sourceId && plan.mode!=='refine') sampleKleinSource(graph,plan.denoise) + // Retain the masked latent and any already lower refine denoise. + if(plan.sourceId && plan.mode==='refine') graph['17'].inputs.denoise=Math.min(graph['17'].inputs.denoise,plan.denoise) + const sampleLatent=assertKleinSample(graph,plan.sourceId) + return {sampleLatent,sampleDenoise:sampleLatent==='source'?(graph.source_sigmas?.inputs.denoise ?? graph['17']?.inputs.denoise):null} +} diff --git a/tests/studio2-identity-edit.test.mjs b/tests/studio2-identity-edit.test.mjs index 323de93..80e46a4 100644 --- a/tests/studio2-identity-edit.test.mjs +++ b/tests/studio2-identity-edit.test.mjs @@ -1,10 +1,11 @@ +import {kleinIdentityPlan,applyKleinIdentity,assertKleinSample} from '../shared/studio2/klein-identity.mjs' import {stylePrompt} from '../shared/studio2/styles.mjs' import test from 'node:test' import assert from 'node:assert/strict' import fs from 'node:fs' import ts from 'typescript' import {validateRequest,compilePrompt} from '../shared/studio2/contracts.mjs' -import {sampleKleinSource,attachHeroReference,applyResolvedImageSize} from '../shared/studio2/graphs.mjs' +import {attachHeroReference,applyResolvedImageSize} from '../shared/studio2/graphs.mjs' import {extendSelected,changeSource,restoreAsset} from '../shared/studio2/input.mjs' const template=JSON.parse(fs.readFileSync(new URL('../server/assets/klein_v2_edit.json',import.meta.url))) const source=fs.readFileSync(new URL('../server/utils/studio2/runner.ts',import.meta.url),'utf8') @@ -14,7 +15,7 @@ const js=ts.transpileModule(fn.getText(tree),{compilerOptions:{target:ts.ScriptT async function prepare(request){ const uploads=[] const scope={stylePrompt,useRuntimeConfig:()=>({comfyFilenamePrefix:'preview'}),upload:async(r,label)=>{uploads.push(label);return label+'.png'},readFileSync:()=>Buffer.from('fixture'),stillPath:()=>'',getStill:(_,id)=>({id}),resolveRequestSize:async()=>{},saveRecord(){}, - buildImageV2Workflow:params=>{const graph=structuredClone(template);graph['1'].inputs.image=params.imageAName;return {graph}},sampleKleinSource,attachHeroReference,applyResolvedImageSize,ensureComfyLoraNames:async()=>{},resolveGraphLoraNames(){}} + buildImageV2Workflow:params=>{const graph=JSON.parse(fs.readFileSync(new URL('../server/assets/klein_v2_'+params.mode+'.json',import.meta.url)));if(graph['1']) graph['1'].inputs.image=params.imageAName;return {graph}},kleinIdentityPlan,applyKleinIdentity,attachHeroReference,applyResolvedImageSize,ensureComfyLoraNames:async()=>{},resolveGraphLoraNames(){}} const run=new Function(...Object.keys(scope),`${js};return prepareGraph`)(...Object.values(scope)) const record={request,owner:'o',id:'job',index:0};return {graph:await run(record),record,uploads} } @@ -22,8 +23,8 @@ for(const sourceId of ['hero','source','']) test(`Klein edit samples source with const q=validateRequest({folderId:'f',mode:'edit',engine:'flux',imageAId:sourceId,identityStillId:'hero',lockFace:true,lockOutfit:false,prompt:'she at a pool'}) const {graph,record,uploads}=await prepare(q) assert.deepEqual(graph['19'].inputs.latent_image,['11',0]);assert.equal(graph['14'],undefined) - assert.deepEqual(graph['19'].inputs.sigmas,['source_sigmas',1]);assert.equal(graph.source_sigmas.inputs.denoise,.65) - assert.equal(record.sourceStillId,sourceId||'hero');assert.equal(record.sampleLatent,'source image latent') + assert.deepEqual(graph['19'].inputs.sigmas,['source_sigmas',1]);assert.equal(graph.source_sigmas.inputs.denoise,.45) + assert.equal(record.sourceStillId,sourceId||'hero');assert.equal(record.sampleLatent,'source') assert.equal(uploads.includes('hero'),sourceId==='source');assert.equal(!!graph.hero_load,sourceId==='source') assert.match(q.compiledPrompt,/face, glasses, hair/);assert.match(q.compiledPrompt,/do not restore the hero outfit/) }) @@ -44,3 +45,28 @@ test('Extend normalizes engine and inherits clip without replacing action or cam assert.equal(extendSelected({...form,startClipId:'clip'},null).startClipId,'clip') assert.throws(()=>validateRequest({folderId:'f',mode:'extend',engine:'minimax',identityStillId:'hero',prompt:'Go'}),/Select the clip in the family strip or Previous clip/) }) + +for (const imageAId of ['', 'hero', 'other']) test('Generate samples hero exclusively: '+imageAId,async()=>{ + const q=validateRequest({folderId:'f',mode:'generate',engine:'flux',imageAId,identityStillId:'hero',lockFace:true,lockOutfit:false,prompt:'she at a pool',settings:{loraStack:[]}}) + const {graph,record,uploads}=await prepare(q) + assert.equal(record.graphId,'klein_v2_edit.json') + assert.equal(record.sourceStillId,'hero') + assert.equal(record.sampleLatent,'source') + assert.equal(record.sampleDenoise,.45) + assert.deepEqual(graph['19'].inputs.latent_image,['11',0]) + assert.deepEqual(uploads,['source']) + assert.equal(graph.hero_load,undefined) + assert.deepEqual(q.settings.loraStack,[]) + assert.doesNotMatch(q.compiledPrompt,/Picture 1/) +}) +test('unlocked source uses .65; text-only uses empty latent',async()=>{ + const q=validateRequest({folderId:'f',mode:'edit',engine:'flux',imageAId:'source',lockFace:false,prompt:'pool'}) + assert.equal((await prepare(q)).record.sampleDenoise,.65) + const empty=validateRequest({folderId:'f',mode:'generate',engine:'flux',prompt:'pool'}) + assert.equal((await prepare(empty)).record.sampleLatent,'empty') +}) +test('reject empty sample with source and exclude other engines',()=>{ + assert.throws(()=>assertKleinSample(template,'hero'),/empty-latent sampling was rejected/) + assert.equal(kleinIdentityPlan({engine:'krea',mode:'generate',identityStillId:'hero'}),null) + assert.equal(kleinIdentityPlan({engine:'minimax',mode:'extend'}),null) +})