Wire every Image v2 LoRA, including Concept, through one sampler chain.

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

server/utils/imageWorkflowV2.ts

@@ -8,8 +8,8 @@ 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 { applyImageV2UserLoras, cachedComfyLoraNames } from '~/server/utils/loras'
+import { applyImageV2UserLoras } from '~/server/utils/loras'
-import { resolveComfyLoraName, loraIdentityKey, type LoraStackItem } from '~/utils/loras'
+import { type LoraStackItem } from '~/utils/loras'
 import { clampImageCfg, clampImageScaleMegapixels, clampImageSteps } from '~/utils/generationPresets'
 import { isXaigenStudio } from '~/utils/videoModels'
 import {
@@ -27,11 +27,8 @@ import {
   IMAGE_V2_GENERATE_WIDTH,
   clampImageV2Denoise,
   clampImageV2Size,
-  clampImageV2Strength,
   composeImageV2Prompt,
   imageV2WorkflowFile,
-  isImageV2ConceptLora,
-  isImageV2ConsistencyLora,
   mergeLegacyImageV2Loras,
   parseImageV2Engine,
   type ImageV2Engine,
@@ -46,19 +43,12 @@ 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'
 const NEGATIVE = '10'
 const NOISE = '15'
 const SCHEDULER = '17'
 const CFG = '18'
 const SAVE = '21'
-const UNET = '4'
-const CLIP = '5'
-const SNOFS = '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'
@@ -111,16 +101,6 @@ function loadImageNames(graph: WorkflowGraph) {
     }))
 }
 
-function resolveRequiredLora(wanted: string, label: string) {
-  const names = cachedComfyLoraNames('image')
-  const resolved = names.length ? resolveComfyLoraName(wanted, names) : wanted
-  const hit = names.some(name => loraIdentityKey(name) === loraIdentityKey(wanted) || loraIdentityKey(name) === loraIdentityKey(resolved))
-  if (names.length && !hit) {
-    throw createError({ statusCode: 503, statusMessage: `Missing ${label} LoRA (${wanted}) on Beast Comfy` })
-  }
-  return resolved
-}
-
 function patchScaleMegapixels(graph: WorkflowGraph, megapixels: number) {
   const mp = clampImageScaleMegapixels(megapixels)
   for (const node of Object.values(graph)) {
@@ -137,26 +117,6 @@ 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]
-    }
-  }
-}
-
-function bypassLoraToBaseLoaders(graph: WorkflowGraph, id: 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 ? [CLIP, 0] : [UNET, 0]
-    }
-  }
-}
-
 function modelBasename(name: string) {
   return String(name || '').replace(/\\/g, '/').split('/').pop() || ''
 }
