aigen/server/utils/imageWorkflow.ts

import template from '../assets/workflow_flux2_klein_edit.json'
import { IMAGE_SCALE_TO_TOTAL_PIXELS } from '~/server/utils/comfy'
import { applyOptionalLoraLoaders } from '~/server/utils/loras'
import { clampImageScaleMegapixels } from '~/utils/generationPresets'
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'
const LOAD_IMAGE = 'LoadImage'

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[]
  scaleToTotalPixels?: boolean
  scaleMegapixels?: number
}

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 || '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]
  }

  applyScaleToTotalPixels(graph, Boolean(params.scaleToTotalPixels), params.scaleMegapixels)

  // User LoRAs only. Empty stack bypasses the optional LoraLoader (no snofs default).
  applyOptionalLoraLoaders(graph, params.loraStack || params.loraName, 'image')
  return graph
}

function linkSource(value: unknown): string | null {
  return Array.isArray(value) && typeof value[0] === 'string' ? value[0] : null
}

function stripScaleNodes(graph: WorkflowGraph) {
  for (const [id, node] of Object.entries(graph)) {
    if (node.class_type !== IMAGE_SCALE_TO_TOTAL_PIXELS) continue
    const src = Array.isArray(node.inputs.image) ? node.inputs.image : null
    for (const [otherId, other] of Object.entries(graph)) {
      if (otherId === id) continue
      for (const [key, value] of Object.entries(other.inputs)) {
        if (linkSource(value) !== id || !Array.isArray(value)) continue
        other.inputs[key] = src ? [src[0], value[1] === 0 ? src[1] : value[1]] : value
      }
    }
    delete graph[id]
  }
}

function insertScaleNodes(graph: WorkflowGraph, megapixels: number) {
  const loaders = Object.entries(graph).filter(([, node]) => node.class_type === LOAD_IMAGE)
  for (const [id] of loaders) {
    const scaleId = `scale:${id}`
    if (graph[scaleId]) continue
    graph[scaleId] = {
      class_type: IMAGE_SCALE_TO_TOTAL_PIXELS,
      inputs: {
        upscale_method: 'lanczos',
        megapixels,
        resolution_steps: 1,
        image: [id, 0]
      },
      _meta: { title: 'Scale Image to Total Pixels' }
    }
    for (const [otherId, other] of Object.entries(graph)) {
      if (otherId === scaleId) continue
      for (const [key, value] of Object.entries(other.inputs)) {
        if (!Array.isArray(value) || value[0] !== id || value[1] !== 0) continue
        other.inputs[key] = [scaleId, 0]
      }
    }
  }
}

function applyScaleToTotalPixels(graph: WorkflowGraph, enabled: boolean, megapixels?: number) {
  stripScaleNodes(graph)
  if (!enabled) return
  insertScaleNodes(graph, clampImageScaleMegapixels(megapixels))
}

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',
  'scale:76': 'Scaling image 1',
  'scale:81': 'Scaling 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'
}