From 1c6944e53441eb559b6abdcb365a1d160fc22b6c Mon Sep 17 00:00:00 2001 From: Towsty Date: Sat, 29 Aug 2026 21:06:14 -0500 Subject: [PATCH] Reconnect Klein Concept and Consistency to the sampler from the LoRA chain. Co-authored-by: Cursor --- server/utils/imageWorkflowV2.ts | 85 +++++++++++++++++++++++---------- 1 file changed, 61 insertions(+), 24 deletions(-) diff --git a/server/utils/imageWorkflowV2.ts b/server/utils/imageWorkflowV2.ts index f009d0c..a7d7106 100644 --- a/server/utils/imageWorkflowV2.ts +++ b/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,