From 28b919ecee5c07cca742f6f3e542d5e944ba9142 Mon Sep 17 00:00:00 2001
From: Towsty <ianjmcarr@gmail.com>
Date: Fri, 11 Sep 2026 12:23:34 -0500
Subject: [PATCH] Sample Klein edits from source and bind Studio 2 extensions
 correctly

---
 components/studio2/Bench.vue         | 16 +++++-----
 components/studio2/Details.vue       |  3 ++
 server/utils/studio2/runner.ts       | 19 +++++++++---
 shared/studio2/contracts.mjs         | 21 +++++++++----
 shared/studio2/graphs.mjs            | 10 +++++++
 shared/studio2/input.mjs             | 21 +++++++++----
 shared/studio2/upload.mjs            |  2 +-
 tests/studio2-controls.test.mjs      |  4 +--
 tests/studio2-identity-edit.test.mjs | 45 ++++++++++++++++++++++++++++
 tests/studio2-upload.test.mjs        |  2 +-
 tests/studio2.test.mjs               |  2 +-
 11 files changed, 118 insertions(+), 27 deletions(-)
 create mode 100644 tests/studio2-identity-edit.test.mjs

diff --git a/components/studio2/Bench.vue b/components/studio2/Bench.vue
index b290609..d751762 100644
--- a/components/studio2/Bench.vue
+++ b/components/studio2/Bench.vue
@@ -36,10 +36,12 @@
     <div class="flex justify-between"><h2 class="text-xl font-semibold">Prompt</h2><button class="text-sm text-amber-200" @click="advanced=!advanced">{{ advanced?'Simple':'Advanced' }}</button></div>
     <label class="block">Action<textarea v-model="form.promptSections.action" rows="4" class="mt-1 w-full rounded-xl bg-zinc-900 p-3" placeholder="What happens in this shot?" /></label>
     <Studio2PromptControls v-if="video" v-model="form.promptSections" />
-    <label class="block">Locks<textarea v-model="form.promptSections.locks" rows="2" class="mt-1 w-full rounded-xl bg-zinc-900 p-3" /></label>
+    <label v-if="video" class="block">Locks<textarea v-model="form.promptSections.locks" rows="2" class="mt-1 w-full rounded-xl bg-zinc-900 p-3" /></label>
     <details class="rounded-xl border border-white/10 p-3"><summary>{{ video?'More shots':'Iteration prompts' }}</summary><textarea v-model="batchText" rows="5" class="mt-2 w-full rounded-xl bg-zinc-900 p-3" placeholder="One action prompt per line. Hero stays attached to every output." /></details>
     <button class="text-sm text-amber-200" @click="settingsOpen=true">Generation settings</button>
     <p v-if="error" role="alert" class="text-sm text-red-300">{{ error }}</p>
+    <p v-if="form.mode==='extend'" class="text-sm text-amber-200">{{ form.startClipId ? 'Extend · '+(clips.find((c:any)=>c.id===form.startClipId)?.name || form.startClipId)+' · last frame' : 'Select the clip in the family strip or Previous clip.' }}</p>
+    <p v-if="!video && form.engine==='krea'" class="text-sm text-zinc-400">Krea does not attach the hero identity reference.</p>
     <button :disabled="submitting" @click="queue" class="w-full rounded-2xl bg-amber-300 p-4 font-semibold text-black disabled:opacity-50">{{ submitting?'Queuing…':'Queue' }}</button>
     <Studio2JobProgress v-for="job in folderJobs.slice(0,6)" :key="job.id" :job="job" :now="now" @select="focusJob(job)" @cancel="cancel(job)" />
     <details v-if="advanced"><summary>Compiled prompt</summary><pre class="whitespace-pre-wrap text-xs text-zinc-400">{{ compiled }}</pre></details>
@@ -55,7 +57,7 @@
 </template>
 <script setup lang="ts">
 import {compilePrompt,sectionsFrom} from '~/shared/studio2/contracts.mjs'
