aigen/server/utils/imageWorkflowV2.ts

import editTemplate from '../assets/klein_v2_edit.json'
import composeTemplate from '../assets/klein_v2_compose.json'
import refineTemplate from '../assets/klein_v2_refine.json'
import generateTemplate from '../assets/klein_v2_generate.json'
import { IMAGE_SCALE_TO_TOTAL_PIXELS } from '~/server/utils/comfy'
import { cachedComfyLoraNames } from '~/server/utils/loras'
import { resolveComfyLoraName, loraIdentityKey } from '~/utils/loras'
import { clampImageCfg, clampImageScaleMegapixels, clampImageSteps } from '~/utils/generationPresets'
import {
  IMAGE_V2_CONSISTENCY_CLIP,
  IMAGE_V2_CONSISTENCY_LORA,
  IMAGE_V2_CONSISTENCY_MODEL,
  IMAGE_V2_SNOFS_CLIP,
  IMAGE_V2_SNOFS_LORA,
  IMAGE_V2_SNOFS_MODEL,
  IMAGE_V2_DENOISE_DEFAULT,
  IMAGE_V2_GENERATE_HEIGHT,
  IMAGE_V2_GENERATE_WIDTH,
  clampImageV2Denoise,
  clampImageV2Size,
  clampImageV2Strength,
  composeImageV2Prompt,
  type ImageV2Mode,
  type ImageV2Task
} from '~/utils/imageV2'

type WorkflowNode = { class_type: string; inputs: Record<string, unknown>; _meta?: { title?: string } }
type WorkflowGraph = Record<string, WorkflowNode>

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 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'
export const IMAGE_V2_GENERATE_WORKFLOW = 'klein_v2_generate.json'

export interface ImageV2BuildParams {
  mode: ImageV2Mode
  task: ImageV2Task
  prompt: string
  negative?: string
  imageAName?: string
  imageBName?: string
  maskName?: string
  strength?: number
  width?: number
  height?: number
  turbo?: boolean
  snofsModel?: number
  snofsClip?: number
  consistencyModel?: number
  consistencyClip?: number
  steps: number
  cfg: number
  seed: number
  megapixels?: number
  filenamePrefix?: string
}

function setInput(graph: WorkflowGraph, id: string, key: string, value: unknown) {
  if (graph[id]) graph[id].inputs[key] = value
}

function loadImageNames(graph: WorkflowGraph) {
  return Object.entries(graph)
    .filter(([, node]) => node.class_type === 'LoadImage')
    .map(([id, node]) => ({
      id,
      title: String(node._meta?.title || id),
      image: String(node.inputs.image || '')
    }))
}

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)) {
    if (node.class_type !== IMAGE_SCALE_TO_TOTAL_PIXELS) continue
    node.inputs.megapixels = mp
    node.inputs.upscale_method = 'lanczos'
  }
}

function graphHasMaskInput(graph: WorkflowGraph) {
  return Object.values(graph).some((node) => {
    const mask = node.inputs?.mask
    return mask !== undefined && mask !== null && mask !== ''
  })
}

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

export function assertImageV2Graph(graph: WorkflowGraph, mode: ImageV2Mode, imageBName?: string) {
  const loaders = loadImageNames(graph)
  if (mode === 'generate') {
    if (loaders.length) {
      throw createError({
        statusCode: 500,
        statusMessage: 'Generate graph has a required LoadImage. Refusing to run an edit fallback.'
      })
    }
    const latent = Object.values(graph).find(node => node.class_type === 'EmptyFlux2LatentImage')
    if (!latent) {
      throw createError({ statusCode: 500, statusMessage: 'Generate graph is missing EmptyFlux2LatentImage.' })
    }
  }
  if (mode === 'refine') {
    const mask = graph[LOAD_MASK]
    if (!mask || mask.class_type !== 'LoadImage' || !String(mask.inputs.image || '').trim()) {
      throw createError({
        statusCode: 500,
        statusMessage: 'Refine job is missing the mask image. Refusing to run.'
      })
    }
    if (!graphHasMaskInput(graph)) {
      throw createError({
        statusCode: 500,
        statusMessage: 'Refine graph has no mask input. Refusing to run.'
      })
    }
    if (!('denoise' in (graph[SCHEDULER]?.inputs || {}))) {
      throw createError({
        statusCode: 500,
        statusMessage: 'Refine graph has no denoise on the sampler. Refusing to run.'
      })
    }
  }
  if (mode === 'compose') {
    if (loaders.length < 2) {
      throw createError({
        statusCode: 500,
        statusMessage: 'Compose graph has no second image input. Refusing to run a one-image fallback.'
      })
    }
    const b = graph[LOAD_B]
    if (!b || b.class_type !== 'LoadImage' || !String(b.inputs.image || '').trim()) {
      throw createError({
        statusCode: 500,
        statusMessage: 'Compose job is missing Load Image B. Refusing to run.'
      })
    }
  }
  if (mode !== 'refine' && mode !== 'generate' && imageBName && loaders.length < 2) {
    throw createError({
      statusCode: 500,
      statusMessage: 'image_b was sent but the executed graph has no second image input.'
    })
  }
  const promptNode = graph[PROMPT]
  if (!promptNode || promptNode.class_type !== 'CLIPTextEncode') {
    throw createError({ statusCode: 500, statusMessage: 'v2 graph is missing the positive CLIPTextEncode node.' })
  }
  if (Array.isArray(promptNode.inputs.text)) {
    throw createError({ statusCode: 500, statusMessage: 'v2 prompt is a subgraph link. Refusing to run with a leftover widget prompt.' })
  }
  if (!String(promptNode.inputs.text || '').trim()) {
    throw createError({ statusCode: 400, statusMessage: 'v2 prompt was not patched onto the graph.' })
  }
}

