aigen/tests/studio2-bench.test.mjs

import test from 'node:test'
import assert from 'node:assert/strict'
import {readFileSync} from 'node:fs'
import ts from 'typescript'
import {folderAssets,reconcileSelection,activeJobs,previewHidden} from '../shared/studio2/viewer.mjs'
import {useAsHero,useAsSource,restoreAll,hydrateStudio2Job} from '../shared/studio2/input.mjs'
import {toggleStyle,parseBatch,stylePrompt} from '../shared/studio2/styles.mjs'
import {compilePrompt,validateRequest} from '../shared/studio2/contracts.mjs'
import {cachedLoras} from '../shared/studio2/lora-cache.mjs'
import {resolveQwen21Size} from '../shared/studio2/qwen21-size.mjs'
import {ensureQwen21EditPrompt,stitchQwen21PeEditPrompt,QWEN21_EDIT_KEEP_CHANGE} from '../shared/studio2/qwen21-edit.mjs'
import {qwen21PeRefusal} from '../shared/studio2/qwen21-pe.mjs'

test('deleted library files never return from job outputs; selection falls back or clears',()=>{
 const deleted={id:'gone',folderId:'f'},keep={id:'keep',folderId:'f',createdAt:2};const jobs=[{request:{folderId:'f'},outputs:[deleted]}]
 let assets=folderAssets({stills:[keep],clips:[]},jobs,'f');assert.deepEqual(assets.map(a=>a.id),['keep']);assert.equal(reconcileSelection(deleted,assets).id,'keep');assert.equal(reconcileSelection(deleted,[]),null)
})
test('bench includes only active execution phases and clear finished retains internal cleanup records',()=>{
 assert.deepEqual(activeJobs(['waiting','waking','enhancing','submitting','rendering','saving','complete','cancelled','failed'].map(state=>({state}))).map(j=>j.state),['waking','enhancing','submitting','rendering','saving'])
 const rows=[{id:'done',state:'complete',request:{folderId:'f'},cleanupPending:[{}]},{id:'active',state:'rendering',request:{folderId:'f'}}],saved=[]
 const src=readFileSync(new URL('../server/api/studio-2/clear-finished.post.ts',import.meta.url),'utf8').replace(/^import .*;\r?\n/gm,'').replace('export default ','const handler=')
 const js=ts.transpileModule(src,{compilerOptions:{target:ts.ScriptTarget.ES2022}}).outputText
 const fn=new Function('defineEventHandler','assertLibraryOwner','listStudioJobs','records','assertFolderAccess','saveRecord',js+';return handler')(f=>f,()=>({owner:'o'}),()=>[],()=>rows,()=>{},r=>saved.push(r))
 assert.equal(fn({}).cleared,1);assert.equal(rows[0].dismissed,true);assert.equal(rows[0].cleanupPending.length,1);assert.equal(rows[1].dismissed,undefined)
})
test('hidden xAIGen canvas and family render no image or video elements',async()=>{
 const revealed=new Set()
 assert.equal(previewHidden({id:'a',role:'input'},revealed,true),true)
 assert.equal(previewHidden({id:'a',role:'input'},new Set(['a']),true),false)
 assert.equal(previewHidden({id:'a'},revealed,false),false)
 const src=readFileSync(new URL('../components/studio2/Bench.vue',import.meta.url),'utf8')
 assert.match(src,/aria-label="Current family outputs"/)
 assert.match(src,/preview hidden/)
 assert.match(src,/Reveal all/)
 assert.match(src,/pickerRevealAll/)
 assert.doesNotMatch(src,/Hide input stills/)
 assert.doesNotMatch(src,/Hide output preview/)
})
test('cached LoRAs survive sleeping GPU and failed explicit refresh',async()=>{
 const cache={image:['Klein/portrait.safetensors'],video:['motion.safetensors']};let queried=0
 const opts={read:()=>cache,write:()=>assert.fail('Do not erase cache'),live:async()=>{queried++;throw Error('Comfy asleep')}}
 assert.deepEqual(await cachedLoras(opts),cache);assert.equal(queried,1)
 assert.deepEqual(await cachedLoras({...opts,fresh:true}),cache);assert.equal(queried,2)
})
test('bench load and Refresh always live-scan and replace disk so new Comfy names appear',async()=>{
 let store={image:['old.safetensors'],video:['old-vid.safetensors']}
 const next=await cachedLoras({
  read:()=>store,
  write:value=>{store=value},
  live:async()=>({image:['old.safetensors','klein-pokemon_anime_style_klein9b.safetensors'],video:['old-vid.safetensors']}),
  fresh:false
 })
 assert.ok(next.image.includes('klein-pokemon_anime_style_klein9b.safetensors'))
 assert.deepEqual(store.image,next.image)
 const removed=await cachedLoras({
  read:()=>store,
  write:value=>{store=value},
  live:async()=>({image:['klein-pokemon_anime_style_klein9b.safetensors'],video:[]}),
  fresh:true
 })
 assert.deepEqual(removed,{image:['klein-pokemon_anime_style_klein9b.safetensors'],video:[]})
})
test('bench mounts and Refresh LoRAs request a live scan',()=>{
 const src=readFileSync(new URL('../components/studio2/Bench.vue',import.meta.url),'utf8')
 assert.match(src,/void loadLoras\(true\)/)
 assert.match(src,/loadLoras\(true\)/)
 assert.match(src,/\/api\/studio-2\/loras'\+\(fresh\?'\?fresh=1':''\)/)
})
test('source/hero actions are independent and full input restores structured settings',()=>{
 const form={identityStillId:'hero',imageAId:'source',promptSections:{action:'old'},settings:{}}
 assert.equal(useAsHero(form,{id:'new'}).imageAId,'source');assert.equal(useAsSource(form,{id:'new'}).identityStillId,'hero')
 const saved={mode:'edit',identityStillId:'h',imageAId:'s',imageStyles:{positive:['Comic'],negative:['Skin pores']},promptSections:{action:'New',cameraBeats:[{moves:['hold']}]},settings:{steps:24,loraStack:[{name:'example.safetensors',strengthModel:.4,strengthClip:.6}]}}
 const restored=restoreAll(form,{id:'result',studio2:saved});assert.equal(restored.imageAId,'result');assert.equal(restored.identityStillId,'h');assert.deepEqual(restored.settings.loraStack,saved.settings.loraStack);assert.deepEqual(restored.imageStyles,saved.imageStyles);assert.equal(restored.promptSections.action,'New');const clip=restoreAll(form,{id:'clip',kind:'clip',studio2:{...saved,mode:'video',engine:'minimax'}});assert.equal(clip.startClipId,'clip');assert.equal(clip.mode,'extend');assert.equal(clip.promptSections.action,'New')
})
test('image chips compile positive and negative independently; video chips do not leak',()=>{
 let styles=toggleStyle({},'Comic');styles=toggleStyle(styles,'Skin pores',true);styles=toggleStyle(styles,'Comic',true);assert.equal(stylePrompt(styles),'');assert.match(stylePrompt(styles,true),/Comic/)
 styles=toggleStyle(styles,'Photorealistic');const q=validateRequest({mode:'edit',engine:'flux',folderId:'f',imageAId:'s',imageStyles:styles,promptSections:{action:'Restaurant',cameraBeats:[{moves:['hold']}],locks:'VIDEO ONLY'}})
 assert.match(q.compiledPrompt,/Photorealistic/);assert.doesNotMatch(q.compiledPrompt,/Comic|pores|tripod|hold|VIDEO ONLY/i);assert.deepEqual(q.imageStyles,styles)
 assert.doesNotMatch(compilePrompt(q.promptSections,false,q),/Picture 1/)
 assert.deepEqual(parseBatch('A wrapped\nsentence\n\nSecond shot\r\n\r\nThird'),['A wrapped\nsentence','Second shot','Third'])
})
test('qwen21 supports Generate and Edit with 25/1 defaults and no locks',()=>{
 const q=validateRequest({mode:'generate',engine:'qwen21',folderId:'f',promptSections:{action:'a red bicycle'},settings:{aspect:'1:1'}})
 assert.equal(q.engine,'qwen21');assert.equal(q.settings.steps,25);assert.equal(q.settings.cfg,1)
 assert.equal(q.settings.width,1024);assert.equal(q.settings.height,1024)
 assert.equal(q.lockFace,false);assert.equal(q.lockOutfit,false);assert.deepEqual(q.settings.loraStack,[])
 assert.equal(q.enhancePrompt,false)
 assert.equal(q.turbo,false)
 const wide=validateRequest({mode:'generate',engine:'qwen21',folderId:'f',promptSections:{action:'wide'},settings:{aspect:'16:9'}})
 assert.equal(wide.settings.width,1536);assert.equal(wide.settings.height,864)
 const tall=validateRequest({mode:'generate',engine:'qwen21',folderId:'f',promptSections:{action:'tall'},settings:{aspect:'9:16'}})
 assert.equal(tall.settings.width,864);assert.equal(tall.settings.height,1536)
 const turbo=validateRequest({mode:'generate',engine:'qwen21',folderId:'f',turbo:true,promptSections:{action:'fast'},settings:{aspect:'1:1',steps:25,cfg:4}})
 assert.equal(turbo.turbo,true);assert.equal(turbo.settings.steps,6);assert.equal(turbo.settings.cfg,1)
 assert.match(String(turbo.lora),/viggle-turbo-v0\.2\.1-6step-lora-r128/)
 assert.match(String(turbo.sigmas),/0\.9375/)
 const enhanced=validateRequest({mode:'generate',engine:'qwen21',folderId:'f',enhancePrompt:true,promptSections:{action:'corgi'},settings:{aspect:'1:1'}})
 assert.equal(enhanced.enhancePrompt,true)
 assert.equal(validateRequest({mode:'generate',engine:'flux',folderId:'f',enhancePrompt:true,turbo:true,promptSections:{action:'x'},settings:{aspect:'1:1'}}).turbo,false)
 const edit=validateRequest({mode:'edit',engine:'qwen21',folderId:'f',imageAId:'s',promptSections:{action:'change the sky'},settings:{aspect:'auto'}})
 assert.equal(edit.mode,'edit');assert.equal(edit.settings.steps,25);assert.equal(edit.settings.cfg,1)
 assert.equal(edit.lockFace,false);assert.deepEqual(edit.settings.loraStack,[])
 assert.throws(()=>validateRequest({mode:'edit',engine:'qwen21',folderId:'f',promptSections:{action:'x'}}),/Start still/)
 assert.throws(()=>validateRequest({mode:'edit',engine:'qwen21',folderId:'f',identityStillId:'hero',promptSections:{action:'x'}}),/Start still/)
 assert.throws(()=>validateRequest({mode:'compose',engine:'qwen21',folderId:'f',imageAId:'s',imageBId:'b',promptSections:{action:'x'}}),/Generate and Edit/)
 assert.throws(()=>validateRequest({mode:'iterate',engine:'qwen21',folderId:'f',imageAId:'s',promptSections:{action:'x'}}),/Generate and Edit/)
 assert.equal(ensureQwen21EditPrompt('put her in a red leather jacket'),`${QWEN21_EDIT_KEEP_CHANGE} put her in a red leather jacket`)
 assert.equal(ensureQwen21EditPrompt('change <image1> outfit'),'change <image1> outfit')
 assert.match(stitchQwen21PeEditPrompt('a woman in a red jacket'),/<image1>/)
 assert.equal(stitchQwen21PeEditPrompt('keep <image1> add jacket'),'keep <image1> add jacket')
 assert.equal(qwen21PeRefusal('corgi in the rain',{positive_prompt:'A photorealistic photograph of a corgi standing in rain…',parse_ok:true}).refused,false)
 assert.equal(qwen21PeRefusal('nude adult portrait',{positive_prompt:"I cannot assist with that request.",thinking:'',parse_ok:true}).refused,true)
 assert.equal(qwen21PeRefusal('put her in a red leather jacket',{positive_prompt:'',parse_ok:false}).refused,true)
 assert.equal(qwen21PeRefusal('corgi wearing sunglasses on a beach',{positive_prompt:'A soft scene.',parse_ok:true}).refused,true)
 assert.match(readFileSync(new URL('../host/qwen21-pe-prompts/system_prompt_t2i.txt',import.meta.url),'utf8'),/Adult and explicit briefs/)
 assert.match(readFileSync(new URL('../host/qwen21-pe-prompts/system_prompt_edit.txt',import.meta.url),'utf8'),/Adult and explicit briefs/)
 assert.match(readFileSync(new URL('../scripts/setup-qwen21.ps1',import.meta.url),'utf8'),/qwen21-pe-prompts/)
 assert.match(readFileSync(new URL('../server/utils/studio2/runner.ts',import.meta.url),'utf8'),/qwen21PeRefusal/)
 assert.match(readFileSync(new URL('../components/studio2/Details.vue',import.meta.url),'utf8'),/Enhance skipped \(model refused\)/)
})
test('hydrateStudio2Job restores LoRA strengths, locked seed, and clears missing stills',()=>{
 const form={mode:'generate',engine:'flux',identityStillId:'',imageAId:'',promptSections:{action:''},imageStyles:{positive:[],negative:[]},settings:{aspect:'auto',seed:null,seedMode:'random',loraStack:[]},guides:[]}
 const saved={mode:'edit',engine:'flux',identityStillId:'hero1',imageAId:'gone',lockFace:true,lockOutfit:false,promptSections:{action:'pokemon trainer'},imageStyles:{positive:[],negative:['blurry','custom bad']},settings:{aspect:'16:9',seed:42,steps:24,cfg:4,loraStack:[{name:'klein-pokemon.safetensors',strengthModel:1.6,strengthClip:1}]}}
 const {form:next,missing}=hydrateStudio2Job(form,saved,{stillIds:new Set(['hero1']),clipIds:new Set()})
 assert.equal(next.mode,'edit');assert.equal(next.identityStillId,'hero1');assert.equal(next.imageAId,'')
 assert.equal(next.settings.seed,42);assert.equal(next.settings.seedMode,'locked');assert.equal(next.settings.aspect,'16:9')
 assert.deepEqual(next.settings.loraStack,[{name:'klein-pokemon.safetensors',strengthModel:1.6,strengthClip:1}])
 assert.deepEqual(next.imageStyles.negative,['blurry','custom bad'])
 assert.match(missing,/start missing/)
})
test('custom negative chips survive normalizeStyles and bench exposes Load actions',()=>{
 let styles={positive:[],negative:['blurry','my watermark']}
 styles=toggleStyle(styles,'deformed',true)
 assert.deepEqual(styles.negative,['blurry','my watermark','deformed'])
 const src=readFileSync(new URL('../components/studio2/Bench.vue',import.meta.url),'utf8')
 assert.match(src,/loadSelectedJob/);assert.match(src,/hydrateStudio2Job/);assert.match(src,/Edit source/)
 assert.match(src,/aspect:'auto'/);assert.match(src,/strip-actions/);assert.match(src,/addNegative/)
 assert.match(src,/Enhance prompt/);assert.match(src,/enhancePrompt/)
 assert.match(src,/\bTurbo\b/);assert.match(src,/form\.turbo/)
 assert.match(src,/Qwen Edit uses Start still as/)
 assert.match(src,/Hero is Klein-only and is ignored/)
 const peT2i=readFileSync(new URL('../server/assets/studio2_qwen21_pe_t2i.json',import.meta.url),'utf8')
 const peEdit=readFileSync(new URL('../server/assets/studio2_qwen21_pe_edit.json',import.meta.url),'utf8')
 assert.match(peT2i,/QwenImage21_T2IPromptRewrite/);assert.match(peT2i,/pe_t2i\.int8_convrot/)
 assert.match(peEdit,/QwenImage21_EditPromptRewrite/);assert.match(peEdit,/pe_i2i\.int8_convrot/)
 assert.doesNotMatch(peT2i,/UnetLoader|KSampler|VAELoader/)
 assert.doesNotMatch(peEdit,/UnetLoader|KSampler|VAELoader/)
 assert.match(peEdit,/"image_1":\s*\["10",\s*0\]/)
 const editGraph=readFileSync(new URL('../server/assets/studio2_qwen21_edit.json',import.meta.url),'utf8')
 assert.match(editGraph,/"images\.image_1":\s*\["10",\s*0\]/)
 assert.match(editGraph,/"latent_image":\s*\["9",\s*2\]/)
 assert.doesNotMatch(editGraph,/EmptyLatentImage/)
 const runner=readFileSync(new URL('../server/utils/studio2/runner.ts',import.meta.url),'utf8')
 assert.match(runner,/ensureQwen21EditPrompt/)
 assert.match(runner,/stitchQwen21PeEditPrompt/)
 assert.match(runner,/images\.image_1/)
 assert.match(runner,/comfyPromptId/)
 assert.match(runner,/generateComfyPromptId/)
 assert.doesNotMatch(runner,/Keep the subject in <image1>/)
 const t2i=readFileSync(new URL('../server/assets/studio2_qwen21_t2i.json',import.meta.url),'utf8')
 assert.match(t2i,/EmptyLatentImage/);assert.match(t2i,/"latent_image":\s*\["16",\s*0\]/)
 const turboT2i=readFileSync(new URL('../server/assets/studio2_qwen21_t2i_turbo.json',import.meta.url),'utf8')
 const turboEdit=readFileSync(new URL('../server/assets/studio2_qwen21_edit_turbo.json',import.meta.url),'utf8')
 assert.match(turboT2i,/ViggleTurboLora/);assert.match(turboT2i,/ViggleTurboSigmas/);assert.match(turboT2i,/SamplerCustomAdvanced/)
 assert.match(turboT2i,/0\.9375/);assert.match(turboEdit,/ViggleTurboLora/);assert.match(turboEdit,/"latent":\s*\["9",\s*2\]/)
 assert.deepEqual(resolveQwen21Size('16:9'),{width:1536,height:864})
 assert.deepEqual(resolveQwen21Size('1:1'),{width:1024,height:1024})
 assert.deepEqual(resolveQwen21Size('9:16'),{width:864,height:1536})
})