Add Image v2 Refine for on-device regional edits.

Mask painter and denoise strength sit beside Edit and Compose. The canned hand/chest prompt helper is gone so that text stays yours.

Co-authored-by: Cursor <cursoragent@cursor.com>

server/utils/imageWorkflowV2.ts

@@ -1,5 +1,6 @@
 import editTemplate from '../assets/klein_v2_edit.json'
 import composeTemplate from '../assets/klein_v2_compose.json'
+import refineTemplate from '../assets/klein_v2_refine.json'
 import { IMAGE_SCALE_TO_TOTAL_PIXELS } from '~/server/utils/comfy'
 import { cachedComfyLoraNames } from '~/server/utils/loras'
 import { resolveComfyLoraName, loraIdentityKey } from '~/utils/loras'
@@ -11,6 +12,8 @@ import {
   IMAGE_V2_SNOFS_CLIP,
   IMAGE_V2_SNOFS_LORA,
   IMAGE_V2_SNOFS_MODEL,
+  IMAGE_V2_DENOISE_DEFAULT,
+  clampImageV2Denoise,
   clampImageV2Strength,
   composeImageV2Prompt,
   type ImageV2Mode,
@@ -22,6 +25,8 @@ type WorkflowGraph = Record<string, WorkflowNode>
 
 const LOAD_A = '1'
 const LOAD_B = '22'
+const LOAD_MASK = '30'
+const SCHEDULER_DENOISE = '17'
 const SCALE_A = '2'
 const SCALE_B = '23'
 const PROMPT = '9'
@@ -35,6 +40,7 @@ const CONSISTENCY = '8'
 
 export const IMAGE_V2_EDIT_WORKFLOW = 'klein_v2_edit.json'
 export const IMAGE_V2_COMPOSE_WORKFLOW = 'klein_v2_compose.json'
+export const IMAGE_V2_REFINE_WORKFLOW = 'klein_v2_refine.json'
 
 export interface ImageV2BuildParams {
   mode: ImageV2Mode
@@ -43,6 +49,8 @@ export interface ImageV2BuildParams {
   negative?: string
   imageAName: string
   imageBName?: string
+  maskName?: string
+  strength?: number
   snofsModel?: number
   snofsClip?: number
   consistencyModel?: number
@@ -87,8 +95,36 @@ function patchScaleMegapixels(graph: WorkflowGraph, megapixels: number) {
   }
 }
 
+function graphHasMaskInput(graph: WorkflowGraph) {
+  return Object.values(graph).some((node) => {
+    const mask = node.inputs?.mask
+    return mask !== undefined && mask !== null && mask !== ''
+  })
+}
+
 export function assertImageV2Graph(graph: WorkflowGraph, mode: ImageV2Mode, imageBName?: string) {
   const loaders = loadImageNames(graph)
+  if (mode === 'refine') {
+    const mask = graph[LOAD_MASK]
+    if (!mask || mask.class_type !== 'LoadImage' || !String(mask.inputs.image || '').trim()) {
+      throw createError({
+        statusCode: 500,
+        statusMessage: 'Refine job is missing the mask image. Refusing to run.'
+      })
+    }
+    if (!graphHasMaskInput(graph)) {
+      throw createError({
+        statusCode: 500,
+        statusMessage: 'Refine graph has no mask input. Refusing to run.'
+      })
+    }
+    if (!('denoise' in (graph[SCHEDULER]?.inputs || {}))) {
+      throw createError({
+        statusCode: 500,
+        statusMessage: 'Refine graph has no denoise on the sampler. Refusing to run.'
+      })
+    }
+  }
   if (mode === 'compose') {
     if (loaders.length < 2) {
       throw createError({
@@ -104,7 +140,7 @@ export function assertImageV2Graph(graph: WorkflowGraph, mode: ImageV2Mode, imag
       })
     }
   }
-  if (imageBName && loaders.length < 2) {
+  if (mode !== 'refine' && imageBName && loaders.length < 2) {
     throw createError({
       statusCode: 500,
       statusMessage: 'image_b was sent but the executed graph has no second image input.'
@@ -124,7 +160,11 @@ export function assertImageV2Graph(graph: WorkflowGraph, mode: ImageV2Mode, imag
 
 export function buildImageV2Workflow(params: ImageV2BuildParams) {
   const compose = params.mode === 'compose'
-  const graph = structuredClone(compose ? composeTemplate : editTemplate) as WorkflowGraph
+  const refine = params.mode === 'refine'
+  if (refine && !String(params.maskName || '').trim()) {
+    throw createError({ statusCode: 400, statusMessage: 'Refine requires a mask. Refusing to fall back to Edit.' })
+  }
+  const graph = structuredClone(refine ? refineTemplate : compose ? composeTemplate : editTemplate) as WorkflowGraph
   const prompt = composeImageV2Prompt(params.mode, params.task, params.prompt)
   const negative = String(params.negative || '')
   const snofsModel = clampImageV2Strength(params.snofsModel, IMAGE_V2_SNOFS_MODEL)
@@ -133,10 +173,15 @@ export function buildImageV2Workflow(params: ImageV2BuildParams) {
   const consistencyClip = clampImageV2Strength(params.consistencyClip, IMAGE_V2_CONSISTENCY_CLIP)
   const steps = clampImageSteps(params.steps, 24)
   const cfg = clampImageCfg(params.cfg, 4)
-  const workflowFile = compose ? IMAGE_V2_COMPOSE_WORKFLOW : IMAGE_V2_EDIT_WORKFLOW
+  const strength = refine ? clampImageV2Denoise(params.strength, IMAGE_V2_DENOISE_DEFAULT) : undefined
+  const workflowFile = refine ? IMAGE_V2_REFINE_WORKFLOW : compose ? IMAGE_V2_COMPOSE_WORKFLOW : IMAGE_V2_EDIT_WORKFLOW
 
   setInput(graph, LOAD_A, 'image', params.imageAName)
   if (compose) setInput(graph, LOAD_B, 'image', params.imageBName || '')
+  if (refine) {
+    setInput(graph, LOAD_MASK, 'image', params.maskName || '')
+    setInput(graph, SCHEDULER_DENOISE, 'denoise', strength)
+  }
   setInput(graph, PROMPT, 'text', prompt)
   setInput(graph, NEGATIVE, 'text', negative)
   setInput(graph, NOISE, 'noise_seed', params.seed)
@@ -152,7 +197,7 @@ export function buildImageV2Workflow(params: ImageV2BuildParams) {
   setInput(graph, CONSISTENCY, 'strength_model', consistencyModel)
   setInput(graph, CONSISTENCY, 'strength_clip', consistencyClip)
 
-  assertImageV2Graph(graph, params.mode, params.imageBName)
+  assertImageV2Graph(graph, params.mode, refine ? undefined : params.imageBName)
 
   const loaders = loadImageNames(graph)
   console.log(JSON.stringify({
@@ -160,6 +205,9 @@ export function buildImageV2Workflow(params: ImageV2BuildParams) {
     workflow: workflowFile,
     mode: params.mode,
     task: params.task,
+    canvas: { id: LOAD_A, file: graph[LOAD_A]?.inputs.image },
+    mask: refine ? { id: LOAD_MASK, file: graph[LOAD_MASK]?.inputs.image } : undefined,
+    strength,
     loadImage: Object.fromEntries(loaders.map(item => [item.id, { title: item.title, file: item.image }])),
     loras: {
       snofs: { name: graph[SNOFS]?.inputs.lora_name, model: snofsModel, clip: snofsClip },
@@ -171,14 +219,16 @@ export function buildImageV2Workflow(params: ImageV2BuildParams) {
     megapixels: clampImageScaleMegapixels(params.megapixels ?? 1)
   }))
 
-  return { graph, workflowFile, loaders, prompt }
+  return { graph, workflowFile, loaders, prompt, strength }
 }
 
 export const IMAGE_V2_NODE_LABELS: Record<string, string> = {
   '1': 'Loading image A',
   '22': 'Loading image B',
+  '30': 'Loading mask',
   '2': 'Scaling image A',
   '23': 'Scaling image B',
+  '31': 'Resizing mask',
   '4': 'Loading Flux.2 Klein 9B',
   '5': 'Loading CLIP',
   '6': 'Loading VAE',
@@ -187,6 +237,7 @@ export const IMAGE_V2_NODE_LABELS: Record<string, string> = {
   '9': 'Encoding prompt',
   '11': 'Encoding image A',
   '24': 'Encoding image B',
+  '33': 'Applying mask',
   '19': 'Sampling Klein v2',
   '20': 'Decoding still',
   '21': 'Saving still'