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

Wire every Image v2 LoRA, including Concept, through one sampler chain.

Browse Source 
Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Towsty
2026-08-29 21:11:50 -05:00
co-authored by Cursor
parent 1c6944e534
commit 61d711f989
2 changed files with 80 additions and 207 deletions
Download Patch File Download Diff File Expand all files Collapse all files
server/utils/imageWorkflowV2.ts
+16 -161
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 { 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<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',
'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',