Add Krea as a second Image v2 Generate engine beside Flux.

Generate can queue krea_v2_generate.json without sharing Klein encoders or falling back if Krea assets are missing.

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

server/utils/generationPresets.ts

@@ -21,10 +21,13 @@ import {
   IMAGE_V2_SNOFS_CLIP,
   IMAGE_V2_SNOFS_MODEL,
   IMAGE_V2_STEPS_DEFAULT,
+  IMAGE_V2_KREA_CFG,
+  IMAGE_V2_KREA_STEPS,
   IMAGE_V2_DENOISE_DEFAULT,
   clampImageV2Denoise,
   clampImageV2Size,
   clampImageV2Strength,
+  parseImageV2Engine,
   parseImageV2Mode,
   parseImageV2Task,
   type ImageV2PresetSettings
@@ -115,16 +118,18 @@ 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'
   return {
     mode,
+    engine,
     task: mode === 'compose' ? parseImageV2Task(rec.task, 'scene') : mode === 'refine' ? 'refine' : mode === 'generate' ? 't2i' : '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),
     consistencyModel: clampImageV2Strength(rec.consistencyModel ?? rec.consistency_model, mode === 'generate' ? 0 : IMAGE_V2_CONSISTENCY_MODEL),
     consistencyClip: clampImageV2Strength(rec.consistencyClip ?? rec.consistency_clip, mode === 'generate' ? 0 : IMAGE_V2_CONSISTENCY_CLIP),
-    steps: turbo ? 8 : clampImageSteps(rec.steps, IMAGE_V2_STEPS_DEFAULT),
-    cfg: turbo ? 1 : clampImageCfg(rec.cfg, IMAGE_V2_CFG_DEFAULT),
+    steps: engine === 'krea' ? clampImageSteps(rec.steps, IMAGE_V2_KREA_STEPS) : turbo ? 8 : clampImageSteps(rec.steps, IMAGE_V2_STEPS_DEFAULT),
+    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,

server/utils/imageChainV2.ts

@@ -2,16 +2,17 @@ import { createJob, emitJob, type Job } from '~/server/utils/jobs'
 import { ensureComfyReady } from '~/server/utils/comfyLifecycle'
 import { assertImageScaleToTotalPixelsNode, getComfyHost, uploadImage, queuePrompt, purgeComfyArtifacts } from '~/server/utils/comfy'
 import { withImageComfyHost, waitForImageEdit, downloadEditedImage } from '~/server/utils/imageComfy'
-import { buildImageV2Workflow } from '~/server/utils/imageWorkflowV2'
+import { buildImageV2Workflow, resolveKreaGenerateAssets } from '~/server/utils/imageWorkflowV2'
 import { ensureComfyLoraNames } from '~/server/utils/loras'
 import { imageDimensions } from '~/server/utils/resolution'
 import { emitChainJob } from '~/server/utils/watch'
 import { saveStill } from '~/server/utils/library'
-import type { ImageV2Mode, ImageV2Task } from '~/utils/imageV2'
+import { parseImageV2Engine, type ImageV2Engine, type ImageV2Mode, type ImageV2Task } from '~/utils/imageV2'
 import type { EditImageFile } from '~/server/utils/imageChain'
 
 export type EditV2RunParams = {
   mode: ImageV2Mode
+  engine?: ImageV2Engine
   task: ImageV2Task
   image?: EditImageFile | null
   reference: EditImageFile | null
@@ -64,10 +65,16 @@ 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.')
+      }
       emitChainJob(job, {
         type: 'status',
         message: generate
-          ? 'Queueing Klein v2 generate on Beast...'
+          ? engine === 'krea'
+            ? 'Queueing Krea generate on Beast...'
+            : 'Queueing Flux 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...',
@@ -103,8 +110,10 @@ export async function runEditV2(job: Job, params: EditV2RunParams) {
       }
       await ensureComfyLoraNames('image')
       if (!generate) await assertImageScaleToTotalPixelsNode()
+      const krea = engine === 'krea' ? await resolveKreaGenerateAssets() : null
       const built = buildImageV2Workflow({
         mode: params.mode,
+        engine,
         task: params.task,
         prompt: params.prompt,
         negative: params.negative,
@@ -123,7 +132,11 @@ export async function runEditV2(job: Job, params: EditV2RunParams) {
         width: params.width,
         height: params.height,
         turbo: params.turbo === true,
-        filenamePrefix: `aigen_v2_${job.id.slice(0, 8)}`
+        filenamePrefix: `aigen_v2_${job.id.slice(0, 8)}`,
+        kreaUnetName: krea?.unet,
+        kreaClipName: krea?.clip,
+        kreaVaeName: krea?.vae,
+        kreaConceptLora: krea?.conceptLora
       })
       const queued = await queuePrompt(built.graph, job.clientId)
       job.promptId = queued.prompt_id
@@ -169,6 +182,7 @@ export async function runEditV2(job: Job, params: EditV2RunParams) {
         settings: {
           kind: 'imagev2',
           mode: params.mode,
+          engine,
           task: params.task,
           negative: params.negative,
           snofsModel: params.snofsModel,

server/utils/imageWorkflowV2.ts

@@ -2,7 +2,9 @@ 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 kreaGenerateTemplate from '../assets/krea_v2_generate.json'
 import { IMAGE_SCALE_TO_TOTAL_PIXELS } from '~/server/utils/comfy'
+import { imageComfyFetch } from '~/server/utils/imageComfy'
 import { cachedComfyLoraNames } from '~/server/utils/loras'
 import { resolveComfyLoraName, loraIdentityKey } from '~/utils/loras'
 import { clampImageCfg, clampImageScaleMegapixels, clampImageSteps } from '~/utils/generationPresets'
@@ -10,6 +12,13 @@ import {
   IMAGE_V2_CONSISTENCY_CLIP,
   IMAGE_V2_CONSISTENCY_LORA,
   IMAGE_V2_CONSISTENCY_MODEL,
+  IMAGE_V2_KREA_CFG,
+  IMAGE_V2_KREA_CLIP,
+  IMAGE_V2_KREA_CLIP_TYPE,
+  IMAGE_V2_KREA_CONCEPT_LORA,
+  IMAGE_V2_KREA_STEPS,
+  IMAGE_V2_KREA_UNET_CANDIDATES,
+  IMAGE_V2_KREA_VAE,
   IMAGE_V2_SNOFS_CLIP,
   IMAGE_V2_SNOFS_LORA,
   IMAGE_V2_SNOFS_MODEL,
@@ -20,6 +29,9 @@ import {
   clampImageV2Size,
   clampImageV2Strength,
   composeImageV2Prompt,
+  imageV2WorkflowFile,
+  parseImageV2Engine,
+  type ImageV2Engine,
   type ImageV2Mode,
   type ImageV2Task
 } from '~/utils/imageV2'
@@ -46,10 +58,12 @@ 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 const IMAGE_V2_KREA_GENERATE_WORKFLOW = 'krea_v2_generate.json'
 
 export interface ImageV2BuildParams {
   mode: ImageV2Mode
   task: ImageV2Task
+  engine?: ImageV2Engine
   prompt: string
   negative?: string
   imageAName?: string
@@ -68,6 +82,10 @@ export interface ImageV2BuildParams {
   seed: number
   megapixels?: number
   filenamePrefix?: string
+  kreaUnetName?: string
+  kreaClipName?: string
+  kreaVaeName?: string
+  kreaConceptLora?: string
 }
 
 function setInput(graph: WorkflowGraph, id: string, key: string, value: unknown) {
@@ -120,9 +138,140 @@ function bypassLoraNode(graph: WorkflowGraph, id: string, modelFrom: string, cli
   }
 }
 
-export function assertImageV2Graph(graph: WorkflowGraph, mode: ImageV2Mode, imageBName?: string) {
+function modelBasename(name: string) {
+  return String(name || '').replace(/\\/g, '/').split('/').pop() || ''
+}
+
+function modelIdentity(name: string) {
+  return modelBasename(name).toLowerCase()
+}
+
+async function listComfyModels(folder: string) {
+  const res = await imageComfyFetch(`/models/${encodeURIComponent(folder)}`, {
+    signal: AbortSignal.timeout(8000)
+  }).catch(() => null)
+  if (!res || !res.ok) return [] as string[]
+  const payload = await res.json().catch(() => null)
+  if (Array.isArray(payload)) return payload.map(item => String(item || '')).filter(Boolean)
+  return [] as string[]
+}
+
+function pickComfyModel(available: string[], candidates: string[]) {
+  for (const wanted of candidates) {
+    const id = modelIdentity(wanted)
+    const exact = available.find(name => modelIdentity(name) === id)
+    if (exact) return exact
+    const stem = id.replace(/\.(safetensors|sft|ckpt|pt|bin)$/i, '')
+    const prefix = available.find(name => modelIdentity(name).startsWith(stem))
+    if (prefix) return prefix
+  }
+  return ''
+}
+
+export function configuredKreaConceptLora() {
+  const config = useRuntimeConfig()
+  return String(
+    config.kreaConceptLora
+    || (config.public as { kreaConceptLora?: string } | undefined)?.kreaConceptLora
+    || IMAGE_V2_KREA_CONCEPT_LORA
+    || ''
+  ).trim()
+}
+
+export async function resolveKreaGenerateAssets() {
+  const [unets, clips, vaes, clipInfo] = await Promise.all([
+    listComfyModels('diffusion_models'),
+    listComfyModels('text_encoders').then(async (names) => names.length ? names : listComfyModels('clip')),
+    listComfyModels('vae'),
+    imageComfyFetch('/object_info/CLIPLoader', { signal: AbortSignal.timeout(8000) })
+      .then(async (res) => res.ok ? res.json() : null)
+      .catch(() => null)
+  ])
+  const unet = pickComfyModel(unets, [...IMAGE_V2_KREA_UNET_CANDIDATES])
+  if (!unet) {
+    throw createError({
+      statusCode: 503,
+      statusMessage: 'Krea diffusion model is missing on Beast Comfy (need krea2_turbo_mxfp8, nvfp4, or fp8_scaled).'
+    })
+  }
+  const clip = pickComfyModel(clips, [IMAGE_V2_KREA_CLIP])
+  if (!clip) {
+    throw createError({
+      statusCode: 503,
+      statusMessage: `Krea text encoder is missing on Beast Comfy (${IMAGE_V2_KREA_CLIP}).`
+    })
+  }
+  const vae = pickComfyModel(vaes, [IMAGE_V2_KREA_VAE])
+  if (!vae) {
+    throw createError({
+      statusCode: 503,
+      statusMessage: `Qwen image VAE is missing on Beast Comfy (${IMAGE_V2_KREA_VAE}).`
+    })
+  }
+  const rec = clipInfo && typeof clipInfo === 'object' ? clipInfo as Record<string, unknown> : {}
+  const body = (rec.CLIPLoader && typeof rec.CLIPLoader === 'object' ? rec.CLIPLoader : rec) as {
+    input?: { required?: { type?: unknown } }
+  }
+  const typeSpec = body.input?.required?.type
+  const types = Array.isArray(typeSpec) && Array.isArray(typeSpec[0])
+    ? (typeSpec[0] as unknown[]).map(item => String(item || ''))
+    : []
+  if (types.length && !types.includes(IMAGE_V2_KREA_CLIP_TYPE)) {
+    throw createError({
+      statusCode: 503,
+      statusMessage: 'Beast Comfy CLIPLoader has no krea2 type. Update Comfy before using Krea.'
+    })
+  }
+  return { unet, clip, vae, conceptLora: configuredKreaConceptLora() }
+}
+
+function graphFileNames(graph: WorkflowGraph) {
+  const names: string[] = []
+  for (const node of Object.values(graph)) {
+    for (const key of ['unet_name', 'clip_name', 'vae_name', 'lora_name']) {
+      const value = node.inputs[key]
+      if (typeof value === 'string' && value.trim()) names.push(value)
+    }
+  }
+  return names
+}
+
+function assertKreaGenerateGraph(graph: WorkflowGraph) {
   const loaders = loadImageNames(graph)
-  if (mode === 'generate') {
+  if (loaders.length) {
+    throw createError({
+      statusCode: 500,
+      statusMessage: 'Krea generate graph has a required LoadImage. Refusing to run.'
+    })
+  }
+  const latent = Object.values(graph).find(node => node.class_type === 'EmptyLatentImage')
+  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.' })
+  }
+  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.' })
+  }
+  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.'
+    })
+  }
+}
+
+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 (mode === 'generate') {
     if (loaders.length) {
       throw createError({
         statusCode: 500,
@@ -188,7 +337,92 @@ export function assertImageV2Graph(graph: WorkflowGraph, mode: ImageV2Mode, imag
   }
 }
 
+function applyKreaConceptLora(graph: WorkflowGraph, loraName: string, model: number, clip: number) {
+  graph['7'] = {
+    class_type: 'LoraLoader',
+    inputs: {
+      lora_name: loraName,
+      strength_model: model,
+      strength_clip: clip,
+      model: ['4', 0],
+      clip: ['5', 0]
+    },
+    _meta: { title: 'Concept LoRA' }
+  }
+  setInput(graph, PROMPT, 'clip', ['7', 1])
+  setInput(graph, NEGATIVE, 'clip', ['7', 1])
+  setInput(graph, '15', 'model', ['7', 0])
+}
+
+function buildKreaGenerateWorkflow(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.'
+    })
+  }
+  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 = 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)
+  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, '14', 'width', size.width)
+  setInput(graph, '14', 'height', size.height)
+  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, 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)
+  }
+
+  assertImageV2Graph(graph, 'generate', undefined, 'krea')
+  const loaders = loadImageNames(graph)
+  console.log(JSON.stringify({
+    src: 'image-v2',
+    engine: 'krea',
+    workflow: workflowFile,
+    mode: 'generate',
+    task: 't2i',
+    size: { width: size.width, height: size.height },
+    models: {
+      unet: params.kreaUnetName,
+      clip: params.kreaClipName,
+      vae: params.kreaVaeName
+    },
+    loras: {
+      concept: graph['7']
+        ? { name: graph['7'].inputs.lora_name, model: snofsModel, clip: snofsClip }
+        : { loaded: false, model: 0, clip: 0 }
+    },
+    steps,
+    cfg,
+    seed: params.seed
+  }))
+  return { graph, workflowFile, loaders, prompt, strength: undefined }
+}
+
 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 compose = params.mode === 'compose'
   const refine = params.mode === 'refine'
   const generate = params.mode === 'generate'
@@ -214,9 +448,7 @@ export function buildImageV2Workflow(params: ImageV2BuildParams) {
   const cfg = clampImageCfg(params.cfg, 4)
   const strength = refine ? clampImageV2Denoise(params.strength, IMAGE_V2_DENOISE_DEFAULT) : undefined
   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
+  const workflowFile = imageV2WorkflowFile(params.mode, engine)
 
   if (!generate) setInput(graph, LOAD_A, 'image', params.imageAName || '')
   if (compose) setInput(graph, LOAD_B, 'image', params.imageBName || '')
@@ -249,11 +481,12 @@ export function buildImageV2Workflow(params: ImageV2BuildParams) {
     setInput(graph, CONSISTENCY, 'strength_clip', consistencyClip)
   }
 
-  assertImageV2Graph(graph, params.mode, generate || refine ? undefined : params.imageBName)
+  assertImageV2Graph(graph, params.mode, generate || refine ? undefined : params.imageBName, engine)
 
   const loaders = loadImageNames(graph)
   console.log(JSON.stringify({
     src: 'image-v2',
+    engine,
     workflow: workflowFile,
     mode: params.mode,
     task: params.task,
@@ -285,16 +518,17 @@ export const IMAGE_V2_NODE_LABELS: Record<string, string> = {
   '2': 'Scaling image A',
   '23': 'Scaling image B',
   '31': 'Resizing mask',
-  '4': 'Loading Flux.2 Klein 9B',
+  '4': 'Loading diffusion model',
   '5': 'Loading CLIP',
   '6': 'Loading VAE',
-  '7': 'Applying SNOFS',
+  '7': 'Applying Concept LoRA',
   '8': 'Applying Consistency',
   '9': 'Encoding prompt',
   '11': 'Encoding image A',
   '24': 'Encoding image B',
   '33': 'Applying mask',
-  '14': 'Building empty Klein latent',
+  '14': 'Building empty latent',
+  '15': 'Sampling',
   '19': 'Sampling Klein v2',
   '20': 'Decoding still',
   '21': 'Saving still'

server/utils/studioQueue.ts

@@ -54,6 +54,7 @@ export interface StudioJobPayload {
   imagePipeline?: 'v1' | 'v2'
   v2Mode?: 'edit' | 'compose' | 'refine' | 'generate'
   v2Task?: 'scene' | 'identity' | 'outfit' | 'face_lock' | 'refine' | 't2i'
+  engine?: 'flux' | 'krea'
   snofsModel?: number
   snofsClip?: number
   consistencyModel?: number
@@ -753,6 +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',
         task: payload.v2Task || 'scene',
         image: mode === 'generate' ? null : image,
         reference: mode === 'refine' || mode === 'generate' ? null : reference,