From 30be4bd5b0ef9101d57e5715844451ff8692fc27 Mon Sep 17 00:00:00 2001
From: Towsty <ianjmcarr@gmail.com>
Date: Sat, 26 Sep 2026 14:27:36 -0500
Subject: [PATCH] Wire Qwen Edit Start still to image_1 and PE-I2I identity
 keep.

Edit no longer invents strangers when the prompt omits <image1> or PE drops image tokens; Hero stays Klein-only.

Co-authored-by: Cursor <cursoragent@cursor.com>
---
 components/studio2/Bench.vue   |  6 +++---
 docs/qwen21.md                 |  4 +++-
 server/utils/studio2/runner.ts | 28 +++++++++++++++++++---------
 shared/studio2/contracts.mjs   |  5 ++++-
 shared/studio2/qwen21-edit.mjs | 26 ++++++++++++++++++++++++++
 tests/studio2-bench.test.mjs   | 20 +++++++++++++++++++-
 6 files changed, 74 insertions(+), 15 deletions(-)
 create mode 100644 shared/studio2/qwen21-edit.mjs

diff --git a/components/studio2/Bench.vue b/components/studio2/Bench.vue
index 9588e55..58ad219 100644
--- a/components/studio2/Bench.vue
+++ b/components/studio2/Bench.vue
@@ -109,7 +109,7 @@
        <label class="inline-flex items-center gap-2"><input v-model="form.lockOutfit" type="checkbox" class="accent-[var(--amber)]">Lock outfit</label>
        <p class="text-[11px]">hero encode + lock text — not a MiniMax face/clothes split on stills</p>
       </div>
-      <p v-else-if="form.mode==='edit'" class="text-[var(--type-sm)] text-[var(--muted)]">Qwen edit uses source as &lt;image1&gt;, not Klein hero-ref. No locks or LoRA stack.</p>
+      <p v-else-if="form.mode==='edit'" class="text-[var(--type-sm)] text-[var(--muted)]">Qwen Edit uses Start still as &lt;image1&gt;. Hero is Klein-only and is ignored.</p>
       <p v-else class="text-[var(--type-sm)] text-[var(--muted)]">Qwen 2.1 — no hero, locks, or LoRA stack.</p>
 
       <label>
