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

Show Flux and Krea as Image v2 engine cards, and run Krea on edit, compose, and refine.

Browse Source 
Krea image-to-image uses encode plus denoise graphs, not Klein reference latents, and still fails closed if the Krea models are missing.

Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Towsty
2026-08-29 07:29:31 -05:00
co-authored by Cursor
parent ab5477c999
commit 5fe91e6bae
12 changed files with 608 additions and 119 deletions
Download Patch File Download Diff File Expand all files Collapse all files
server/utils/generationPresets.ts
+9 -2
View File
@@ -22,6 +22,8 @@ import {
IMAGE_V2_SNOFS_MODEL,
IMAGE_V2_STEPS_DEFAULT,
IMAGE_V2_KREA_CFG,
IMAGE_V2_KREA_COMPOSE_DENOISE,
IMAGE_V2_KREA_EDIT_DENOISE,
IMAGE_V2_KREA_STEPS,
IMAGE_V2_DENOISE_DEFAULT,
clampImageV2Denoise,
@@ -118,7 +120,7 @@ function sanitizeImageV2Settings(raw: unknown): ImageV2PresetSettings {
const rec = raw && typeof raw === 'object' ? raw as Record<string, unknown> : {}
const mode = parseImageV2Mode(rec.mode) || 'edit'
const turbo = rec.turbo === true
const engine = mode === 'generate' ? parseImageV2Engine(rec.engine, 'flux') : 'flux'
const engine = parseImageV2Engine(rec.engine, 'flux')
return {
mode,
engine,
@@ -132,7 +134,12 @@ function sanitizeImageV2Settings(raw: unknown): ImageV2PresetSettings {
cfg: engine === 'krea' ? clampImageCfg(rec.cfg, IMAGE_V2_KREA_CFG) : turbo ? 1 : clampImageCfg(rec.cfg, IMAGE_V2_CFG_DEFAULT),
megapixels: mode === 'generate' ? undefined : clampImageScaleMegapixels(rec.megapixels ?? rec.scaleMegapixels, 1),
turbo,
strength: mode === 'refine' ? clampImageV2Denoise(rec.strength, IMAGE_V2_DENOISE_DEFAULT) : undefined,
strength: mode === 'generate' || (mode !== 'refine' && engine !== 'krea')
? undefined
: clampImageV2Denoise(
rec.strength,
mode === 'compose' ? IMAGE_V2_KREA_COMPOSE_DENOISE : mode === 'edit' ? IMAGE_V2_KREA_EDIT_DENOISE : IMAGE_V2_DENOISE_DEFAULT
),
width: mode === 'generate' ? clampImageV2Size(rec.width, rec.height, rec.aspect).width : undefined,
height: mode === 'generate' ? clampImageV2Size(rec.width, rec.height, rec.aspect).height : undefined
}
server/utils/imageChainV2.ts
+8 -12
View File
@@ -65,16 +65,12 @@ export async function runEditV2(job: Job, params: EditV2RunParams) {
await withImageComfyHost(job.imageComfyHost, async () => {
const generate = params.mode === 'generate'
const engine = generate ? parseImageV2Engine(params.engine, 'flux') : 'flux'
if (engine === 'krea' && params.mode !== 'generate') {
throw new Error('Krea is available on Generate first.')
}
const engine = parseImageV2Engine(params.engine, 'flux')
const engineLabel = engine === 'krea' ? 'Krea' : 'Flux'
emitChainJob(job, {
type: 'status',
message: generate
? engine === 'krea'
? 'Queueing Krea generate on Beast...'
: 'Queueing Flux generate on Beast...'
? `Queueing ${engineLabel} generate on Beast...`
: params.mode === 'refine'
? 'Uploading canvas and mask to Beast...'
: params.mode === 'compose' ? 'Uploading stills A and B to Beast...' : 'Uploading still A to Beast...',
@@ -101,10 +97,10 @@ export async function runEditV2(job: Job, params: EditV2RunParams) {
emitChainJob(job, {
type: 'status',
message: params.mode === 'refine'
? 'Queueing Klein v2 refine on Beast...'
? `Queueing ${engineLabel} v2 refine on Beast...`
: params.mode === 'compose'
? `Queueing Klein v2 compose (${params.task}) on Beast...`
: 'Queueing Klein v2 edit on Beast...',
? `Queueing ${engineLabel} v2 compose (${params.task}) on Beast...`
: `Queueing ${engineLabel} v2 edit on Beast...`,
progress: 12
})
}
@@ -165,7 +161,7 @@ export async function runEditV2(job: Job, params: EditV2RunParams) {
isCancelled: () => job.status === 'cancelled'
})
emitChainJob(job, { type: 'status', message: 'Saving Klein v2 still...', progress: 94 })
emitChainJob(job, { type: 'status', message: `Saving ${engineLabel} v2 still...`, progress: 94 })
const buffer = await downloadEditedImage(output)
const size = imageDimensions(buffer)
const still = await saveStill({
@@ -193,7 +189,7 @@ export async function runEditV2(job: Job, params: EditV2RunParams) {
cfg: params.cfg,
megapixels: params.mode === 'generate' ? undefined : params.megapixels,
turbo: params.turbo === true,
strength: params.mode === 'refine' ? params.strength : undefined,
strength: params.mode === 'generate' ? undefined : params.strength,
width: params.mode === 'generate' ? params.width : undefined,
height: params.mode === 'generate' ? params.height : undefined,
seed: params.seed,
server/utils/imageWorkflowV2.ts
+136 -30
View File
@@ -3,6 +3,9 @@ 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 kreaGenerateTemplate from '../assets/krea_v2_generate.json'
import kreaEditTemplate from '../assets/krea_v2_edit.json'
import kreaComposeTemplate from '../assets/krea_v2_compose.json'
import kreaRefineTemplate from '../assets/krea_v2_refine.json'
import { IMAGE_SCALE_TO_TOTAL_PIXELS } from '~/server/utils/comfy'
import { imageComfyFetch } from '~/server/utils/imageComfy'
import { cachedComfyLoraNames } from '~/server/utils/loras'
@@ -16,7 +19,9 @@ import {
IMAGE_V2_KREA_CFG,
IMAGE_V2_KREA_CLIP,
IMAGE_V2_KREA_CLIP_TYPE,
IMAGE_V2_KREA_COMPOSE_DENOISE,
IMAGE_V2_KREA_CONCEPT_LORA,
IMAGE_V2_KREA_EDIT_DENOISE,
IMAGE_V2_KREA_STEPS,
IMAGE_V2_KREA_UNET_CANDIDATES,
IMAGE_V2_KREA_VAE,
@@ -62,6 +67,9 @@ 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 const IMAGE_V2_KREA_GENERATE_WORKFLOW = 'krea_v2_generate.json'
export const IMAGE_V2_KREA_EDIT_WORKFLOW = 'krea_v2_edit.json'
export const IMAGE_V2_KREA_COMPOSE_WORKFLOW = 'krea_v2_compose.json'
export const IMAGE_V2_KREA_REFINE_WORKFLOW = 'krea_v2_refine.json'
export interface ImageV2BuildParams {
mode: ImageV2Mode
@@ -249,6 +257,30 @@ function graphFileNames(graph: WorkflowGraph) {
return names
}
function assertKreaShared(graph: WorkflowGraph) {
if (Object.values(graph).some(node => node.class_type === 'EmptyFlux2LatentImage' || node.class_type === 'Flux2Scheduler' || node.class_type === 'ReferenceLatent')) {
throw createError({ statusCode: 500, statusMessage: 'Krea graph contains a Klein latent, scheduler, or reference node. Refusing to run.' })
}
const clip = graph['5']
if (!clip || clip.class_type !== 'CLIPLoader' || String(clip.inputs.type || '') !== IMAGE_V2_KREA_CLIP_TYPE) {
throw createError({ statusCode: 500, statusMessage: 'Krea graph CLIPLoader type must be krea2.' })
}
const sampler = graph['15']
if (!sampler || sampler.class_type !== 'KSampler') {
throw createError({ statusCode: 500, statusMessage: 'Krea graph is missing KSampler.' })
}
const banned = graphFileNames(graph).filter((name) => {
const id = modelIdentity(name)
return /flux-?2-klein|klein_snofs|flux2-klein-9b-consistency|qwen_3_8b|full_encoder_small_decoder/i.test(id)
})
if (banned.length) {
throw createError({
statusCode: 500,
statusMessage: 'Krea graph includes a Klein model or LoRA. Refusing to run.'
})
}
}
function assertKreaGenerateGraph(graph: WorkflowGraph) {
const loaders = loadImageNames(graph)
if (loaders.length) {
@@ -261,29 +293,29 @@ function assertKreaGenerateGraph(graph: WorkflowGraph) {
if (!latent) {
throw createError({ statusCode: 500, statusMessage: 'Krea generate graph is missing EmptyLatentImage.' })
}
if (Object.values(graph).some(node => node.class_type === 'EmptyFlux2LatentImage' || node.class_type === 'Flux2Scheduler')) {
throw createError({ statusCode: 500, statusMessage: 'Krea generate graph contains a Klein latent or scheduler. Refusing to run.' })
assertKreaShared(graph)
}
function assertKreaItiGraph(graph: WorkflowGraph, mode: ImageV2Mode) {
const loaders = loadImageNames(graph)
if (!loaders.length) {
throw createError({ statusCode: 500, statusMessage: 'Krea image-to-image graph is missing LoadImage.' })
}
const clip = graph['5']
if (!clip || clip.class_type !== 'CLIPLoader' || String(clip.inputs.type || '') !== IMAGE_V2_KREA_CLIP_TYPE) {
throw createError({ statusCode: 500, statusMessage: 'Krea generate graph CLIPLoader type must be krea2.' })
if (!Object.values(graph).some(node => node.class_type === 'VAEEncode')) {
throw createError({ statusCode: 500, statusMessage: 'Krea image-to-image graph is missing VAEEncode.' })
}
const banned = graphFileNames(graph).filter((name) => {
const id = modelIdentity(name)
return /flux-?2-klein|klein_snofs|flux2-klein-9b-consistency|qwen_3_8b|full_encoder_small_decoder/i.test(id)
})
if (banned.length) {
throw createError({
statusCode: 500,
statusMessage: 'Krea generate graph includes a Klein model or LoRA. Refusing to run.'
})
if (mode === 'compose' && !Object.values(graph).some(node => node.class_type === 'ImageStitch')) {
throw createError({ statusCode: 500, statusMessage: 'Krea compose graph is missing ImageStitch. Refusing to drop still B.' })
}
assertKreaShared(graph)
}
export function assertImageV2Graph(graph: WorkflowGraph, mode: ImageV2Mode, imageBName?: string, engine: ImageV2Engine = 'flux') {
const loaders = loadImageNames(graph)
if (mode === 'generate' && engine === 'krea') {
assertKreaGenerateGraph(graph)
} else if (engine === 'krea') {
assertKreaItiGraph(graph, mode)
} else if (mode === 'generate') {
if (loaders.length) {
throw createError({
@@ -310,7 +342,7 @@ export function assertImageV2Graph(graph: WorkflowGraph, mode: ImageV2Mode, imag
statusMessage: 'Refine graph has no mask input. Refusing to run.'
})
}
if (!('denoise' in (graph[SCHEDULER]?.inputs || {}))) {
if (!('denoise' in (graph[SCHEDULER]?.inputs || {})) && !('denoise' in (graph['15']?.inputs || {}))) {
throw createError({
statusCode: 500,
statusMessage: 'Refine graph has no denoise on the sampler. Refusing to run.'
@@ -367,22 +399,99 @@ function applyKreaConceptLora(graph: WorkflowGraph, loraName: string, model: num
setInput(graph, '15', 'model', ['7', 0])
}
function buildKreaGenerateWorkflow(params: ImageV2BuildParams) {
function applyKreaConceptIfNeeded(graph: WorkflowGraph, params: ImageV2BuildParams) {
const conceptName = isXaigenStudio() ? String(params.kreaConceptLora || '').trim() : ''
const snofsModel = conceptName ? clampImageV2Strength(params.snofsModel, 0) : 0
const snofsClip = conceptName ? clampImageV2Strength(params.snofsClip, 0) : 0
if (!conceptName || (snofsModel <= 0 && snofsClip <= 0)) return { snofsModel: 0, snofsClip: 0 }
const resolved = resolveRequiredLora(conceptName, 'Krea concept')
if (/klein_snofs|flux2-klein-9b-consistency|flux-?2-klein/i.test(loraIdentityKey(resolved))) {
throw createError({ statusCode: 500, statusMessage: 'Refusing to load a Klein LoRA into the Krea graph.' })
}
applyKreaConceptLora(graph, resolved, snofsModel, snofsClip)
return { snofsModel, snofsClip, name: resolved }
}
function requireKreaModels(params: ImageV2BuildParams) {
if (!params.kreaUnetName || !params.kreaClipName || !params.kreaVaeName) {
throw createError({
statusCode: 503,
statusMessage: 'Krea models were not resolved. Refusing to fall back to Flux.'
})
}
}
function buildKreaItiWorkflow(params: ImageV2BuildParams) {
requireKreaModels(params)
const refine = params.mode === 'refine'
const compose = params.mode === 'compose'
if (refine && !String(params.maskName || '').trim()) {
throw createError({ statusCode: 400, statusMessage: 'Refine requires a mask. Refusing to fall back to Edit.' })
}
if (compose && !String(params.imageBName || '').trim()) {
throw createError({ statusCode: 400, statusMessage: 'Compose requires still B. Refusing to fall back to Edit.' })
}
const graph = structuredClone(
refine ? kreaRefineTemplate : compose ? kreaComposeTemplate : kreaEditTemplate
) as WorkflowGraph
const prompt = composeImageV2Prompt(params.mode, params.task, params.prompt)
const negative = String(params.negative || '')
const steps = clampImageSteps(params.steps, IMAGE_V2_KREA_STEPS)
const cfg = clampImageCfg(params.cfg, IMAGE_V2_KREA_CFG)
const denoise = clampImageV2Denoise(
params.strength,
refine ? IMAGE_V2_DENOISE_DEFAULT : compose ? IMAGE_V2_KREA_COMPOSE_DENOISE : IMAGE_V2_KREA_EDIT_DENOISE
)
const workflowFile = imageV2WorkflowFile(params.mode, 'krea')
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, '4', 'unet_name', params.kreaUnetName)
setInput(graph, '5', 'clip_name', params.kreaClipName)
setInput(graph, '5', 'type', IMAGE_V2_KREA_CLIP_TYPE)
setInput(graph, '6', 'vae_name', params.kreaVaeName)
setInput(graph, PROMPT, 'text', prompt)
setInput(graph, NEGATIVE, 'text', negative)
setInput(graph, '15', 'seed', params.seed)
setInput(graph, '15', 'steps', steps)
setInput(graph, '15', 'cfg', cfg)
setInput(graph, '15', 'sampler_name', 'euler')
setInput(graph, '15', 'scheduler', 'simple')
setInput(graph, '15', 'denoise', denoise)
setInput(graph, SAVE, 'filename_prefix', params.filenamePrefix || `aigen-v2-krea-${params.mode}`)
patchScaleMegapixels(graph, params.megapixels ?? 1)
const concept = applyKreaConceptIfNeeded(graph, params)
assertImageV2Graph(graph, params.mode, compose ? params.imageBName : undefined, 'krea')
const loaders = loadImageNames(graph)
console.log(JSON.stringify({
src: 'image-v2',
engine: 'krea',
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: denoise,
loadImage: Object.fromEntries(loaders.map(item => [item.id, { title: item.title, file: item.image }])),
loras: concept.name ? { concept: { name: concept.name, model: concept.snofsModel, clip: concept.snofsClip } } : { loaded: false },
steps,
cfg,
seed: params.seed,
megapixels: clampImageScaleMegapixels(params.megapixels ?? 1)
}))
return { graph, workflowFile, loaders, prompt, strength: denoise }
}
function buildKreaGenerateWorkflow(params: ImageV2BuildParams) {
requireKreaModels(params)
const graph = structuredClone(kreaGenerateTemplate) as WorkflowGraph
const prompt = composeImageV2Prompt('generate', 't2i', params.prompt)
const negative = String(params.negative || '')
const size = clampImageV2Size(params.width, params.height)
const steps = clampImageSteps(params.steps, IMAGE_V2_KREA_STEPS)
const cfg = clampImageCfg(params.cfg, IMAGE_V2_KREA_CFG)
const conceptName = isXaigenStudio() ? String(params.kreaConceptLora || '').trim() : ''
const snofsModel = conceptName ? clampImageV2Strength(params.snofsModel, 0) : 0
const snofsClip = conceptName ? clampImageV2Strength(params.snofsClip, 0) : 0
const workflowFile = IMAGE_V2_KREA_GENERATE_WORKFLOW
setInput(graph, '4', 'unet_name', params.kreaUnetName)
@@ -399,13 +508,7 @@ function buildKreaGenerateWorkflow(params: ImageV2BuildParams) {
setInput(graph, '15', 'sampler_name', 'euler')
setInput(graph, '15', 'scheduler', 'simple')
setInput(graph, SAVE, 'filename_prefix', params.filenamePrefix || 'aigen-v2-krea')
if (conceptName && (snofsModel > 0 || snofsClip > 0)) {
const resolved = resolveRequiredLora(conceptName, 'Krea concept')
if (/klein_snofs|flux2-klein-9b-consistency|flux-?2-klein/i.test(loraIdentityKey(resolved))) {
throw createError({ statusCode: 500, statusMessage: 'Refusing to load a Klein LoRA into the Krea graph.' })
}
applyKreaConceptLora(graph, resolved, snofsModel, snofsClip)
}
const concept = applyKreaConceptIfNeeded(graph, params)
assertImageV2Graph(graph, 'generate', undefined, 'krea')
const loaders = loadImageNames(graph)
@@ -422,8 +525,8 @@ function buildKreaGenerateWorkflow(params: ImageV2BuildParams) {
vae: params.kreaVaeName
},
loras: {
concept: graph['7']
? { name: graph['7'].inputs.lora_name, model: snofsModel, clip: snofsClip }
concept: concept.name
? { name: concept.name, model: concept.snofsModel, clip: concept.snofsClip }
: { loaded: false, model: 0, clip: 0 }
},
steps,
@@ -434,8 +537,10 @@ function buildKreaGenerateWorkflow(params: ImageV2BuildParams) {
}
export function buildImageV2Workflow(params: ImageV2BuildParams) {
const engine = params.mode === 'generate' ? parseImageV2Engine(params.engine, 'flux') : 'flux'
if (params.mode === 'generate' && engine === 'krea') return buildKreaGenerateWorkflow(params)
const engine = parseImageV2Engine(params.engine, 'flux')
if (engine === 'krea') {
return params.mode === 'generate' ? buildKreaGenerateWorkflow(params) : buildKreaItiWorkflow(params)
}
const compose = params.mode === 'compose'
const refine = params.mode === 'refine'
const generate = params.mode === 'generate'
@@ -536,6 +641,7 @@ export const IMAGE_V2_NODE_LABELS: Record<string, string> = {
'30': 'Loading mask',
'2': 'Scaling image A',
'23': 'Scaling image B',
'27': 'Stitching stills',
'31': 'Resizing mask',
'4': 'Loading diffusion model',
'5': 'Loading CLIP',
server/utils/studioQueue.ts
+3 -3
View File
@@ -754,7 +754,7 @@ async function startStudioEditJob(item: StudioJob) {
await markStudioLive(item.ownerKey, item.id, live.id)
void runEditV2(live, {
mode,
engine: mode === 'generate' && payload.engine === 'krea' ? 'krea' : 'flux',
engine: payload.engine === 'krea' ? 'krea' : 'flux',
task: payload.v2Task || 'scene',
image: mode === 'generate' ? null : image,
reference: mode === 'refine' || mode === 'generate' ? null : reference,
@@ -766,8 +766,8 @@ async function startStudioEditJob(item: StudioJob) {
cfg: payload.cfg,
snofsModel: payload.snofsModel ?? (isXaigenStudio() ? 0.65 : 0),
snofsClip: payload.snofsClip ?? (isXaigenStudio() ? 0.35 : 0),
consistencyModel: payload.consistencyModel ?? (mode === 'generate' ? 0 : 0.7),
consistencyClip: payload.consistencyClip ?? (mode === 'generate' ? 0 : 0.7),
consistencyModel: payload.consistencyModel ?? (mode === 'generate' || payload.engine === 'krea' ? 0 : 0.7),
consistencyClip: payload.consistencyClip ?? (mode === 'generate' || payload.engine === 'krea' ? 0 : 0.7),
megapixels: payload.scaleMegapixels ?? 1,
strength: payload.refineStrength,
width: payload.width,