From e190496eb44a029c8cbc6ffd5251ab73a7219090 Mon Sep 17 00:00:00 2001 From: Towsty Date: Sat, 12 Sep 2026 12:18:14 -0500 Subject: [PATCH] Restore Klein empty-latent reference sampling --- shared/studio2/klein-identity.mjs | 44 +++++++++++++++------------- tests/studio2-identity-edit.test.mjs | 24 ++++++++------- 2 files changed, 38 insertions(+), 30 deletions(-) diff --git a/shared/studio2/klein-identity.mjs b/shared/studio2/klein-identity.mjs index 219acb0..e584a1c 100644 --- a/shared/studio2/klein-identity.mjs +++ b/shared/studio2/klein-identity.mjs @@ -1,33 +1,37 @@ -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} + return {sourceId,mode} } -/** 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) +/** Follow the live conditioning chain, including optional hero/compose references. */ +function hasSourceReference(graph, link, referenceId) { + const seen=new Set() + while(Array.isArray(link) && !seen.has(String(link[0]))) { + const id=String(link[0]);seen.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 + if(id===referenceId) return node?.class_type==='ReferenceLatent' && node.inputs.latent?.[0]==='11' && graph['11']?.class_type==='VAEEncode' + link=node?.inputs?.conditioning } - throw new Error('Klein sample latent could not be verified against the source still.') + return false +} + +export function assertKleinSample(graph, sourceId) { + const sampler=graph['19']?.inputs + const latent=graph[sampler?.latent_image?.[0]] + if(latent?.class_type==='EmptyFlux2LatentImage') { + if(!sourceId) return 'empty' + const guider=graph[sampler.guider?.[0]]?.inputs + if(hasSourceReference(graph,guider?.positive,'12') && hasSourceReference(graph,guider?.negative,'13')) return 'empty+reference' + throw new Error('Klein source reference conditioning is missing from the sampler.') + } + // Refine retains its dedicated masked-source workflow. + if(latent?.class_type==='SetLatentNoiseMask' && latent.inputs.samples?.[0]==='11' && graph['11']?.class_type==='VAEEncode') return 'masked source' + throw new Error('Klein edit must sample empty latent with source references, not the source image latent.') } 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} + return {sampleLatent,sampleDenoise:sampleLatent==='masked source' ? graph['17']?.inputs.denoise ?? null : null} } diff --git a/tests/studio2-identity-edit.test.mjs b/tests/studio2-identity-edit.test.mjs index 80e46a4..526a9f3 100644 --- a/tests/studio2-identity-edit.test.mjs +++ b/tests/studio2-identity-edit.test.mjs @@ -22,9 +22,9 @@ async function prepare(request){ for(const sourceId of ['hero','source','']) test(`Klein edit samples source with hero ${sourceId || 'as sole source'}`,async()=>{ 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,.45) - assert.equal(record.sourceStillId,sourceId||'hero');assert.equal(record.sampleLatent,'source') + assert.deepEqual(graph['19'].inputs.latent_image,['14',0]);assert.equal(graph['14'].class_type,'EmptyFlux2LatentImage') + assert.deepEqual(graph['19'].inputs.sigmas,['17',0]);assert.equal(graph.source_sigmas,undefined) + assert.equal(record.sourceStillId,sourceId||'hero');assert.equal(record.sampleLatent,'empty+reference') 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/) }) @@ -51,22 +51,26 @@ for (const imageAId of ['', 'hero', 'other']) test('Generate samples hero exclus 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.equal(record.sampleLatent,'empty+reference') + assert.equal(record.sampleDenoise,null) + assert.deepEqual(graph['19'].inputs.latent_image,['14',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()=>{ +test('reference sampling ignores face-lock denoise; 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) + assert.equal((await prepare(q)).record.sampleDenoise,null) 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/) +test('validate live references and reject direct source sampling',()=>{ + assert.equal(assertKleinSample(template,'hero'),'empty+reference') + const broken=structuredClone(template);broken['18'].inputs.positive=['9',0] + assert.throws(()=>assertKleinSample(broken,'hero'),/reference conditioning is missing/) + const copy=structuredClone(template);copy['19'].inputs.latent_image=['11',0] + assert.throws(()=>assertKleinSample(copy,'hero'),/not the source image latent/) assert.equal(kleinIdentityPlan({engine:'krea',mode:'generate',identityStillId:'hero'}),null) assert.equal(kleinIdentityPlan({engine:'minimax',mode:'extend'}),null) })