102 lines
3.2 KiB
TypeScript
102 lines
3.2 KiB
TypeScript
import template from '../assets/workflow_flux2_klein_edit.json'
|
|
import { applyUserLoraToGraph } from '~/server/utils/loras'
|
|
import type { LoraStackItem } from '~/utils/loras'
|
|
|
|
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
|
|
loraName?: string
|
|
loraStack?: LoraStackItem[]
|
|
}
|
|
|
|
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]
|
|
}
|
|
|
|
applyUserLoraToGraph(graph, params.loraStack || params.loraName)
|
|
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'
|
|
}
|