export function buildImageV2Workflow(params: ImageV2BuildParams) {
  const compose = params.mode === 'compose'
  const refine = params.mode === 'refine'
  const generate = params.mode === 'generate'
  if (refine && !String(params.maskName || '').trim()) {
    throw createError({ statusCode: 400, statusMessage: 'Refine requires a mask. Refusing to fall back to Edit.' })
  }
  const graph = structuredClone(
    generate ? generateTemplate : refine ? refineTemplate : compose ? composeTemplate : editTemplate
  ) as WorkflowGraph
  const prompt = composeImageV2Prompt(params.mode, params.task, params.prompt)
  const negative = String(params.negative || '')
  const snofsModel = clampImageV2Strength(params.snofsModel, IMAGE_V2_SNOFS_MODEL)
  const snofsClip = clampImageV2Strength(params.snofsClip, IMAGE_V2_SNOFS_CLIP)
  const consistencyModel = clampImageV2Strength(
    params.consistencyModel,
    generate ? 0 : IMAGE_V2_CONSISTENCY_MODEL
  )
  const consistencyClip = clampImageV2Strength(
    params.consistencyClip,
    generate ? 0 : IMAGE_V2_CONSISTENCY_CLIP
  )
  const steps = clampImageSteps(params.steps, 24)
  const cfg = clampImageCfg(params.cfg, 4)
  const strength = refine ? clampImageV2Denoise(params.strength, IMAGE_V2_DENOISE_DEFAULT) : undefined
  const size = generate ? clampImageV2Size(params.width, params.height) : null
  const workflowFile = generate
    ? IMAGE_V2_GENERATE_WORKFLOW
    : refine ? IMAGE_V2_REFINE_WORKFLOW : compose ? IMAGE_V2_COMPOSE_WORKFLOW : IMAGE_V2_EDIT_WORKFLOW

  if (!generate) setInput(graph, LOAD_A, 'image', params.imageAName || '')
  if (compose) setInput(graph, LOAD_B, 'image', params.imageBName || '')
  if (refine) {
    setInput(graph, LOAD_MASK, 'image', params.maskName || '')
    setInput(graph, SCHEDULER_DENOISE, 'denoise', strength)
  }
  if (generate && size) {
    setInput(graph, '14', 'width', size.width)
    setInput(graph, '14', 'height', size.height)
    setInput(graph, SCHEDULER, 'width', size.width)
    setInput(graph, SCHEDULER, 'height', size.height)
  }
  setInput(graph, PROMPT, 'text', prompt)
  setInput(graph, NEGATIVE, 'text', negative)
  setInput(graph, NOISE, 'noise_seed', params.seed)
  setInput(graph, SCHEDULER, 'steps', steps)
  setInput(graph, CFG, 'cfg', cfg)
  setInput(graph, SAVE, 'filename_prefix', params.filenamePrefix || (generate ? 'aigen-v2-generate' : 'aigen-v2'))
  if (!generate) patchScaleMegapixels(graph, params.megapixels ?? 1)

  setInput(graph, SNOFS, 'lora_name', resolveRequiredLora(IMAGE_V2_SNOFS_LORA, 'SNOFS'))
  setInput(graph, SNOFS, 'strength_model', snofsModel)
  setInput(graph, SNOFS, 'strength_clip', snofsClip)
  if (generate && consistencyModel <= 0 && consistencyClip <= 0) {
    bypassLoraNode(graph, CONSISTENCY, SNOFS, SNOFS)
  } else {
    setInput(graph, CONSISTENCY, 'lora_name', resolveRequiredLora(IMAGE_V2_CONSISTENCY_LORA, 'Consistency'))
    setInput(graph, CONSISTENCY, 'strength_model', consistencyModel)
    setInput(graph, CONSISTENCY, 'strength_clip', consistencyClip)
  }

  assertImageV2Graph(graph, params.mode, generate || refine ? undefined : params.imageBName)

  const loaders = loadImageNames(graph)
  console.log(JSON.stringify({
    src: 'image-v2',
    workflow: workflowFile,
    mode: params.mode,
    task: params.task,
    canvas: generate ? undefined : { id: LOAD_A, file: graph[LOAD_A]?.inputs.image },
    mask: refine ? { id: LOAD_MASK, file: graph[LOAD_MASK]?.inputs.image } : undefined,
    size: generate ? { width: size?.width ?? IMAGE_V2_GENERATE_WIDTH, height: size?.height ?? IMAGE_V2_GENERATE_HEIGHT } : undefined,
    strength,
    turbo: params.turbo === true,
    loadImage: Object.fromEntries(loaders.map(item => [item.id, { title: item.title, file: item.image }])),
    loras: {
      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 }
    },
    steps,
    cfg,
    seed: params.seed,
    megapixels: generate ? undefined : clampImageScaleMegapixels(params.megapixels ?? 1)
  }))

  return { graph, workflowFile, loaders, prompt, strength }
}

export const IMAGE_V2_NODE_LABELS: Record<string, string> = {
  '1': 'Loading image A',
  '22': 'Loading image B',
  '30': 'Loading mask',
  '2': 'Scaling image A',
  '23': 'Scaling image B',
  '31': 'Resizing mask',
  '4': 'Loading Flux.2 Klein 9B',
  '5': 'Loading CLIP',
  '6': 'Loading VAE',
  '7': 'Applying SNOFS',
  '8': 'Applying Consistency',
  '9': 'Encoding prompt',
  '11': 'Encoding image A',
  '24': 'Encoding image B',
  '33': 'Applying mask',
  '14': 'Building empty Klein latent',
  '19': 'Sampling Klein v2',
  '20': 'Decoding still',
  '21': 'Saving still'
}