towsty/aigen
1
0
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases Wiki Activity

Let Image v2 pick a LoRA or LoRA chain.

Browse Source 
Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Towsty
2026-08-29 17:28:32 -05:00
co-authored by Cursor
parent 852b07bb97
commit cdc69d2bd6
8 changed files with 159 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
server/utils/imageChainV2.ts
+6 -3
View File
@@ -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
server/utils/imageWorkflowV2.ts
+18 -5
View File
@@ -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<string, string> = {
'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',
server/utils/loras.ts
+102
View File
@@ -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)
server/utils/studioQueue.ts
+4 -2
View File
@@ -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') {