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

Add Image v2 Generate for Klein text-to-image.

Browse Source 
New sibling graph and mode so T2I does not borrow Edit, Compose, or Refine. Stills are ignored; size is width x height only.

Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Towsty
2026-08-28 22:56:33 -05:00
co-authored by Cursor
parent a2895b8054
commit b35f9af10d
9 changed files with 511 additions and 112 deletions
Download Patch File Download Diff File Expand all files Collapse all files
server/utils/imageWorkflowV2.ts
+73 -16
View File
@@ -1,6 +1,7 @@
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 generateTemplate from '../assets/klein_v2_generate.json'
import { IMAGE_SCALE_TO_TOTAL_PIXELS } from '~/server/utils/comfy'
import { cachedComfyLoraNames } from '~/server/utils/loras'
import { resolveComfyLoraName, loraIdentityKey } from '~/utils/loras'
@@ -13,7 +14,10 @@ import {
IMAGE_V2_SNOFS_LORA,
IMAGE_V2_SNOFS_MODEL,
IMAGE_V2_DENOISE_DEFAULT,
IMAGE_V2_GENERATE_HEIGHT,
IMAGE_V2_GENERATE_WIDTH,
clampImageV2Denoise,
clampImageV2Size,
clampImageV2Strength,
composeImageV2Prompt,
type ImageV2Mode,
@@ -41,16 +45,20 @@ 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 const IMAGE_V2_GENERATE_WORKFLOW = 'klein_v2_generate.json'
export interface ImageV2BuildParams {
mode: ImageV2Mode
task: ImageV2Task
prompt: string
negative?: string
imageAName: string
imageAName?: string
imageBName?: string
maskName?: string
strength?: number
width?: number
height?: number
turbo?: boolean
snofsModel?: number
snofsClip?: number
consistencyModel?: number
@@ -102,8 +110,30 @@ function graphHasMaskInput(graph: WorkflowGraph) {
})
}
function bypassLoraNode(graph: WorkflowGraph, id: string, modelFrom: string, clipFrom: string) {
delete graph[id]
for (const node of Object.values(graph)) {
for (const [key, value] of Object.entries(node.inputs)) {
if (!Array.isArray(value) || value[0] !== id) continue
node.inputs[key] = value[1] === 1 ? [clipFrom, 1] : [modelFrom, 0]
}
}
}
export function assertImageV2Graph(graph: WorkflowGraph, mode: ImageV2Mode, imageBName?: string) {
const loaders = loadImageNames(graph)
if (mode === 'generate') {
if (loaders.length) {
throw createError({
statusCode: 500,
statusMessage: 'Generate graph has a required LoadImage. Refusing to run an edit fallback.'
})
}
const latent = Object.values(graph).find(node => node.class_type === 'EmptyFlux2LatentImage')
if (!latent) {
throw createError({ statusCode: 500, statusMessage: 'Generate graph is missing EmptyFlux2LatentImage.' })
}
}
if (mode === 'refine') {
const mask = graph[LOAD_MASK]
if (!mask || mask.class_type !== 'LoadImage' || !String(mask.inputs.image || '').trim()) {
@@ -140,7 +170,7 @@ export function assertImageV2Graph(graph: WorkflowGraph, mode: ImageV2Mode, imag
})
}
}
if (mode !== 'refine' && imageBName && loaders.length < 2) {
if (mode !== 'refine' && mode !== 'generate' && imageBName && loaders.length < 2) {
throw createError({
statusCode: 500,
statusMessage: 'image_b was sent but the executed graph has no second image input.'
@@ -161,43 +191,65 @@ export function assertImageV2Graph(graph: WorkflowGraph, mode: ImageV2Mode, imag
export function buildImageV2Workflow(params: ImageV2BuildParams) {
const compose = params.mode === 'compose'
const refine = params.mode === 'refine'
const generate = params.mode === 'generate'
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 graph = structuredClone(
generate ? generateTemplate : 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)
const snofsClip = clampImageV2Strength(params.snofsClip, IMAGE_V2_SNOFS_CLIP)
const consistencyModel = clampImageV2Strength(params.consistencyModel, IMAGE_V2_CONSISTENCY_MODEL)
const consistencyClip = clampImageV2Strength(params.consistencyClip, IMAGE_V2_CONSISTENCY_CLIP)
const consistencyModel = clampImageV2Strength(
params.consistencyModel,
generate ? 0 : IMAGE_V2_CONSISTENCY_MODEL
)
const consistencyClip = clampImageV2Strength(
params.consistencyClip,
generate ? 0 : IMAGE_V2_CONSISTENCY_CLIP
)
const steps = clampImageSteps(params.steps, 24)
const cfg = clampImageCfg(params.cfg, 4)
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
const size = generate ? clampImageV2Size(params.width, params.height) : null
const workflowFile = generate
? IMAGE_V2_GENERATE_WORKFLOW
: refine ? IMAGE_V2_REFINE_WORKFLOW : compose ? IMAGE_V2_COMPOSE_WORKFLOW : IMAGE_V2_EDIT_WORKFLOW
setInput(graph, LOAD_A, 'image', params.imageAName)
if (!generate) 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)
}
if (generate && size) {
setInput(graph, '14', 'width', size.width)
setInput(graph, '14', 'height', size.height)
setInput(graph, SCHEDULER, 'width', size.width)
setInput(graph, SCHEDULER, 'height', size.height)
}
setInput(graph, PROMPT, 'text', prompt)
setInput(graph, NEGATIVE, 'text', negative)
setInput(graph, NOISE, 'noise_seed', params.seed)
setInput(graph, SCHEDULER, 'steps', steps)
setInput(graph, CFG, 'cfg', cfg)
setInput(graph, SAVE, 'filename_prefix', params.filenamePrefix || 'aigen-v2')
patchScaleMegapixels(graph, params.megapixels ?? 1)
setInput(graph, SAVE, 'filename_prefix', params.filenamePrefix || (generate ? 'aigen-v2-generate' : 'aigen-v2'))
if (!generate) patchScaleMegapixels(graph, params.megapixels ?? 1)
setInput(graph, SNOFS, 'lora_name', resolveRequiredLora(IMAGE_V2_SNOFS_LORA, 'SNOFS'))
setInput(graph, SNOFS, 'strength_model', snofsModel)
setInput(graph, SNOFS, 'strength_clip', snofsClip)
setInput(graph, CONSISTENCY, 'lora_name', resolveRequiredLora(IMAGE_V2_CONSISTENCY_LORA, 'Consistency'))
setInput(graph, CONSISTENCY, 'strength_model', consistencyModel)
setInput(graph, CONSISTENCY, 'strength_clip', consistencyClip)
if (generate && consistencyModel <= 0 && consistencyClip <= 0) {
bypassLoraNode(graph, CONSISTENCY, SNOFS, SNOFS)
} else {
setInput(graph, CONSISTENCY, 'lora_name', resolveRequiredLora(IMAGE_V2_CONSISTENCY_LORA, 'Consistency'))
setInput(graph, CONSISTENCY, 'strength_model', consistencyModel)
setInput(graph, CONSISTENCY, 'strength_clip', consistencyClip)
}
assertImageV2Graph(graph, params.mode, refine ? undefined : params.imageBName)
assertImageV2Graph(graph, params.mode, generate || refine ? undefined : params.imageBName)
const loaders = loadImageNames(graph)
console.log(JSON.stringify({
@@ -205,18 +257,22 @@ export function buildImageV2Workflow(params: ImageV2BuildParams) {
workflow: workflowFile,
mode: params.mode,
task: params.task,
canvas: { id: LOAD_A, file: graph[LOAD_A]?.inputs.image },
canvas: generate ? undefined : { id: LOAD_A, file: graph[LOAD_A]?.inputs.image },
mask: refine ? { id: LOAD_MASK, file: graph[LOAD_MASK]?.inputs.image } : undefined,
size: generate ? { width: size?.width ?? IMAGE_V2_GENERATE_WIDTH, height: size?.height ?? IMAGE_V2_GENERATE_HEIGHT } : undefined,
strength,
turbo: params.turbo === true,
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 },
consistency: { name: graph[CONSISTENCY]?.inputs.lora_name, model: consistencyModel, clip: consistencyClip }
consistency: graph[CONSISTENCY]
? { name: graph[CONSISTENCY]?.inputs.lora_name, model: consistencyModel, clip: consistencyClip }
: { loaded: false, model: 0, clip: 0 }
},
steps,
cfg,
seed: params.seed,
megapixels: clampImageScaleMegapixels(params.megapixels ?? 1)
megapixels: generate ? undefined : clampImageScaleMegapixels(params.megapixels ?? 1)
}))
return { graph, workflowFile, loaders, prompt, strength }
@@ -238,6 +294,7 @@ export const IMAGE_V2_NODE_LABELS: Record<string, string> = {
'11': 'Encoding image A',
'24': 'Encoding image B',
'33': 'Applying mask',
'14': 'Building empty Klein latent',
'19': 'Sampling Klein v2',
'20': 'Decoding still',
'21': 'Saving still'