-import {aspectOptions,resolveSize,bindClip,restoreAsset,extendSelected} from '~/shared/studio2/input.mjs'
+import {aspectOptions,resolveSize,bindClip,restoreAsset,extendSelected,changeSource} from '~/shared/studio2/input.mjs'
 import {folderAssets,newestOutput,assetLabel} from '~/shared/studio2/viewer.mjs'
 import {applyUploadedStill} from '~/shared/studio2/upload.mjs'
 const uploading=ref(false),uploadMessage=ref('')
@@ -76,14 +78,14 @@ const sourceAsset=computed(()=>form.startClipId?clips.value.find((c:any)=>c.id==
 const sourceSize=computed(()=>measured[sourceAsset.value?.id] || sourceAsset.value)
 const resolvedSize=computed(()=>resolveSize(form.settings.aspect,sourceSize.value))
 watch(()=>sourceAsset.value?.id,id=>{const asset=sourceAsset.value;if(!import.meta.client||!asset||measured[id])return;const media=asset.id===form.startClipId?document.createElement('video'):new Image();if(media.tagName==='VIDEO')(media as HTMLVideoElement).preload='metadata';const done=()=>{measured[asset.id]={width:(media as HTMLVideoElement).videoWidth||(media as HTMLImageElement).naturalWidth,height:(media as HTMLVideoElement).videoHeight||(media as HTMLImageElement).naturalHeight}};media.addEventListener(media.tagName==='VIDEO'?'loadedmetadata':'load',done,{once:true});media.src=media.tagName==='VIDEO'?`/api/library/clips/${asset.id}/video`:stillUrl(asset.id)})
-function chooseTask(task:string){if(task==='extend'){Object.assign(form,extendSelected(form,selected.value?.kind==='clip'?selected.value:null));return}const was=video.value;form.mode=task;form.startClipId='';if(video.value!==was){form.engine=video.value?'minimax':'flux';form.settings.steps=video.value?8:24;form.settings.cfg=video.value?1.5:4}}
-function chooseClip(id:string){if(!id){form.startClipId='';return}const clip=clips.value.find((c:any)=>c.id===id);if(clip)Object.assign(form,bindClip(form,clip))}
-function selectOutput(item:any){selected.value=item}
-function sourceChanged(){form.startClipId='';if(form.mode==='extend')form.mode='video';form.settings.aspect='auto'}
+function chooseTask(task:string){if(task==='extend'){Object.assign(form,extendSelected(form,selected.value?.kind==='clip'?selected.value:null,useRuntimeConfig().public.studio==='xaigen'));return}const was=video.value;form.mode=task;form.startClipId='';if(video.value!==was){form.engine=video.value?'minimax':'flux';form.settings.steps=video.value?8:24;form.settings.cfg=video.value?1.5:4}}
+function chooseClip(id:string){if(!id){form.startClipId='';return}const clip=clips.value.find((c:any)=>c.id===id);if(clip)Object.assign(form,bindClip(form,clip,useRuntimeConfig().public.studio==='xaigen'))}
+function selectOutput(item:any){selected.value=item;if(form.mode==='extend' && item?.kind==='clip')Object.assign(form,extendSelected(form,item,useRuntimeConfig().public.studio==='xaigen'))}
+function sourceChanged(){Object.assign(form,changeSource(form))}
 function restore(item:any){Object.assign(form,restoreAsset(form,item));form.guides=form.guides||[];batchText.value=''}
 function focusJob(job:any){selectOutput(job.outputs?.at(-1)||null)}
 async function loadLoras(){loraNames.value=await $fetch('/api/loras').catch(()=>({image:[],video:[]}))}
-async function queue(){submitting.value=true;error.value='';try{if(form.settings.seedMode!=='locked')form.settings.seed=randomSeed();const queued=await $fetch<any>('/api/studio-2/jobs',{method:'POST',body:{...form,settings:{...form.settings,...resolvedSize.value},batch:batchText.value.split('\n').map(s=>s.trim()).filter(Boolean)}});form.settings.seed=queued.seed;form.settings.width=queued.width;form.settings.height=queued.height;await refresh()}catch(e:any){error.value=e.data?.statusMessage||e.message}finally{submitting.value=false}}
+async function queue(){submitting.value=true;error.value='';try{if(form.settings.seedMode!=='locked')form.settings.seed=randomSeed();const queued=await $fetch<any>('/api/studio-2/jobs',{method:'POST',body:{...form,promptSections:video.value?form.promptSections:{action:form.promptSections.action},settings:{...form.settings,...resolvedSize.value},batch:batchText.value.split('\n').map(s=>s.trim()).filter(Boolean)}});form.settings.seed=queued.seed;form.settings.width=queued.width;form.settings.height=queued.height;await refresh()}catch(e:any){error.value=e.data?.statusMessage||e.message}finally{submitting.value=false}}
 async function cancel(job:any){await $fetch('/api/studio-2/cancel',{method:'POST',body:{id:job.id}});await refresh()}
 async function upload(event:Event,target:'source'|'hero'){
  const input=event.target as HTMLInputElement,file=input.files?.[0];if(!file)return
diff --git a/components/studio2/Details.vue b/components/studio2/Details.vue
index 0db6b90..a0ea6bf 100644
--- a/components/studio2/Details.vue
+++ b/components/studio2/Details.vue
@@ -8,6 +8,9 @@ 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:'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'},
diff --git a/server/utils/studio2/runner.ts b/server/utils/studio2/runner.ts
index aae292e..6c453ce 100644
--- a/server/utils/studio2/runner.ts
+++ b/server/utils/studio2/runner.ts
@@ -5,7 +5,7 @@ import { queueSeeds } from '~/shared/studio2/seed.mjs';
 import { readFileSync, mkdirSync, existsSync, unlinkSync } from 'node:fs';
 import { join } from 'node:path';
 import template from '../../assets/studio2_minimax_native.json';
-import { nativeVideoGraph, attachHeroReference, applyResolvedImageSize } from '~/shared/studio2/graphs.mjs';
+import { nativeVideoGraph, attachHeroReference, applyResolvedImageSize, sampleKleinSource } 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';
@@ -68,7 +68,15 @@ async function prepareGraph(r: any) {
     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 hero = q.lockFace === false && q.lockOutfit === false ? '' : await load(q.identityStillId, 'hero'), a = await load(q.imageAId, 'source'), b = await load(q.imageBId, 'compose'), mask = await load(q.maskId, 'mask');
+    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 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 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);
+
     let graph: any;
     if (['video', 'extend'].includes(q.mode)) {
         if (q.engine !== 'minimax')
@@ -97,9 +105,12 @@ async function prepareGraph(r: any) {
         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, 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 (q.engine === 'flux')
             attachHeroReference(graph, hero);
-        else if (hero && ['edit', 'compose', 'iterate'].includes(q.mode))
+        r.sampleLatent=usesSourceLatent?'source image latent':q.mode==='refine'?'masked source latent':'empty latent';
+        r.sampleDenoise=usesSourceLatent ? .65 : 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');
         resolveGraphLoraNames(graph, 'image');
@@ -212,7 +223,7 @@ async function run(r: any, job: Job) {
             saveRecord(r);
             r.savedAt = Date.now();
             r.wallTime = (r.savedAt - r.startedAt) / 1000;
-            const metadata = { ...structuredClone(q), id: r.id, kind: video ? 'video' : 'image', graphId: r.graphId, promptId: r.promptId, queuedAt: r.queuedAt, startedAt: r.startedAt, savedAt: r.savedAt, gpuSeconds: r.gpuSeconds, wallTime: r.wallTime, outputWidth: asset.width || s.width, outputHeight: asset.height || s.height, purgeResult: 'Left on host' };
+            const metadata = { ...structuredClone(q), id: r.id, sourceStillId:r.sourceStillId,heroReferenceAttached:r.heroReferenceAttached,sampleLatent:r.sampleLatent,sampleDenoise:r.sampleDenoise,kind: video ? 'video' : 'image', graphId: r.graphId, promptId: r.promptId, queuedAt: r.queuedAt, startedAt: r.startedAt, savedAt: r.savedAt, gpuSeconds: r.gpuSeconds, wallTime: r.wallTime, outputWidth: asset.width || s.width, outputHeight: asset.height || s.height, purgeResult: 'Left on host' };
             await attachStudio2Metadata(r.owner, asset.id, metadata);
             if (!r.outputs.some((o:any)=>o.id===asset.id)) r.outputs.push({id:asset.id,kind:video?'clip':'still',studio2:metadata});
             saveRecord(r);
diff --git a/shared/studio2/contracts.mjs b/shared/studio2/contracts.mjs
index a359d6b..51d69fc 100644
--- a/shared/studio2/contracts.mjs
+++ b/shared/studio2/contracts.mjs
@@ -10,7 +10,15 @@ export function appearanceLock({lockFace = true, lockOutfit = true} = {}) {
   return 'The start frame is the only appearance reference for person and clothing, as well as pose, place, and camera. Do not use the hero as an appearance reference.'
 }
 export function compilePrompt(value, video = false, locks = {}) {
-  return [compileSections(value), video ? appearanceLock(locks) : ''].filter(Boolean).join('\n\n')
+  if (video) return [compileSections(value),appearanceLock(locks)].filter(Boolean).join('\n\n')
+  const hasImage=!!(locks.imageAId || locks.identityStillId)
+  const attached=hasImage && locks.engine !== 'krea'
+  const reference=locks.mode!=='generate' && locks.identityStillId && locks.imageAId && locks.identityStillId!==locks.imageAId ? 'second (hero) reference' : 'image reference'
+  return [sectionsFrom(value).action,
+    attached && locks.lockFace !== false ? `Preserve the person’s face, glasses, hair, and body identity from the ${reference}.` : '',
+    attached && locks.lockOutfit !== false ? `Preserve the wardrobe from the ${reference}.` : '',
+    attached && locks.lockOutfit === false ? 'Clothing may follow the source and requested scene; do not restore the hero outfit.' : ''
+  ].filter(Boolean).join('\n\n')
 }
 export function restoreInput(item) {
   const saved = structuredClone(item.studio2 || item)
@@ -28,12 +36,13 @@ export function validateRequest(raw, xaigen = false) {
   if (Array.isArray(raw.batch) && raw.batch.length>30) fail('A batch supports up to 31 shots including the first prompt.')
   if (Array.isArray(raw.guides) && raw.guides.length>8) fail('Use up to eight guide stills.')
   if (!raw.folderId) fail('Choose a folder.')
-  if ((video || ['edit','compose','iterate'].includes(mode)) && (raw.lockFace !== false || raw.lockOutfit !== false) && !raw.identityStillId) fail('Choose a sticky hero still.')
+  if ((video || ['edit','compose','iterate'].includes(mode)) && (raw.lockFace !== false || raw.lockOutfit !== false) && !raw.identityStillId && (video || !raw.imageAId)) fail('Choose a sticky hero still.')
   if (mode === 'compose' && !raw.imageBId) fail('Compose requires a second still.')
   if (mode === 'refine' && !raw.maskId) fail('Refine requires a mask.')
-  if (['edit','compose','refine'].includes(mode) && !raw.imageAId) fail('Choose the source still.')
-  if (mode === 'extend' && !raw.startClipId) fail('Choose the video to extend.')
-  const promptSections = sectionsFrom(raw.promptSections || {action:raw.prompt || ''})
+  if (['edit','compose','refine'].includes(mode) && !raw.imageAId && !(mode==='edit' && engine==='flux' && raw.identityStillId)) fail('Choose the source still.')
+  if (mode === 'extend' && !raw.startClipId) fail('Select the clip in the family strip or Previous clip.')
+  const allSections = sectionsFrom(raw.promptSections || {action:raw.prompt || ''})
+  const promptSections = video ? allSections : {action:allSections.action}
   if (!promptSections.action.trim()) fail('Enter an action prompt.')
   const number=(key, fallback, min, max) => { const n=Number(raw.settings?.[key] ?? fallback); if (!Number.isFinite(n)||n<min||n>max) fail(`Invalid ${key}.`); return n }
   for (const key of ['steps']) if (!Number.isInteger(number(key,key==='steps'?24:1,0,2147483647))) fail(`Invalid ${key}.`)
@@ -51,7 +60,7 @@ export function validateRequest(raw, xaigen = false) {
     promptSections, compiledPrompt:compilePrompt(promptSections,video,raw),
     settings: { aspect:['auto','16:9','9:16','1:1'].includes(raw.settings?.aspect)?raw.settings.aspect:'auto', fps:number('fps',24,1,60), width:number('width',1344,256,4096),height:number('height',768,256,4096),steps:number('steps',video?8:24,1,100),cfg:number('cfg',video?1.5:4,0,20),seedMode:raw.settings?.seedMode==='locked'?'locked':'random',seed:shotSeeds[0],duration:number('duration',5,1,15),turbo:raw.settings?.turbo !== false, loraStack:Array.isArray(raw.settings?.loraStack)?raw.settings.loraStack:[] },
     guides: (Array.isArray(raw.guides)?raw.guides:[]).slice(0,8).map(g=>{ if (!g.stillId || !Number.isInteger(g.frame) || g.frame<1 || g.frame>=frames) fail('Each guide requires a still and positive frame index.'); return {stillId:String(g.stillId),frame:g.frame} }),
-    batch: (Array.isArray(raw.batch)?raw.batch:[]).slice(0,30).map(s=>{const next=sectionsFrom({...promptSections,...(typeof s==='string'?{action:s}:s)});if(!next.action.trim())fail('Every shot needs an action prompt.');return next}) }
+    batch: (Array.isArray(raw.batch)?raw.batch:[]).slice(0,30).map(s=>{const next=sectionsFrom({...promptSections,...(typeof s==='string'?{action:s}:s)});if(!next.action.trim())fail('Every shot needs an action prompt.');return video?next:{action:next.action}}) }
 }
 export function scopedFile(file, prefix) {
   const clean = s => String(s || '').replaceAll('\\','/')
diff --git a/shared/studio2/graphs.mjs b/shared/studio2/graphs.mjs
index 4a5e17a..6c8f32d 100644
--- a/shared/studio2/graphs.mjs
+++ b/shared/studio2/graphs.mjs
@@ -34,6 +34,16 @@ export function attachHeroReference(g, image) {
   }
   return g
 }
+export function sampleKleinSource(graph, denoise = .65) {
+  if (!graph['11'] || !graph['19'] || !graph['17']) throw new Error('Klein source sampling nodes are missing.')
+  graph['19'].inputs.latent_image=['11',0]
+  delete graph['14']
+  const steps=graph['17'].inputs.steps
+  graph['17'].inputs.steps=Math.ceil(steps/denoise)
+  graph.source_sigmas={class_type:'SplitSigmasDenoise',inputs:{sigmas:['17',0],denoise}}
+  graph['19'].inputs.sigmas=['source_sigmas',1]
+  return graph
+}
 /** The shared image builders have their own defaults and megapixel scalers.
  * Studio 2 must send its already resolved canvas size through those nodes. */
 export function applyResolvedImageSize(graph, {width,height}) {
diff --git a/shared/studio2/input.mjs b/shared/studio2/input.mjs
index c95b951..122c1c3 100644
--- a/shared/studio2/input.mjs
+++ b/shared/studio2/input.mjs
@@ -7,21 +7,28 @@ export function resolveSize(aspect, source, fallback = {width:1344,height:768})
   const factor = Math.min(1344/Math.max(source.width,source.height),768/Math.min(source.width,source.height))
   return {width:Math.max(256,Math.round(source.width*factor/32)*32),height:Math.max(256,Math.round(source.height*factor/32)*32)}
 }
-export function bindClip(form, clip) {
+export function bindClip(form, clip, xaigen = false) {
   if (!clip?.id) return form
-  const saved = clip.studio2 || {}, settings = saved.settings || clip.settings || clip
+  const saved = clip.studio2 || clip.job || {}, settings = saved.settings || clip.settings || clip
   const next = {...form, startClipId:clip.id, settings:{...form.settings}}
   next.settings.aspect=['auto','16:9','9:16','1:1'].includes(settings.aspect)?settings.aspect:'auto'
   for (const key of ['duration','aspect','width','height','fps','loraStack','turbo','steps','cfg']) {
     if (settings[key] != null && key!=='aspect') next.settings[key] = structuredClone(settings[key])
   }
-  next.engine = saved.engine || clip.engine || 'minimax'
+  next.engine = xaigen && (saved.engine || clip.engine)==='ltx' ? 'ltx' : 'minimax'
   if (saved.identityStillId) next.identityStillId = saved.identityStillId
   // Missing legacy metadata must not reset a user's current choice.
   for (const key of ['lockFace','lockOutfit']) if (typeof saved[key] === 'boolean') next[key] = saved[key]
   return next
 }
 export function restoreAsset(form, item) {
+  if (item.kind==='clip' || item.studio2?.kind==='video') {
+    const next=bindClip(form,item,form.engine==='ltx')
+    const seed=item.studio2?.settings?.seed ?? item.settings?.seed ?? item.seed
+    if (seed != null) next.settings.seed=seed
+    next.settings.seedMode='random'
+    return {...next,mode:'extend'}
+  }
   const restored = restoreInput(item)
   const next = {...form,...restored,settings:{...form.settings,...restored.settings}}
   next.settings.seedMode='random'
@@ -31,7 +38,11 @@ export function restoreAsset(form, item) {
   next.settings.aspect=item.studio2?.settings?.aspect || item.settings?.aspect || 'auto'
   return next
 }
-export function extendSelected(form, clip) {
+export function extendSelected(form, clip, xaigen = false) {
   // Switching task on the already bound clip must preserve the user's overrides.
-  return {...(clip?.id && form.startClipId !== clip.id ? bindClip(form,clip) : form),mode:'extend'}
+  const next=clip?.id && form.startClipId !== clip.id ? bindClip(form,clip,xaigen) : form
+  return {...next,mode:'extend',engine:xaigen && next.engine==='ltx'?'ltx':'minimax'}
+}
+export function changeSource(form) {
+  return form.mode==='extend' ? form : {...form,startClipId:'',settings:{...form.settings,aspect:'auto'}}
 }
diff --git a/shared/studio2/upload.mjs b/shared/studio2/upload.mjs
index e395366..485d15d 100644
--- a/shared/studio2/upload.mjs
+++ b/shared/studio2/upload.mjs
@@ -3,6 +3,6 @@ export function applyUploadedStill(form, response, target) {
   if (!still?.id) throw new Error('The upload did not return a saved still. Please try again.')
   const next={...form,settings:{...form.settings}}
   if (target==='hero') next.identityStillId=still.id
-  else {next.imageAId=still.id;next.startClipId='';if(next.mode==='extend')next.mode='video';next.settings.aspect='auto'}
+  else {next.imageAId=still.id;if(next.mode!=='extend'){next.startClipId='';next.settings.aspect='auto'}}
   return next
 }
diff --git a/tests/studio2-controls.test.mjs b/tests/studio2-controls.test.mjs
index 709430a..1a1f2af 100644
--- a/tests/studio2-controls.test.mjs
+++ b/tests/studio2-controls.test.mjs
@@ -37,8 +37,8 @@ test('appearance roles distinguish wardrobe lock and both-off graph excludes her
  assert.equal(graph.hero,undefined);assert.equal(graph['155'].inputs['ref_images.ref_image_0'],undefined);assert.equal(graph.load_start.inputs.image,'start.png');assert.equal(r.identityStillId,'hero')
 })
 test('sound/dialogue/music compile without empty invented audio',()=>{
- const empty=compilePrompt({});assert.doesNotMatch(empty,/Soundscape:|<d>/);assert.match(empty,/non_diegetic_music: N\/A/)
- const text=compilePrompt({soundscapeChips:['footsteps'],dialogue:{mode:'spoken off-screen',line:'Hello'}});assert.match(text,/lips remain closed/);assert.match(text,/<Subject 1> \(S1\) says <d>\[English\]Hello<\/d>/)
+ const empty=compilePrompt({},true);assert.doesNotMatch(empty,/Soundscape:|<d>/);assert.match(empty,/non_diegetic_music: N\/A/)
+ const text=compilePrompt({soundscapeChips:['footsteps'],dialogue:{mode:'spoken off-screen',line:'Hello'}},true);assert.match(text,/lips remain closed/);assert.match(text,/<Subject 1> \(S1\) says <d>\[English\]Hello<\/d>/)
 })
 test('aspect sizing follows source and named presets',()=>{
  assert.deepEqual(resolveSize('9:16',{width:1344,height:768}),{width:768,height:1344});const size=resolveSize('auto',{width:768,height:1024});assert.equal(size.width/size.height,.75)
diff --git a/tests/studio2-identity-edit.test.mjs b/tests/studio2-identity-edit.test.mjs
new file mode 100644
index 0000000..b5aab45
--- /dev/null
+++ b/tests/studio2-identity-edit.test.mjs
@@ -0,0 +1,45 @@
+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 {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={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(){}}
+ 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,.65)
+ assert.equal(record.sourceStillId,sourceId||'hero');assert.equal(record.sampleLatent,'source image latent')
+ 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/)
+})
diff --git a/tests/studio2-upload.test.mjs b/tests/studio2-upload.test.mjs
index 5e99831..658bbd6 100644
--- a/tests/studio2-upload.test.mjs
+++ b/tests/studio2-upload.test.mjs
@@ -4,7 +4,7 @@ import {applyUploadedStill} from '../shared/studio2/upload.mjs'
 test('uploaded source is selected without changing hero or output',()=>{
  const form={imageAId:'old',identityStillId:'hero',selected:'result',startClipId:'clip',mode:'extend',settings:{aspect:'16:9'}}
  const next=applyUploadedStill(form,{still:{id:'uploaded'}},'source')
- assert.equal(next.imageAId,'uploaded');assert.equal(next.identityStillId,'hero');assert.equal(next.selected,'result');assert.equal(next.startClipId,'');assert.equal(next.settings.aspect,'auto')
+ assert.equal(next.imageAId,'uploaded');assert.equal(next.identityStillId,'hero');assert.equal(next.selected,'result');assert.equal(next.startClipId,'clip');assert.equal(next.mode,'extend');assert.equal(next.settings.aspect,'16:9')
 })
 test('uploaded hero changes identity only and missing saves are reported',()=>{
  const form={identityStillId:'old',imageAId:'source',settings:{aspect:'1:1'}}
diff --git a/tests/studio2.test.mjs b/tests/studio2.test.mjs
index a52ce47..9c28c4f 100644
--- a/tests/studio2.test.mjs
+++ b/tests/studio2.test.mjs
@@ -12,7 +12,7 @@ test('compose without B and refine without mask are 400',()=>{
  for(const mode of ['compose','refine']) assert.throws(()=>validateRequest({...base,mode,engine:'flux',imageAId:'a'}),e=>e.statusCode===400)
 })
 test('Use as input round-trips camera chip and never invents one for flat text',()=>{
- const r=validateRequest({...base,promptSections:{action:'walk',cameraChip:{family:'dolly in',amplitude:'subtle',speed:'slow'}}}); assert.deepEqual(restoreInput({studio2:r}).promptSections,r.promptSections);assert.equal(restoreInput({prompt:'old text'}).promptSections.cameraChip,undefined);assert.match(compilePrompt(r.promptSections),/dolly in/)
+ const r=validateRequest({...base,promptSections:{action:'walk',cameraChip:{family:'dolly in',amplitude:'subtle',speed:'slow'}}}); assert.deepEqual(restoreInput({studio2:r}).promptSections,r.promptSections);assert.equal(restoreInput({prompt:'old text'}).promptSections.cameraChip,undefined);assert.match(compilePrompt(r.promptSections,true),/dolly in/)
 })
 test('purge accepts only exact preview path boundary',()=>{
  assert.equal(scopedFile({subfolder:'video/MiniMax_H3_preview/job',filename:'a.png'},'video/MiniMax_H3_preview'),true)