Add Image v2 as a sibling of Image edit, with separate Klein graphs.

Keep the v1 Image tab on its dual-branch template. The host agent follows Desktop's live port so Coolify stays on :8198. Co-authored-by: Cursor <cursoragent@cursor.com>
2026-08-28 21:22:37 -05:00
parent 55f8cb1def
commit 21fe2412f9
18 changed files with 1678 additions and 67 deletions
@@ -0,0 +1,193 @@
+import editTemplate from '../assets/klein_v2_edit.json'
+import composeTemplate from '../assets/klein_v2_compose.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,
+  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 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 interface ImageV2BuildParams {
+  mode: ImageV2Mode
+  task: ImageV2Task
+  prompt: string
+  negative?: string
+  imageAName: string
+  imageBName?: string
+  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'
+  }
+}
+
+export function assertImageV2Graph(graph: WorkflowGraph, mode: ImageV2Mode, imageBName?: string) {
+  const loaders = loadImageNames(graph)
+  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 (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 graph = structuredClone(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, IMAGE_V2_CONSISTENCY_MODEL)
+  const consistencyClip = clampImageV2Strength(params.consistencyClip, IMAGE_V2_CONSISTENCY_CLIP)
+  const steps = clampImageSteps(params.steps, 24)
+  const cfg = clampImageCfg(params.cfg, 4)
+  const workflowFile = compose ? IMAGE_V2_COMPOSE_WORKFLOW : IMAGE_V2_EDIT_WORKFLOW
+
+  setInput(graph, LOAD_A, 'image', params.imageAName)
+  if (compose) setInput(graph, LOAD_B, 'image', params.imageBName || '')
+  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 || 'aigen-v2')
+  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)
+  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, params.imageBName)
+
+  const loaders = loadImageNames(graph)
+  console.log(JSON.stringify({
+    src: 'image-v2',
+    workflow: workflowFile,
+    mode: params.mode,
+    task: params.task,
+    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: { name: graph[CONSISTENCY]?.inputs.lora_name, model: consistencyModel, clip: consistencyClip }
+    },
+    steps,
+    cfg,
+    seed: params.seed,
+    megapixels: clampImageScaleMegapixels(params.megapixels ?? 1)
+  }))
+
+  return { graph, workflowFile, loaders, prompt }
+}
+
+export const IMAGE_V2_NODE_LABELS: Record<string, string> = {
+  '1': 'Loading image A',
+  '22': 'Loading image B',
+  '2': 'Scaling image A',
+  '23': 'Scaling image B',
+  '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',
+  '19': 'Sampling Klein v2',
+  '20': 'Decoding still',
+  '21': 'Saving still'
+}