Sample Klein hero identity from image latents

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)
+})