Reconnect Klein Concept and Consistency to the sampler from the LoRA chain.

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

server/utils/imageWorkflowV2.ts

@@ -13,9 +13,6 @@ import { resolveComfyLoraName, loraIdentityKey, type LoraStackItem } from '~/uti
 import { clampImageCfg, clampImageScaleMegapixels, clampImageSteps } from '~/utils/generationPresets'
 import { isXaigenStudio } from '~/utils/videoModels'
 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,
@@ -25,9 +22,6 @@ import {
   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,
   IMAGE_V2_DENOISE_DEFAULT,
   IMAGE_V2_GENERATE_HEIGHT,
   IMAGE_V2_GENERATE_WIDTH,
@@ -36,6 +30,8 @@ import {
   clampImageV2Strength,
   composeImageV2Prompt,
   imageV2WorkflowFile,
+  isImageV2ConceptLora,
+  isImageV2ConsistencyLora,
   mergeLegacyImageV2Loras,
   parseImageV2Engine,
   type ImageV2Engine,
@@ -438,6 +434,57 @@ 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({
@@ -584,17 +631,6 @@ export function buildImageV2Workflow(params: ImageV2BuildParams) {
   ) as WorkflowGraph
   const prompt = composeImageV2Prompt(params.mode, params.task, params.prompt)
   const negative = String(params.negative || '')
-  const conceptAllowed = isXaigenStudio()
-  const snofsModel = conceptAllowed ? clampImageV2Strength(params.snofsModel, IMAGE_V2_SNOFS_MODEL) : 0
-  const snofsClip = conceptAllowed ? clampImageV2Strength(params.snofsClip, IMAGE_V2_SNOFS_CLIP) : 0
-  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
@@ -621,9 +657,7 @@ 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)
 
-  if (graph[CONSISTENCY]) bypassLoraToBaseLoaders(graph, CONSISTENCY)
-  if (graph[SNOFS]) bypassLoraToBaseLoaders(graph, SNOFS)
-  const userLoras = applyImageV2UserLoras(graph, imageV2UserStack(params, 'flux'), 'flux')
+  const fluxLoras = applyFluxTemplateLoras(graph, params)
 
   assertImageV2Graph(graph, params.mode, generate || refine ? undefined : params.imageBName, engine)
 
@@ -641,11 +675,14 @@ export function buildImageV2Workflow(params: ImageV2BuildParams) {
     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: graph[CONSISTENCY]
-        ? { name: graph[CONSISTENCY]?.inputs.lora_name, model: consistencyModel, clip: consistencyClip }
+      snofs: fluxLoras.concept
+        ? { name: graph[SNOFS]?.inputs.lora_name, model: fluxLoras.concept.strengthModel, clip: fluxLoras.concept.strengthClip }
         : { loaded: false, model: 0, clip: 0 },
-      user: userLoras
+      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,