Restore Klein empty-latent reference sampling

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. */
+/** Follow the live conditioning chain, including optional hero/compose references. */
-export function assertKleinSample(graph, sourceId) {
+function hasSourceReference(graph, link, referenceId) {
- let link=graph['19']?.inputs?.latent_image
+ const seen=new Set()
- const visited=new Set()
+ while(Array.isArray(link) && !seen.has(String(link[0]))) {
- while(Array.isArray(link) && !visited.has(String(link[0]))) {
+  const id=String(link[0]);seen.add(id)
-  const id=String(link[0]);visited.add(id)
   const node=graph[id]
-  if(node?.class_type==='VAEEncode' && id==='11') return 'source'
+  if(id===referenceId) return node?.class_type==='ReferenceLatent' && node.inputs.latent?.[0]==='11' && graph['11']?.class_type==='VAEEncode'
-  if(node?.class_type?.startsWith('Empty')) {
+  link=node?.inputs?.conditioning
-   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
+ 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.')
  }
- throw new Error('Klein sample latent could not be verified against the source still.')
+ // 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}
 }

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.latent_image,['14',0]);assert.equal(graph['14'].class_type,'EmptyFlux2LatentImage')
- assert.deepEqual(graph['19'].inputs.sigmas,['source_sigmas',1]);assert.equal(graph.source_sigmas.inputs.denoise,.45)
+ assert.deepEqual(graph['19'].inputs.sigmas,['17',0]);assert.equal(graph.source_sigmas,undefined)
- assert.equal(record.sourceStillId,sourceId||'hero');assert.equal(record.sampleLatent,'source')
+ 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.sampleLatent,'empty+reference')
- assert.equal(record.sampleDenoise,.45)
+ assert.equal(record.sampleDenoise,null)
- assert.deepEqual(graph['19'].inputs.latent_image,['11',0])
+ 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',()=>{
+test('validate live references and reject direct source sampling',()=>{
- assert.throws(()=>assertKleinSample(template,'hero'),/empty-latent sampling was rejected/)
+ 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)
 })