towsty/aigen
1
0
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases Wiki Activity

Add Image v2 Refine for on-device regional edits.

Browse Source 
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>
This commit is contained in:
Towsty
2026-08-28 22:03:20 -05:00
co-authored by Cursor
parent 21fe2412f9
commit a2895b8054
10 changed files with 825 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files
server/utils/generationPresets.ts
+5 -2
View File
@@ -21,6 +21,8 @@ import {
IMAGE_V2_SNOFS_CLIP,
IMAGE_V2_SNOFS_MODEL,
IMAGE_V2_STEPS_DEFAULT,
IMAGE_V2_DENOISE_DEFAULT,
clampImageV2Denoise,
clampImageV2Strength,
parseImageV2Mode,
parseImageV2Task,
@@ -114,7 +116,7 @@ function sanitizeImageV2Settings(raw: unknown): ImageV2PresetSettings {
const turbo = rec.turbo === true
return {
mode,
task: mode === 'compose' ? parseImageV2Task(rec.task, 'scene') : 'scene',
task: mode === 'compose' ? parseImageV2Task(rec.task, 'scene') : mode === 'refine' ? 'refine' : 'scene',
negative: String(rec.negative || '').slice(0, 2000),
snofsModel: clampImageV2Strength(rec.snofsModel ?? rec.snofs_model, IMAGE_V2_SNOFS_MODEL),
snofsClip: clampImageV2Strength(rec.snofsClip ?? rec.snofs_clip, IMAGE_V2_SNOFS_CLIP),
@@ -123,7 +125,8 @@ function sanitizeImageV2Settings(raw: unknown): ImageV2PresetSettings {
steps: turbo ? 8 : clampImageSteps(rec.steps, IMAGE_V2_STEPS_DEFAULT),
cfg: turbo ? 1 : clampImageCfg(rec.cfg, IMAGE_V2_CFG_DEFAULT),
megapixels: clampImageScaleMegapixels(rec.megapixels ?? rec.scaleMegapixels, 1),
turbo
turbo,
strength: mode === 'refine' ? clampImageV2Denoise(rec.strength, IMAGE_V2_DENOISE_DEFAULT) : undefined
}
}
server/utils/imageChainV2.ts
+27 -8
View File
@@ -15,6 +15,7 @@ export type EditV2RunParams = {
task: ImageV2Task
image: EditImageFile
reference: EditImageFile | null
mask?: EditImageFile | null
prompt: string
negative: string
steps: number
@@ -25,11 +26,15 @@ export type EditV2RunParams = {
consistencyModel: number
consistencyClip: number
megapixels: number
strength?: number
}
export async function runEditV2(job: Job, params: EditV2RunParams) {
const library = job.library
if (!library) throw new Error('Edit job is missing library metadata')
if (params.mode === 'refine' && !params.mask) {
throw new Error('Refine requires a mask. Refusing to fall back to Edit.')
}
if (params.mode === 'compose' && !params.reference) {
throw new Error('Compose requires image B. Refusing to fall back to one-image generation.')
}
@@ -52,23 +57,35 @@ export async function runEditV2(job: Job, params: EditV2RunParams) {
await withImageComfyHost(job.imageComfyHost, async () => {
emitChainJob(job, {
type: 'status',
message: params.mode === 'compose' ? 'Uploading stills A and B to Beast...' : 'Uploading still A to Beast...',
message: params.mode === 'refine'
? 'Uploading canvas and mask to Beast...'
: params.mode === 'compose' ? 'Uploading stills A and B to Beast...' : 'Uploading still A to Beast...',
progress: 8
})
const uploaded = await uploadImage(params.image, job.id)
const uploadedRef = params.reference
const uploadedRef = params.mode === 'refine'
? null
: params.reference
? await uploadImage({
...params.reference,
filename: `ref_${params.reference.filename || 'image_b.png'}`
}, job.id)
: null
const uploadedMask = params.mode === 'refine' && params.mask
? await uploadImage({
...params.reference,
filename: `ref_${params.reference.filename || 'image_b.png'}`
...params.mask,
filename: `mask_${params.mask.filename || 'refine-mask.png'}`
}, job.id)
: null
if (job.status === 'cancelled') throw new Error('Job interrupted.')
emitChainJob(job, {
type: 'status',
message: params.mode === 'compose'
? `Queueing Klein v2 compose (${params.task}) on Beast...`
: 'Queueing Klein v2 edit on Beast...',
message: params.mode === 'refine'
? 'Queueing Klein v2 refine on Beast...'
: params.mode === 'compose'
? `Queueing Klein v2 compose (${params.task}) on Beast...`
: 'Queueing Klein v2 edit on Beast...',
progress: 12
})
await ensureComfyLoraNames('image')
@@ -80,6 +97,8 @@ export async function runEditV2(job: Job, params: EditV2RunParams) {
negative: params.negative,
imageAName: uploaded.name,
imageBName: uploadedRef?.name,
maskName: uploadedMask?.name,
strength: params.strength,
snofsModel: params.snofsModel,
snofsClip: params.snofsClip,
consistencyModel: params.consistencyModel,
@@ -137,7 +156,7 @@ export async function runEditV2(job: Job, params: EditV2RunParams) {
video: { filename: output.filename, subfolder: output.subfolder, type: output.type },
imageName: uploaded.name,
imageSubfolder: uploaded.subfolder,
extraImageNames: uploadedRef?.name ? [uploadedRef.name] : [],
extraImageNames: [uploadedRef?.name, uploadedMask?.name].filter((name): name is string => Boolean(name)),
promptId: job.promptId
})
server/utils/imageWorkflowV2.ts
+56 -5
View File
@@ -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'
server/utils/studioQueue.ts
+21 -5
View File
@@ -52,12 +52,15 @@ export interface StudioJobPayload {
scaleToTotalPixels?: boolean
scaleMegapixels?: number
imagePipeline?: 'v1' | 'v2'
v2Mode?: 'edit' | 'compose'
v2Task?: 'scene' | 'identity' | 'outfit' | 'face_lock'
v2Mode?: 'edit' | 'compose' | 'refine'
v2Task?: 'scene' | 'identity' | 'outfit' | 'face_lock' | 'refine'
snofsModel?: number
snofsClip?: number
consistencyModel?: number
consistencyClip?: number
maskStillId?: string
maskStillFilename?: string
refineStrength?: number
}
export interface StudioJob {
@@ -694,7 +697,18 @@ async function startStudioEditJob(item: StudioJob) {
: null
if (payload.imagePipeline === 'v2') {
const mode = payload.v2Mode === 'compose' ? 'compose' : 'edit'
const mode = payload.v2Mode === 'compose' ? 'compose' : payload.v2Mode === 'refine' ? 'refine' : 'edit'
const maskId = payload.maskStillId
const mask = mode === 'refine' && maskId && existsSync(stillPath(item.ownerKey, maskId))
? {
filename: payload.maskStillFilename || 'refine-mask.png',
data: readFileSync(stillPath(item.ownerKey, maskId)),
type: 'image/png'
}
: null
if (mode === 'refine' && !mask) {
throw new Error('Refine requires a mask. Refusing to fall back to Edit.')
}
if (mode === 'compose' && !reference) {
throw new Error('Compose requires image B. Refusing to fall back to one-image generation.')
}
@@ -732,7 +746,8 @@ async function startStudioEditJob(item: StudioJob) {
mode,
task: payload.v2Task || 'scene',
image,
reference,
reference: mode === 'refine' ? null : reference,
mask,
prompt: payload.prompt,
negative: payload.negative || '',
steps: payload.steps,
@@ -742,7 +757,8 @@ async function startStudioEditJob(item: StudioJob) {
snofsClip: payload.snofsClip ?? 0.35,
consistencyModel: payload.consistencyModel ?? 0.7,
consistencyClip: payload.consistencyClip ?? 0.7,
megapixels: payload.scaleMegapixels ?? 1
megapixels: payload.scaleMegapixels ?? 1,
strength: payload.refineStrength
}).catch((error) => {
const message = error instanceof Error ? error.message : String(error)
if (live && live.status !== 'error' && live.status !== 'cancelled' && live.status !== 'deferred') {