@@ -521,7 +521,7 @@ const editSourceRequired=computed(()=>['edit','compose','refine'].includes(form.
 const heroSlotLabel=computed(()=>editSourceRequired.value?'Identity / reference':'Hero still')
 const sourceSlotLabel=computed(()=>{
  if(form.mode==='extend')return 'Previous clip'
- if(editSourceRequired.value)return 'Edit source'
+ if(editSourceRequired.value)return form.engine==='qwen21'?'Start still':'Edit source'
  return 'Start still'
 })
 const slotPairHint=computed(()=>{
@@ -529,7 +529,7 @@ const slotPairHint=computed(()=>{
   ? 'Krea Refine: Edit source is the picture; Mask is the refine region.'
   : 'Refine: Edit source is the picture; Mask marks the region to change.'
  if(!editSourceRequired.value)return ''
- if(form.engine==='qwen21')return 'Qwen: Edit source is <image1>. Hero is unused unless the graph maps it to <image2>.'
+ if(form.engine==='qwen21')return 'Qwen Edit uses Start still as <image1>. Hero is Klein-only and is ignored.'
  return 'Klein: Hero = identity encode. Edit source = the picture you are changing.'
 })
 const queueCards=computed(()=>{
diff --git a/docs/qwen21.md b/docs/qwen21.md
index c8a1383..98d33cc 100644
--- a/docs/qwen21.md
+++ b/docs/qwen21.md
@@ -35,12 +35,14 @@ Restart Comfy **only when idle** (`COMFY_CONTROL_URL/status` → `gpu.busy=false
 
 - Engine key: `qwen21` · UI label: **Qwen 2.1**
 - **Generate** (T2I) and **Edit** (same checkpoint, second graph). Compose / Iterate / Video / Extend / Music stay disabled unless a graph exists.
-- Edit requires Start still → `images.image_1`; prompt slots use `<image1>` / `<image2>`. If the user omits `<image1>`, the runner prepends `Keep the subject in <image1>.`
+- Edit requires Start still → `images.image_1`; prompt slots use `<image1>` / `<image2>`. If the user omits `<image1>`, the runner prepends `Keep the identity, face, body, and pose of the person in <image1>. Apply only this change:` then their text. Hero is Klein-only and is ignored (never image_1).
 - Face / outfit locks stay Klein semantics — they do not drive Qwen slots.
 - Sampler defaults: euler / simple / cfg **1** / steps **25**. Edit uses `QwenImage21Cache` (device auto, dtype int8). T2I uses `ModelSamplingAuraFlow` shift **3.1**.
 - Default Generate canvas follows the bench Aspect control on the Qwen-safe ~1 MP grid (`EmptyLatentImage`): 1:1 → 1024×1024, 16:9 → 1536×864, 9:16 → 864×1536 (and the other table rows). PE `wh_ratio` is advisory only and never sizes the canvas. Do not use native 2K bins on 16 GB.
 - Edit follows `image_1` via the encode node’s latent (resolution long-edge ~1024). Do not inject a picker EmptyLatentImage onto the edit sampler.
 - Optional **Enhance prompt** (off by default): runs a separate PE-only Comfy graph (`CLIPLoader` + rewrite node), then frees VRAM and queues the existing T2I/Edit graph with the rewritten prompt. Never loads PE CLIP + DiT together on 16 GB. Fail closed if `parse_ok` is false or the rewrite is empty.
+  - Generate + Enhance → PE-T2I only (`pe_t2i`). No images.
+  - Edit + Enhance → PE-I2I only (`pe_i2i`) with the same Start still on `image_1`. If the rewrite drops every `<imageN>` token, stitch `Keep the identity…` onto the rewrite before Edit. Never fall back to PE-T2I on Edit.
 - PE weights (int8 only): `text_encoders\qwen3.5_9b_qwen_image_2.1_pe_{t2i,i2i}.int8_convrot.safetensors` — do not replace the image TE `qwen3vl_8b_int8_convrot.safetensors`.
 - Optional **Turbo** (off by default): Viggle DMD LoRA `Qwen-Image-2.1-viggle-turbo-v0.2.1-6step-lora-r128.safetensors` on the same Q8 GGUF via `ViggleTurboLora` + `ViggleTurboSigmas` (`1.0, 0.9375, 0.875, 0.75, 0.5, 0.25`), 6 steps, CFG 1, empty negative. Not a new engine. Enhance prompt stays compatible and is recommended with Turbo.
 - Custom node: `ComfyUI-Viggle-Turbo` (`viggle_turbo.py`). Do not merge the LoRA with stock `LoraLoaderModelOnly` (lossy on int8/bf16).
diff --git a/server/utils/studio2/runner.ts b/server/utils/studio2/runner.ts
index dba2c51..feedaae 100644
--- a/server/utils/studio2/runner.ts
+++ b/server/utils/studio2/runner.ts
@@ -17,6 +17,7 @@ import qwen21PeEditTemplate from '../../assets/studio2_qwen21_pe_edit.json';
 import { nativeVideoGraph, attachHeroReference, applyResolvedImageSize } from '~/shared/studio2/graphs.mjs';
 import { compilePrompt, scopedFile } from '~/shared/studio2/contracts.mjs';
 import { resolveQwen21Size } from '~/shared/studio2/qwen21-size.mjs';
+import { ensureQwen21EditPrompt, stitchQwen21PeEditPrompt } from '~/shared/studio2/qwen21-edit.mjs';
 import { QWEN21_TURBO_LORA, QWEN21_TURBO_SIGMAS, QWEN21_TURBO_STEPS } from '~/shared/studio2/qwen21-turbo.mjs';
 import { createJob, restoreJob, getJob, emitJob, type Job } from '../jobs';
 import { markStudioLive, onLiveVideoSettled, type StudioJob } from '../studioQueue';
@@ -98,16 +99,18 @@ async function runQwen21PromptEnhance(r: any, job: Job) {
     const q = r.request, s = q.settings
     if (!q.enhancePrompt || q.engine !== 'qwen21' || !['generate', 'edit'].includes(q.mode)) return
     const edit = q.mode === 'edit'
-    if (edit && !q.imageAId) throw new Error('Choose the source still.')
+    // Edit → PE-I2I only (pe_i2i). Generate → PE-T2I only. Never cross-feed.
+    if (edit && !q.imageAId)
+        throw new Error('Qwen Edit requires Start still (image_1). Hero is ignored.')
     update(r, 'enhancing')
     emitJob(job, { type: 'status', message: 'Enhancing prompt…' })
     let prompt = String(q.compiledPrompt || '')
-    if (edit && !/<image1>/i.test(prompt))
-        prompt = `Keep the subject in <image1>. ${prompt}`.trim()
+    if (edit) prompt = ensureQwen21EditPrompt(prompt)
     const graph = structuredClone(edit ? qwen21PeEditTemplate : qwen21PeT2iTemplate)
     graph['9'].inputs.prompt = prompt
     graph['9'].inputs.seed = s.seed
     if (edit) {
+        // Same Start still pixels as Edit image_1 — never Hero.
         const a = await upload(r, 'pe-source', readFileSync(stillPath(r.owner, getStill(r.owner, q.imageAId).id)))
         graph['10'].inputs.image = a
         graph['9'].inputs.image_1 = ['10', 0]
@@ -131,7 +134,10 @@ async function runQwen21PromptEnhance(r: any, job: Job) {
     if (!result.parse_ok || !result.positive_prompt)
         throw new Error('Prompt enhance failed to parse a rewrite. Try again or turn Enhance prompt off.')
     r.request.promptRaw = q.compiledPrompt
-    r.request.compiledPrompt = result.positive_prompt
+    // PE-I2I: if rewrite dropped every <imageN>, stitch keep-identity onto the rewrite.
+    r.request.compiledPrompt = edit
+        ? stitchQwen21PeEditPrompt(result.positive_prompt)
+        : result.positive_prompt
     r.request.enhance = {
         wh_ratio: result.wh_ratio || '',
         ...(edit ? { ratio_follow: result.ratio_follow || '' } : {}),
@@ -172,7 +178,10 @@ 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))) : '';
     const video=['video','extend'].includes(q.mode);
     const identityPlan=kleinIdentityPlan(q);
-    const sourceId=identityPlan ? identityPlan.sourceId : q.imageAId;
+    // Qwen Edit: Start still (imageAId) only — never Hero / identityStillId.
+    const sourceId = q.engine === 'qwen21'
+        ? (q.imageAId || '')
+        : (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');
@@ -247,12 +256,13 @@ async function prepareGraph(r: any) {
         }
         const negative = turbo ? '' : stylePrompt(q.imageStyles, true)
         if (q.mode === 'edit') {
-            if (!a) throw new Error('Choose the source still.')
-            let prompt = String(q.compiledPrompt || '')
-            if (!/<image1>/i.test(prompt))
-                prompt = `Keep the subject in <image1>. ${prompt}`.trim()
+            // Start still → image_1 only. Never Hero. Never T2I / EmptyLatent.
+            if (!q.imageAId || !a)
+                throw new Error('Qwen Edit requires Start still (image_1). Hero is ignored.')
+            const prompt = ensureQwen21EditPrompt(String(q.compiledPrompt || ''))
             graph = structuredClone(turbo ? qwen21EditTurboTemplate : qwen21EditTemplate)
             graph['10'].inputs.image = a
+            graph['9'].inputs['images.image_1'] = ['10', 0]
             if (b) {
                 graph['11'] = {
                     inputs: { image: b },
diff --git a/shared/studio2/contracts.mjs b/shared/studio2/contracts.mjs
index 7709c04..7713f11 100644
--- a/shared/studio2/contracts.mjs
+++ b/shared/studio2/contracts.mjs
@@ -40,7 +40,10 @@ 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 (['edit','compose','refine'].includes(mode) && !raw.imageAId && !(mode==='edit' && engine==='flux' && raw.identityStillId)) fail('Choose the source still.')
+  if (['edit','compose','refine'].includes(mode) && !raw.imageAId && !(mode==='edit' && engine==='flux' && raw.identityStillId)) {
+    if (engine === 'qwen21') fail('Qwen Edit requires Start still (image_1). Hero is ignored.')
+    fail('Choose the source still.')
+  }
   if ((video || ['edit','compose','iterate'].includes(mode)) && engine !== 'qwen21' && (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.')
diff --git a/shared/studio2/qwen21-edit.mjs b/shared/studio2/qwen21-edit.mjs
new file mode 100644
index 0000000..6bf5222
--- /dev/null
+++ b/shared/studio2/qwen21-edit.mjs
@@ -0,0 +1,26 @@
+/** Qwen Edit / PE-I2I prompt helpers. Start still = <image1>; Hero is never a slot. */
+
+export const QWEN21_EDIT_KEEP =
+  'Keep the identity, face, body, and pose of the person in <image1>.'
+
+export const QWEN21_EDIT_KEEP_CHANGE =
+  `${QWEN21_EDIT_KEEP} Apply only this change:`
+
+/** User omitted <image1> — prepend the keep-identity instruction once. */
+export function ensureQwen21EditPrompt(prompt) {
+  const text = String(prompt || '').trim()
+  if (!text) return text
+  if (/<image1>/i.test(text)) return text
+  return `${QWEN21_EDIT_KEEP_CHANGE} ${text}`.trim()
+}
+
+/**
+ * After PE-I2I: if the rewrite dropped every <imageN> token, stitch the
+ * keep-identity line onto the rewrite (do not discard the rewrite).
+ */
+export function stitchQwen21PeEditPrompt(rewritten) {
+  const text = String(rewritten || '').trim()
+  if (!text) return text
+  if (/<image\d+>/i.test(text)) return text
+  return `${QWEN21_EDIT_KEEP} ${text}`.trim()
+}
diff --git a/tests/studio2-bench.test.mjs b/tests/studio2-bench.test.mjs
index 6624c46..a35074c 100644
--- a/tests/studio2-bench.test.mjs
+++ b/tests/studio2-bench.test.mjs
@@ -8,6 +8,7 @@ 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'
 
 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]}]
@@ -98,9 +99,14 @@ test('qwen21 supports Generate and Edit with 25/1 defaults and no locks',()=>{
  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'}}),/source still/)
+ 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')
 })
 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:[]}
@@ -121,12 +127,24 @@ test('custom negative chips survive normalizeStyles and bench exposes Load actio
  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.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')