Add a Scale to 1 MP toggle on image edit before Klein.

Co-authored-by: Cursor <cursoragent@cursor.com>

server/utils/imageWorkflow.ts

@@ -1,5 +1,7 @@
 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'
@@ -16,6 +18,7 @@ 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>
@@ -31,6 +34,8 @@ export interface EditParams {
   filenamePrefix?: string
   loraName?: string
   loraStack?: LoraStackItem[]
+  scaleToTotalPixels?: boolean
+  scaleMegapixels?: number
 }
 
 export function buildEditWorkflow(params: EditParams) {
@@ -65,11 +70,63 @@ export function buildEditWorkflow(params: EditParams) {
     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
 }
@@ -85,6 +142,8 @@ function setInput(graph: WorkflowGraph, id: string, key: string, value: unknown)
 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',