73 lines
6.0 KiB
JavaScript
73 lines
6.0 KiB
JavaScript
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 {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')
|
|
const tree=ts.createSourceFile('runner.ts',source,ts.ScriptTarget.Latest,true)
|
|
const fn=tree.statements.find(n=>ts.isFunctionDeclaration(n)&&n.name.text==='prepareGraph')
|
|
const js=ts.transpileModule(fn.getText(tree),{compilerOptions:{target:ts.ScriptTarget.ES2022}}).outputText
|
|
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=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}
|
|
}
|
|
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.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/)
|
|
})
|
|
test('image requests omit video sections while form keeps them',()=>{
|
|
const sections={action:'sitting in a restaurant',cameraBeats:[{moves:['hold']}],soundscapeChips:['footsteps'],dialogue:{mode:'whisper',line:'Hello'},musicChip:'low underscore'}
|
|
const q=validateRequest({folderId:'f',mode:'edit',engine:'flux',identityStillId:'hero',promptSections:sections})
|
|
assert.doesNotMatch(q.compiledPrompt,/tripod|hold|soundscape|music|Hello|Picture 1/i)
|
|
assert.deepEqual(q.promptSections,{action:sections.action});assert.equal(sections.cameraBeats.length,1)
|
|
assert.equal(compilePrompt(sections,false,{engine:'krea',identityStillId:'hero'}),sections.action)
|
|
})
|
|
test('Extend normalizes engine and inherits clip without replacing action or camera',()=>{
|
|
const form={mode:'edit',engine:'flux',imageAId:'still',promptSections:{action:'New action',cameraBeats:[{moves:['hold']}]},settings:{loraStack:[]}}
|
|
const clip={id:'clip',kind:'clip',studio2:{engine:'minimax',identityStillId:'hero',lockOutfit:false,promptSections:{action:'Old action'},settings:{duration:8,aspect:'1:1',width:960,height:960,fps:24,turbo:false,steps:8,cfg:1.5,loraStack:[{name:'a',strengthModel:.6,strengthClip:.3}]}}}
|
|
const next=extendSelected(form,clip);assert.equal(next.engine,'minimax');assert.equal(next.startClipId,'clip');assert.deepEqual(next.settings,clip.studio2.settings)
|
|
assert.deepEqual(next.promptSections,form.promptSections);assert.equal(next.lockOutfit,false)
|
|
assert.equal(changeSource({...next,imageAId:'new-still'}).startClipId,'clip');assert.equal(changeSource(next).mode,'extend')
|
|
assert.deepEqual(restoreAsset(form,clip).promptSections,form.promptSections)
|
|
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)
|
|
})
|