Files
aigen/server/utils/imageWorkflow.ts
T

97 lines
3.0 KiB
TypeScript

import template from '../assets/workflow_flux2_klein_edit.json'
const IMAGE1 = '76'
const IMAGE2 = '81'
const SINGLE_SAVE = '9'
const DUAL_SAVE = '94'
const SINGLE_PROMPT = '75:74'
const SINGLE_NEGATIVE = '75:67'
const SINGLE_NOISE = '75:73'
const SINGLE_SCHEDULER = '75:62'
const SINGLE_CFG = '75:63'
const DUAL_PROMPT = '92:113'
const DUAL_NEGATIVE = '92:87'
const DUAL_NOISE = '92:105'
const DUAL_SCHEDULER = '92:115'
const DUAL_CFG = '92:114'
type WorkflowNode = { class_type: string; inputs: Record<string, unknown>; _meta?: { title?: string } }
type WorkflowGraph = Record<string, WorkflowNode>
export interface EditParams {
prompt: string
negative?: string
imageName: string
referenceImageName?: string
steps: number
seed: number
cfg: number
filenamePrefix?: string
}
export function buildEditWorkflow(params: EditParams) {
const graph = structuredClone(template) as WorkflowGraph
const dual = Boolean(params.referenceImageName)
const prompt = String(params.prompt || '').trim()
const negative = String(params.negative || '')
const prefix = params.filenamePrefix || 'aigen-edit'
setImage(graph, IMAGE1, params.imageName)
setText(graph, SINGLE_PROMPT, prompt)
setText(graph, DUAL_PROMPT, prompt)
setText(graph, SINGLE_NEGATIVE, negative)
setText(graph, DUAL_NEGATIVE, negative)
setInput(graph, SINGLE_NOISE, 'noise_seed', params.seed)
setInput(graph, DUAL_NOISE, 'noise_seed', params.seed)
setInput(graph, SINGLE_SCHEDULER, 'steps', params.steps)
setInput(graph, DUAL_SCHEDULER, 'steps', params.steps)
setInput(graph, SINGLE_CFG, 'cfg', params.cfg)
setInput(graph, DUAL_CFG, 'cfg', params.cfg)
setInput(graph, SINGLE_SAVE, 'filename_prefix', prefix)
setInput(graph, DUAL_SAVE, 'filename_prefix', prefix)
if (dual) {
setImage(graph, IMAGE2, params.referenceImageName || '')
}
for (const id of Object.keys(graph)) {
if (id === IMAGE1) continue
const dualNode = id === IMAGE2 || id === DUAL_SAVE || id.startsWith('92:')
const singleNode = id === SINGLE_SAVE || id.startsWith('75:')
if (dual ? singleNode : dualNode) delete graph[id]
}
return graph
}
function setImage(graph: WorkflowGraph, id: string, name: string) {
if (graph[id]) graph[id].inputs.image = name
}
function setText(graph: WorkflowGraph, id: string, text: string) {
if (graph[id]) graph[id].inputs.text = text
}
function setInput(graph: WorkflowGraph, id: string, key: string, value: unknown) {
if (graph[id]) graph[id].inputs[key] = value
}
export const EDIT_NODE_LABELS: Record<string, string> = {
'76': 'Loading image 1',
'81': 'Loading image 2',
'75:70': 'Loading Flux.2 Klein',
'75:71': 'Loading CLIP',
'75:72': 'Loading VAE',
'75:74': 'Encoding prompt',
'75:64': 'Sampling edit',
'75:65': 'Decoding still',
'9': 'Saving still',
'92:106': 'Loading Flux.2 Klein',
'92:111': 'Loading CLIP',
'92:107': 'Loading VAE',
'92:113': 'Encoding prompt',
'92:103': 'Sampling two-image edit',
'92:104': 'Decoding still',
'94': 'Saving still'
}