diff --git a/pages/index.vue b/pages/index.vue index d284d56..4a23728 100644 --- a/pages/index.vue +++ b/pages/index.vue @@ -1975,7 +1975,6 @@ import { filterLorasForImageEngine, isKleinImageLora, kleinLoraOnKreaMessage, - loraDisplayName, loraStacksEqual, normalizeLoraStack, persistLoraFields, @@ -2442,7 +2441,7 @@ const videoLoraOptions = computed(() => mergeLoraOptions(videoLoras.value, [ const imageLoraOptions = computed(() => mergeLoraOptions(imageLoras.value, imageLoraStack.value.map(item => item.name))) const imageV2LoraOptions = computed(() => mergeLoraOptions( filterLorasForImageEngine(imageLoras.value, v2Engine.value), - imageV2LoraStack.value.map(item => item.name) + filterLorasForImageEngine(imageV2LoraStack.value.map(item => item.name), v2Engine.value) )) const useIdentityRefs = ref(false) const identityRefs = ref<(File | null)[]>([null, null, null, null]) diff --git a/server/api/v2/generate.post.ts b/server/api/v2/generate.post.ts index dacf4c1..19be2c9 100644 --- a/server/api/v2/generate.post.ts +++ b/server/api/v2/generate.post.ts @@ -27,6 +27,7 @@ import { type ImageV2Task } from '~/utils/imageV2' import { isXaigenStudio } from '~/utils/videoModels' +import { listStudioLoras, parsePostedLoraStack, persistLoraFields, assertImageV2LoraStack } from '~/server/utils/loras' import { getStill, rememberInputStill, stillPath } from '~/server/utils/library' import { existsSync, readFileSync } from 'node:fs' @@ -222,6 +223,8 @@ export default defineEventHandler(async (event) => { const conceptAllowed = isXaigenStudio() const generateSnofsFallback = !conceptAllowed || engine === 'krea' ? 0 : IMAGE_V2_SNOFS_MODEL const generateSnofsClipFallback = !conceptAllowed || engine === 'krea' ? 0 : IMAGE_V2_SNOFS_CLIP + await listStudioLoras() + const loraStack = assertImageV2LoraStack(parsePostedLoraStack(fields.loraStack || fields.lora, 'image'), engine) const still = imageA ? await rememberInputStill({ @@ -298,7 +301,8 @@ export default defineEventHandler(async (event) => { consistencyClip: generateConsistencyClip, maskStillId: savedMask?.id, maskStillFilename: savedMask?.filename, - refineStrength + refineStrength, + ...persistLoraFields(loraStack) } }) await kickStudioQueue() diff --git a/server/utils/imageChainV2.ts b/server/utils/imageChainV2.ts index f705c64..9be3dfe 100644 --- a/server/utils/imageChainV2.ts +++ b/server/utils/imageChainV2.ts @@ -3,7 +3,7 @@ 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, IMAGE_V2_NODE_LABELS, resolveKreaGenerateAssets } from '~/server/utils/imageWorkflowV2' -import { ensureComfyLoraNames } from '~/server/utils/loras' +import { ensureComfyLoraNames, persistLoraFields } from '~/server/utils/loras' import { imageDimensions } from '~/server/utils/resolution' import { emitChainJob } from '~/server/utils/watch' import { saveStill } from '~/server/utils/library' @@ -33,6 +33,7 @@ export type EditV2RunParams = { turbo?: boolean sourceStillId?: string referenceStillId?: string + loraStack?: import('~/utils/loras').LoraStackItem[] } export async function runEditV2(job: Job, params: EditV2RunParams) { @@ -132,7 +133,8 @@ export async function runEditV2(job: Job, params: EditV2RunParams) { kreaUnetName: krea?.unet, kreaClipName: krea?.clip, kreaVaeName: krea?.vae, - kreaConceptLora: krea?.conceptLora + kreaConceptLora: krea?.conceptLora, + loraStack: params.loraStack }) const queued = await queuePrompt(built.graph, job.clientId) job.promptId = queued.prompt_id @@ -199,7 +201,8 @@ export async function runEditV2(job: Job, params: EditV2RunParams) { height: params.mode === 'generate' ? params.height : undefined, seed: params.seed, sourceStillId: params.sourceStillId || library.stillId, - referenceStillId: params.referenceStillId + referenceStillId: params.referenceStillId, + ...persistLoraFields(params.loraStack) } }) job.stillId = still?.id diff --git a/server/utils/imageWorkflowV2.ts b/server/utils/imageWorkflowV2.ts index 351ac30..b1c52b8 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 { cachedComfyLoraNames } from '~/server/utils/loras' -import { resolveComfyLoraName, loraIdentityKey } from '~/utils/loras' +import { applyImageV2UserLoras, cachedComfyLoraNames } from '~/server/utils/loras' +import { resolveComfyLoraName, loraIdentityKey, type LoraStackItem } from '~/utils/loras' import { clampImageCfg, clampImageScaleMegapixels, clampImageSteps } from '~/utils/generationPresets' import { isXaigenStudio } from '~/utils/videoModels' import { @@ -97,6 +97,7 @@ export interface ImageV2BuildParams { kreaClipName?: string kreaVaeName?: string kreaConceptLora?: string + loraStack?: LoraStackItem[] } function setInput(graph: WorkflowGraph, id: string, key: string, value: unknown) { @@ -462,6 +463,7 @@ function buildKreaItiWorkflow(params: ImageV2BuildParams) { 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, params.loraStack, 'krea') assertImageV2Graph(graph, params.mode, compose ? params.imageBName : undefined, 'krea') const loaders = loadImageNames(graph) @@ -475,7 +477,10 @@ 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.name ? { concept: { name: concept.name, model: concept.snofsModel, clip: concept.snofsClip } } : { loaded: false }, + loras: { + concept: concept.name ? { name: concept.name, model: concept.snofsModel, clip: concept.snofsClip } : { loaded: false }, + user: userLoras + }, steps, cfg, seed: params.seed, @@ -509,6 +514,7 @@ function buildKreaGenerateWorkflow(params: ImageV2BuildParams) { setInput(graph, '15', 'scheduler', 'simple') setInput(graph, SAVE, 'filename_prefix', params.filenamePrefix || 'aigen-v2-krea') const concept = applyKreaConceptIfNeeded(graph, params) + const userLoras = applyImageV2UserLoras(graph, params.loraStack, 'krea') assertImageV2Graph(graph, 'generate', undefined, 'krea') const loaders = loadImageNames(graph) @@ -527,7 +533,8 @@ function buildKreaGenerateWorkflow(params: ImageV2BuildParams) { loras: { concept: concept.name ? { name: concept.name, model: concept.snofsModel, clip: concept.snofsClip } - : { loaded: false, model: 0, clip: 0 } + : { loaded: false, model: 0, clip: 0 }, + user: userLoras }, steps, cfg, @@ -604,6 +611,7 @@ export function buildImageV2Workflow(params: ImageV2BuildParams) { setInput(graph, CONSISTENCY, 'strength_model', consistencyModel) setInput(graph, CONSISTENCY, 'strength_clip', consistencyClip) } + const userLoras = applyImageV2UserLoras(graph, params.loraStack, 'flux') assertImageV2Graph(graph, params.mode, generate || refine ? undefined : params.imageBName, engine) @@ -624,7 +632,8 @@ export function buildImageV2Workflow(params: ImageV2BuildParams) { snofs: { name: graph[SNOFS]?.inputs.lora_name, model: snofsModel, clip: snofsClip }, consistency: graph[CONSISTENCY] ? { name: graph[CONSISTENCY]?.inputs.lora_name, model: consistencyModel, clip: consistencyClip } - : { loaded: false, model: 0, clip: 0 } + : { loaded: false, model: 0, clip: 0 }, + user: userLoras }, steps, cfg, @@ -648,6 +657,10 @@ 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', '9': 'Encoding prompt', '11': 'Encoding image A', '24': 'Encoding image B', diff --git a/server/utils/loras.ts b/server/utils/loras.ts index a3792ec..2bdcd31 100644 --- a/server/utils/loras.ts +++ b/server/utils/loras.ts @@ -4,7 +4,9 @@ import { XAIGEN_LORA_MESSAGE, MINIMAX_TURBO_LORA, filterLorasForStudio, + isKleinImageLora, isXaigenOnlyLora, + kleinLoraOnKreaMessage, loraIdentityKey, normalizeLoraName, normalizeLoraStack, @@ -378,6 +380,106 @@ export function applyOptionalLoraLoaders(graph: WorkflowGraph, stack?: unknown, } } +function linkRef(value: unknown): [string, number] | null { + return Array.isArray(value) && typeof value[0] === 'string' && Number.isFinite(Number(value[1])) + ? [value[0], Number(value[1])] + : 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 + const klein = items.find(item => isKleinImageLora(item.name)) + if (klein) { + throw createError({ statusCode: 400, statusMessage: kleinLoraOnKreaMessage(klein.name) }) + } + return items +} + +/** + * 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. + */ +export function applyImageV2UserLoras( + graph: WorkflowGraph, + stack?: unknown, + engine: 'flux' | 'krea' = 'flux' +) { + const xaigen = isXaigenStudio() + const items = normalizeLoraStack(stack) + .filter(item => xaigen || !isXaigenOnlyLora(item.name)) + .map(item => ({ + ...item, + name: resolveUserLoraName(item.name, 'image') + })) + .filter(item => item.name && (xaigen || !isXaigenOnlyLora(item.name)) && !alreadyHasLora(graph, item.name)) + if (!items.length) return [] as LoraStackItem[] + + if (engine === 'krea') { + const klein = items.find(item => isKleinImageLora(item.name)) + if (klein) { + throw createError({ statusCode: 400, statusMessage: kleinLoraOnKreaMessage(klein.name) }) + } + } + + 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) + if (!modelFrom || !clipFrom) { + throw createError({ statusCode: 500, statusMessage: 'v2 graph has no model/CLIP feed for user LoRAs.' }) + } + + 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 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], + clip: [sourceId, 1] + }, + _meta: { title: `User LoRA ${index + 2}` } + }) + sourceId = nodeId + } + return items +} + export function applyUserLoraToGraph(graph: WorkflowGraph, stack?: unknown) { const xaigen = isXaigenStudio() const items = normalizeLoraStack(stack) diff --git a/server/utils/studioQueue.ts b/server/utils/studioQueue.ts index ddd4702..a488274 100644 --- a/server/utils/studioQueue.ts +++ b/server/utils/studioQueue.ts @@ -748,7 +748,8 @@ async function startStudioEditJob(item: StudioJob) { familyId: item.familyId, chainIndex: 0, chainStep: 1, - chainTotal: 1 + chainTotal: 1, + ...persistLoraFields(payload.loraStack || payload.loraName) } }) await markStudioLive(item.ownerKey, item.id, live.id) @@ -774,7 +775,8 @@ async function startStudioEditJob(item: StudioJob) { height: payload.height, turbo: payload.turbo === true, sourceStillId: payload.stillId, - referenceStillId: payload.referenceStillId + referenceStillId: payload.referenceStillId, + loraStack: payload.loraStack }).catch((error) => { const message = error instanceof Error ? error.message : String(error) if (live && live.status !== 'error' && live.status !== 'cancelled' && live.status !== 'deferred') { diff --git a/utils/libraryGeneration.ts b/utils/libraryGeneration.ts index 94429db..f6b7f9a 100644 --- a/utils/libraryGeneration.ts +++ b/utils/libraryGeneration.ts @@ -160,6 +160,8 @@ export function stillSettingsRows(still: { width?: number; height?: number; sett if (settings.mode !== 'generate' && (settings.consistencyModel != null || settings.consistencyClip != null)) { rows.push({ label: 'Consistency', value: formatPair(settings.consistencyModel, settings.consistencyClip) }) } + const loras = formatLoras(settings.loraStack, settings.loraName) + if (loras) rows.push({ label: 'LoRA', value: loras }) if (settings.mode === 'refine' && settings.strength != null) { rows.push({ label: 'Strength', value: formatNum(settings.strength) }) } diff --git a/utils/loras.ts b/utils/loras.ts index 507aecd..c33126b 100644 --- a/utils/loras.ts +++ b/utils/loras.ts @@ -213,3 +213,24 @@ export function filterLorasForStudio(names: string[], xaigenStudio: boolean) { export function filterLoraStackForStudio(stack: unknown, xaigenStudio: boolean) { return normalizeLoraStack(stack).filter(item => xaigenStudio || !isXaigenOnlyLora(item.name)) } + +/** Klein / Flux.2 Klein adapters. Must not load on a Krea graph. */ +export function isKleinImageLora(name: string) { + const id = loraIdentityKey(name) + return /flux-?2-klein|klein_snofs|flux2-klein-9b-consistency|(^|[^a-z0-9])klein([^a-z0-9]|$)/i.test(id) +} + +export function filterLorasForImageEngine(names: string[], engine: 'flux' | 'krea') { + if (engine !== 'krea') return names + return names.filter(name => !isKleinImageLora(name)) +} + +export function filterLoraStackForImageEngine(stack: unknown, engine: 'flux' | 'krea') { + const items = normalizeLoraStack(stack) + if (engine !== 'krea') return items + return items.filter(item => !isKleinImageLora(item.name)) +} + +export function kleinLoraOnKreaMessage(name: string) { + return `${loraDisplayName(name)} is a Flux LoRA and cannot be used with Krea.` +}