Sample Klein hero identity from image latents

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'},

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<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');

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

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