@@ -380,47 +340,6 @@ 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 applyKreaConceptIfNeeded(graph: WorkflowGraph, params: ImageV2BuildParams) {
-  const merged = mergeLegacyImageV2Loras(params.loraStack, {
-    engine: 'krea',
-    mode: params.mode,
-    xaigen: isXaigenStudio(),
-    snofsModel: params.snofsModel,
-    snofsClip: params.snofsClip,
-    kreaConcept: params.kreaConceptLora
-  })
-  const conceptName = isXaigenStudio() ? String(params.kreaConceptLora || '').trim() : ''
-  if (conceptName && merged.some(item => loraIdentityKey(item.name) === loraIdentityKey(conceptName))) {
-    return { snofsModel: 0, snofsClip: 0 }
-  }
-  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 imageV2UserStack(params: ImageV2BuildParams, engine: ImageV2Engine) {
   return mergeLegacyImageV2Loras(params.loraStack, {
     engine,
@@ -434,57 +353,6 @@ function imageV2UserStack(params: ImageV2BuildParams, engine: ImageV2Engine) {
   })
 }
 
-function loraIsOn(item?: LoraStackItem | null) {
-  return Boolean(item && (item.strengthModel !== 0 || item.strengthClip !== 0))
-}
-
-function applyFluxTemplateLoras(graph: WorkflowGraph, params: ImageV2BuildParams) {
-  const xaigen = isXaigenStudio()
-  const stack = imageV2UserStack(params, 'flux')
-  const concept = xaigen ? stack.find(item => isImageV2ConceptLora(item.name)) : undefined
-  const consistency = stack.find(item => isImageV2ConsistencyLora(item.name))
-  const extras = stack.filter(item => !isImageV2ConceptLora(item.name) && !isImageV2ConsistencyLora(item.name))
-
-  if (loraIsOn(concept) && graph[SNOFS]) {
-    setInput(graph, SNOFS, 'lora_name', resolveRequiredLora(concept!.name, 'Concept LoRA'))
-    setInput(graph, SNOFS, 'strength_model', concept!.strengthModel)
-    setInput(graph, SNOFS, 'strength_clip', concept!.strengthClip)
-  } else if (graph[SNOFS]) {
-    bypassLoraToBaseLoaders(graph, SNOFS)
-  }
-
-  if (loraIsOn(consistency) && graph[CONSISTENCY]) {
-    setInput(graph, CONSISTENCY, 'lora_name', resolveRequiredLora(consistency!.name, 'Consistency'))
-    setInput(graph, CONSISTENCY, 'strength_model', consistency!.strengthModel)
-    setInput(graph, CONSISTENCY, 'strength_clip', consistency!.strengthClip)
-    if (!graph[SNOFS]) {
-      graph[CONSISTENCY].inputs.model = [UNET, 0]
-      graph[CONSISTENCY].inputs.clip = [CLIP, 0]
-    }
-  } else if (graph[CONSISTENCY]) {
-    if (graph[SNOFS]) bypassLoraNode(graph, CONSISTENCY, SNOFS, SNOFS)
-    else bypassLoraToBaseLoaders(graph, CONSISTENCY)
-  }
-
-  const extraApplied = applyImageV2UserLoras(graph, extras, 'flux')
-  const guider = Object.values(graph).find(node => node.class_type === 'CFGGuider' || node.class_type === 'BasicGuider')
-  const modelFeed = Array.isArray(guider?.inputs.model) ? String(guider.inputs.model[0] || '') : ''
-  const live = stack.filter(loraIsOn)
-  if (live.length && (modelFeed === UNET || !modelFeed)) {
-    throw createError({
-      statusCode: 500,
-      statusMessage: 'Image v2 LoRA chain is not connected to the sampler. Refusing to run Klein without those adapters.'
-    })
-  }
-  return {
-    stack,
-    concept: loraIsOn(concept) ? concept : undefined,
-    consistency: loraIsOn(consistency) ? consistency : undefined,
-    extras: extraApplied,
-    modelFeed
-  }
-}
-
 function requireKreaModels(params: ImageV2BuildParams) {
   if (!params.kreaUnetName || !params.kreaClipName || !params.kreaVaeName) {
     throw createError({
@@ -534,7 +402,6 @@ function buildKreaItiWorkflow(params: ImageV2BuildParams) {
   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)
   const userLoras = applyImageV2UserLoras(graph, imageV2UserStack(params, 'krea'), 'krea')
 
   assertImageV2Graph(graph, params.mode, compose ? params.imageBName : undefined, 'krea')
@@ -549,10 +416,7 @@ function buildKreaItiWorkflow(params: ImageV2BuildParams) {
     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: {
+    loras: { user: userLoras },
-      concept: concept.name ? { name: concept.name, model: concept.snofsModel, clip: concept.snofsClip } : { loaded: false },
-      user: userLoras
-    },
     steps,
     cfg,
     seed: params.seed,
@@ -585,7 +449,6 @@ 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')
-  const concept = applyKreaConceptIfNeeded(graph, params)
   const userLoras = applyImageV2UserLoras(graph, imageV2UserStack(params, 'krea'), 'krea')
 
   assertImageV2Graph(graph, 'generate', undefined, 'krea')
@@ -602,12 +465,7 @@ function buildKreaGenerateWorkflow(params: ImageV2BuildParams) {
       clip: params.kreaClipName,
       vae: params.kreaVaeName
     },
-    loras: {
+    loras: { user: userLoras },
-      concept: concept.name
-        ? { name: concept.name, model: concept.snofsModel, clip: concept.snofsClip }
-        : { loaded: false, model: 0, clip: 0 },
-      user: userLoras
-    },
     steps,
     cfg,
     seed: params.seed
@@ -657,7 +515,9 @@ export function buildImageV2Workflow(params: ImageV2BuildParams) {
   setInput(graph, SAVE, 'filename_prefix', params.filenamePrefix || (generate ? 'aigen-v2-generate' : 'aigen-v2'))
   if (!generate) patchScaleMegapixels(graph, params.megapixels ?? 1)
 
-  const fluxLoras = applyFluxTemplateLoras(graph, params)
+  const userLoras = applyImageV2UserLoras(graph, imageV2UserStack(params, 'flux'), 'flux')
+  const guider = Object.values(graph).find(node => node.class_type === 'CFGGuider' || node.class_type === 'BasicGuider')
+  const modelFeed = Array.isArray(guider?.inputs.model) ? String(guider.inputs.model[0] || '') : ''
 
   assertImageV2Graph(graph, params.mode, generate || refine ? undefined : params.imageBName, engine)
 
@@ -674,16 +534,7 @@ export function buildImageV2Workflow(params: ImageV2BuildParams) {
     strength,
     turbo: params.turbo === true,
     loadImage: Object.fromEntries(loaders.map(item => [item.id, { title: item.title, file: item.image }])),
-    loras: {
+    loras: { user: userLoras, modelFeed },
-      snofs: fluxLoras.concept
-        ? { name: graph[SNOFS]?.inputs.lora_name, model: fluxLoras.concept.strengthModel, clip: fluxLoras.concept.strengthClip }
-        : { loaded: false, model: 0, clip: 0 },
-      consistency: fluxLoras.consistency
-        ? { name: graph[CONSISTENCY]?.inputs.lora_name, model: fluxLoras.consistency.strengthModel, clip: fluxLoras.consistency.strengthClip }
-        : { loaded: false, model: 0, clip: 0 },
-      user: fluxLoras.extras,
-      modelFeed: fluxLoras.modelFeed
-    },
     steps,
     cfg,
     seed: params.seed,
@@ -706,10 +557,14 @@ export const IMAGE_V2_NODE_LABELS: Record<string, string> = {
   '6': 'Loading VAE',
   '7': 'Applying Concept LoRA',
   '8': 'Applying Consistency',
-  'user:lora:0': 'Applying LoRA',
+  '70': 'Applying LoRA',
-  'user:lora:1': 'Applying LoRA',
+  '71': 'Applying LoRA',
-  'user:lora:2': 'Applying LoRA',
+  '72': 'Applying LoRA',
-  'user:lora:3': 'Applying LoRA',
+  '73': 'Applying LoRA',
+  '74': 'Applying LoRA',
+  '75': 'Applying LoRA',
+  '76': 'Applying LoRA',
+  '77': 'Applying LoRA',
   '9': 'Encoding prompt',
   '11': 'Encoding image A',
   '24': 'Encoding image B',

server/utils/loras.ts

@@ -386,17 +386,6 @@ function linkRef(value: unknown): [string, number] | null {
     : null
 }
 
-function currentModelClipFeeds(graph: WorkflowGraph) {
-  const prompt = Object.values(graph).find(node => node.class_type === 'CLIPTextEncode' && linkRef(node.inputs.clip))
-  const modelNode = Object.values(graph).find(node => MODEL_FEED_CLASSES.has(node.class_type) && linkRef(node.inputs.model))
-  const model = linkRef(modelNode?.inputs.model)
-  const clip = linkRef(prompt?.inputs.clip)
-  return {
-    model: model && graph[model[0]] ? model : null,
-    clip: clip && graph[clip[0]] ? clip : null
-  }
-}
-
 export function assertImageV2LoraStack(stack: unknown, engine: 'flux' | 'krea') {
   const items = normalizeLoraStack(stack)
   if (engine !== 'krea') return items
@@ -407,16 +396,40 @@ export function assertImageV2LoraStack(stack: unknown, engine: 'flux' | 'krea')
   return items
 }
 
+function stripExistingLoraLoaders(graph: WorkflowGraph) {
+  const remaining = new Set(
+    Object.entries(graph)
+      .filter(([, node]) => node.class_type === LORA_LOADER)
+      .map(([id]) => id)
+  )
+  while (remaining.size) {
+    const id = [...remaining].find(candidate => (
+      ![...remaining].some(other => other !== candidate && linkSource(graph[other]?.inputs.model) === candidate)
+    )) || [...remaining][0]
+    bypassLoraNode(graph, id)
+    remaining.delete(id)
+  }
+}
+
+function modelConsumers(graph: WorkflowGraph) {
+  return Object.values(graph).filter(node => MODEL_FEED_CLASSES.has(node.class_type) && linkRef(node.inputs.model))
+}
+
+function clipConsumers(graph: WorkflowGraph) {
+  return Object.values(graph).filter(node => node.class_type === 'CLIPTextEncode' && linkRef(node.inputs.clip))
+}
+
 /**
- * Append a user LoRA chain after the current model/CLIP feed.
+ * Replace any template LoRA loaders with the posted chain, in list order.
- * Flux v2 already has Concept/Consistency loaders; this stacks after those.
+ * Strength 0 skips that card (off). Empty chain leaves UNET/CLIP wired straight to the sampler.
- * Empty stack is a no-op. Klein LoRAs on a Krea graph throw.
  */
 export function applyImageV2UserLoras(
   graph: WorkflowGraph,
   stack?: unknown,
   engine: 'flux' | 'krea' = 'flux'
 ) {
+  stripExistingLoraLoaders(graph)
+
   const xaigen = isXaigenStudio()
   const items = normalizeLoraStack(stack)
     .filter(item => xaigen || !isXaigenOnlyLora(item.name))
@@ -424,8 +437,7 @@ export function applyImageV2UserLoras(
       ...item,
       name: resolveUserLoraName(item.name, 'image')
     }))
-    .filter(item => item.name && (xaigen || !isXaigenOnlyLora(item.name)) && !alreadyHasLora(graph, item.name))
+    .filter(item => item.name && (xaigen || !isXaigenOnlyLora(item.name)) && (item.strengthModel !== 0 || item.strengthClip !== 0))
-  if (!items.length) return [] as LoraStackItem[]
 
   if (engine === 'krea') {
     const klein = items.find(item => isKleinImageLora(item.name))
@@ -434,48 +446,54 @@ export function applyImageV2UserLoras(
     }
   }
 
-  const feeds = currentModelClipFeeds(graph)
+  const guiders = modelConsumers(graph)
-  const modelFrom = feeds.model || (graph['4'] ? ['4', 0] as [string, number] : null)
+  const encodes = clipConsumers(graph)
-  const clipFrom = feeds.clip || (graph['5'] ? ['5', 0] as [string, number] : null)
+  const modelFrom = linkRef(guiders[0]?.inputs.model) || (graph['4'] ? ['4', 0] as [string, number] : null)
+  const clipFrom = linkRef(encodes[0]?.inputs.clip) || (graph['5'] ? ['5', 0] as [string, number] : null)
   if (!modelFrom || !clipFrom) {
     throw createError({ statusCode: 500, statusMessage: 'v2 graph has no model/CLIP feed for user LoRAs.' })
   }
+  if (!items.length) return [] as LoraStackItem[]
+  if (!guiders.length || !encodes.length) {
+    throw createError({
+      statusCode: 500,
+      statusMessage: 'Image v2 LoRA chain is not connected to the sampler. Refusing to run without those adapters.'
+    })
+  }
 
-  const firstId = 'user:lora:0'
+  let model: [string, number] = modelFrom
-  graph[firstId] = {
+  let clip: [string, number] = clipFrom
-    class_type: LORA_LOADER,
+  const ids: string[] = []
-    inputs: {
+  for (const [index, item] of items.entries()) {
-      lora_name: items[0].name,
+    const nodeId = String(70 + index)
-      strength_model: items[0].strengthModel,
+    if (graph[nodeId]) {
-      strength_clip: items[0].strengthClip,
+      throw createError({ statusCode: 500, statusMessage: `v2 graph already has node ${nodeId}. Refusing to overwrite it with a LoRA.` })
-      model: modelFrom,
-      clip: clipFrom
-    },
-    _meta: { title: items.length === 1 ? 'User LoRA' : 'User LoRA 1' }
-  }
-  for (const [id, node] of Object.entries(graph)) {
-    if (id === firstId) continue
-    for (const [key, value] of Object.entries(node.inputs)) {
-      if (!Array.isArray(value)) continue
-      if (value[0] === modelFrom[0] && value[1] === modelFrom[1]) node.inputs[key] = [firstId, 0]
-      else if (value[0] === clipFrom[0] && value[1] === clipFrom[1]) node.inputs[key] = [firstId, 1]
     }
-  }
+    graph[nodeId] = {
-  let sourceId = firstId
-  for (const [index, item] of items.slice(1).entries()) {
-    const nodeId = `user:lora:${index + 1}`
-    injectAfter(graph, sourceId, nodeId, {
       class_type: LORA_LOADER,
       inputs: {
         lora_name: item.name,
         strength_model: item.strengthModel,
         strength_clip: item.strengthClip,
-        model: [sourceId, 0],
+        model,
-        clip: [sourceId, 1]
+        clip
       },
-      _meta: { title: `User LoRA ${index + 2}` }
+      _meta: { title: items.length === 1 ? 'LoRA' : `LoRA ${index + 1}` }
+    }
+    ids.push(nodeId)
+    model = [nodeId, 0]
+    clip = [nodeId, 1]
+  }
+  for (const node of guiders) node.inputs.model = model
+  for (const node of encodes) node.inputs.clip = clip
+
+  const hooked = guiders.every(node => linkRef(node.inputs.model)?.[0] === ids[ids.length - 1])
+    && encodes.every(node => linkRef(node.inputs.clip)?.[0] === ids[ids.length - 1])
+  if (!hooked) {
+    throw createError({
+      statusCode: 500,
+      statusMessage: 'Image v2 LoRA chain is not connected to the sampler. Refusing to run without those adapters.'
     })
-    sourceId = nodeId
   }
   return items
 }