diff --git a/server/utils/imageWorkflowV2.ts b/server/utils/imageWorkflowV2.ts index a7d7106..0b177f1 100644 --- a/server/utils/imageWorkflowV2.ts +++ b/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 { resolveComfyLoraName, loraIdentityKey, type LoraStackItem } from '~/utils/loras' +import { applyImageV2UserLoras } from '~/server/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: { - concept: concept.name ? { name: concept.name, model: concept.snofsModel, clip: concept.snofsClip } : { loaded: false }, - user: userLoras - }, + loras: { 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: { - concept: concept.name - ? { name: concept.name, model: concept.snofsModel, clip: concept.snofsClip } - : { loaded: false, model: 0, clip: 0 }, - user: userLoras - }, + loras: { 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: { - 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 - }, + loras: { user: userLoras, modelFeed }, steps, cfg, seed: params.seed, @@ -706,10 +557,14 @@ export const IMAGE_V2_NODE_LABELS: Record = { '6': 'Loading VAE', '7': 'Applying Concept LoRA', '8': 'Applying Consistency', - 'user:lora:0': 'Applying LoRA', - 'user:lora:1': 'Applying LoRA', - 'user:lora:2': 'Applying LoRA', - 'user:lora:3': 'Applying LoRA', + '70': 'Applying LoRA', + '71': 'Applying LoRA', + '72': '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', diff --git a/server/utils/loras.ts b/server/utils/loras.ts index 2bdcd31..3ee9944 100644 --- a/server/utils/loras.ts +++ 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. - * Flux v2 already has Concept/Consistency loaders; this stacks after those. - * Empty stack is a no-op. Klein LoRAs on a Krea graph throw. + * Replace any template LoRA loaders with the posted chain, in list order. + * Strength 0 skips that card (off). Empty chain leaves UNET/CLIP wired straight to the sampler. */ 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)) - if (!items.length) return [] as LoraStackItem[] + .filter(item => item.name && (xaigen || !isXaigenOnlyLora(item.name)) && (item.strengthModel !== 0 || item.strengthClip !== 0)) if (engine === 'krea') { const klein = items.find(item => isKleinImageLora(item.name)) @@ -434,48 +446,54 @@ export function applyImageV2UserLoras( } } - const feeds = currentModelClipFeeds(graph) - const modelFrom = feeds.model || (graph['4'] ? ['4', 0] as [string, number] : null) - const clipFrom = feeds.clip || (graph['5'] ? ['5', 0] as [string, number] : null) + const guiders = modelConsumers(graph) + const encodes = clipConsumers(graph) + 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' - graph[firstId] = { - class_type: LORA_LOADER, - inputs: { - lora_name: items[0].name, - strength_model: items[0].strengthModel, - strength_clip: items[0].strengthClip, - 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] + let model: [string, number] = modelFrom + let clip: [string, number] = clipFrom + const ids: string[] = [] + for (const [index, item] of items.entries()) { + const nodeId = String(70 + index) + if (graph[nodeId]) { + throw createError({ statusCode: 500, statusMessage: `v2 graph already has node ${nodeId}. Refusing to overwrite it with a LoRA.` }) } - } - let sourceId = firstId - for (const [index, item] of items.slice(1).entries()) { - const nodeId = `user:lora:${index + 1}` - injectAfter(graph, sourceId, nodeId, { + graph[nodeId] = { class_type: LORA_LOADER, inputs: { lora_name: item.name, strength_model: item.strengthModel, strength_clip: item.strengthClip, - model: [sourceId, 0], - clip: [sourceId, 1] + model, + 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 }