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

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

server/api/edit.post.ts

@@ -1,8 +1,8 @@
 import { addStudioJob, kickStudioQueue, listStudioJobs } from '~/server/utils/studioQueue'
 import { parsePostedLoraStack, listStudioLoras, persistLoraFields } from '~/server/utils/loras'
-import { comfyConfigured } from '~/server/utils/comfy'
+import { assertImageScaleToTotalPixelsNode, comfyConfigured } from '~/server/utils/comfy'
 import { imageDimensions } from '~/server/utils/resolution'
-import { clampImageCfg, clampImageSteps } from '~/utils/generationPresets'
+import { clampImageCfg, clampImageScaleMegapixels, clampImageSteps, parseImageScaleToTotalPixels } from '~/utils/generationPresets'
 
 function parsePasses(raw: string | undefined) {
   if (!raw) return [] as { prompt: string }[]
@@ -66,6 +66,16 @@ export default defineEventHandler(async (event) => {
   const hideThumbnail = fields.hideThumbnail === 'true'
   const steps = clampImageSteps(fields.steps)
   const cfg = clampImageCfg(fields.cfg)
+  const scaleToTotalPixels = parseImageScaleToTotalPixels(fields.scaleToTotalPixels)
+  const scaleMegapixels = clampImageScaleMegapixels(fields.scaleMegapixels)
+  if (scaleToTotalPixels) {
+    try {
+      await assertImageScaleToTotalPixelsNode()
+    } catch (error) {
+      const status = Number((error as { statusCode?: number })?.statusCode || 0)
+      if (status === 400) throw error
+    }
+  }
   const seed = fields.seed && fields.seed !== 'random'
     ? Number(fields.seed)
     : Math.floor(Math.random() * 2_147_483_647)
@@ -130,6 +140,8 @@ export default defineEventHandler(async (event) => {
       passes,
       referenceStillId: savedRef?.id,
       referenceStillFilename: savedRef?.filename,
+      scaleToTotalPixels,
+      scaleMegapixels,
       ...persistLoraFields(loraStack)
     }
   })

server/assets/workflow_flux2_klein_edit.json

@@ -405,21 +405,6 @@
       "title": "Empty Flux 2 Latent"
     }
   },
-  "92:85": {
-    "inputs": {
-      "upscale_method": "lanczos",
-      "megapixels": 1,
-      "resolution_steps": 1,
-      "image": [
-        "81",
-        0
-      ]
-    },
-    "class_type": "ImageScaleToTotalPixels",
-    "_meta": {
-      "title": "Scale Image to Total Pixels"
-    }
-  },
   "92:111": {
     "inputs": {
       "clip_name": "qwen_3_8b_fp8mixed.safetensors",
@@ -562,7 +547,7 @@
   "92:129": {
     "inputs": {
       "pixels": [
-        "92:85",
+        "81",
         0
       ],
       "vae": [

server/utils/comfy.ts

@@ -51,6 +51,34 @@ export async function comfyFetch(path: string, init?: RequestInit) {
   }
 }
 
+/** true = installed, false = missing, null = Comfy unreachable / object_info failed */
+export async function comfyHasClassType(classType: string): Promise<boolean | null> {
+  try {
+    const res = await comfyFetch(`/object_info/${encodeURIComponent(classType)}`, {
+      signal: AbortSignal.timeout(8000)
+    })
+    if (res.status === 404) return false
+    if (!res.ok) return null
+    const body = await res.json().catch(() => null)
+    if (!body || typeof body !== 'object') return false
+    const rec = body as Record<string, unknown>
+    if (rec[classType] && typeof rec[classType] === 'object') return true
+    return String(rec.name || rec.class_type || '') === classType
+  } catch {
+    return null
+  }
+}
+
+export const IMAGE_SCALE_TO_TOTAL_PIXELS = 'ImageScaleToTotalPixels'
+export const IMAGE_SCALE_MISSING_MESSAGE = 'Scale Image to Total Pixels (ImageScaleToTotalPixels) is not installed on Beast. Turn Scale to 1 MP off, or install the node.'
+
+export async function assertImageScaleToTotalPixelsNode() {
+  const present = await comfyHasClassType(IMAGE_SCALE_TO_TOTAL_PIXELS)
+  if (present === false) {
+    throw createError({ statusCode: 400, statusMessage: IMAGE_SCALE_MISSING_MESSAGE })
+  }
+}
+
 export async function uploadImage(file: { filename: string; data: Buffer; type?: string }, jobId?: string) {
   const body = new FormData()
   const blob = new Blob([new Uint8Array(file.data)], { type: file.type || 'application/octet-stream' })

server/utils/generationPresets.ts

@@ -4,6 +4,7 @@ import { ownerRoot } from '~/server/utils/library'
 import { filterLoraStackForStudio, normalizeLoraStack } from '~/utils/loras'
 import {
   clampImageCfg,
+  clampImageScaleMegapixels,
   clampImageSteps,
   clampVideoCfg,
   normalizePresetName,
@@ -89,7 +90,9 @@ function sanitizeImageSettings(raw: unknown): ImagePresetSettings {
   return {
     steps: clampImageSteps(rec.steps),
     cfg: clampImageCfg(rec.cfg),
-    negative: String(rec.negative || '').slice(0, 2000)
+    negative: String(rec.negative || '').slice(0, 2000),
+    scaleToTotalPixels: rec.scaleToTotalPixels === true,
+    scaleMegapixels: clampImageScaleMegapixels(rec.scaleMegapixels)
   }
 }
 

server/utils/imageChain.ts

@@ -1,6 +1,6 @@
 import { createJob, emitJob, type Job } from '~/server/utils/jobs'
 import { ensureComfyReady } from '~/server/utils/comfyLifecycle'
-import { getComfyHost, uploadImage, queuePrompt, purgeComfyArtifacts } from '~/server/utils/comfy'
+import { assertImageScaleToTotalPixelsNode, getComfyHost, uploadImage, queuePrompt, purgeComfyArtifacts } from '~/server/utils/comfy'
 import { withImageComfyHost, waitForImageEdit, downloadEditedImage } from '~/server/utils/imageComfy'
 import { buildEditWorkflow } from '~/server/utils/imageWorkflow'
 import { ensureComfyLoraNames } from '~/server/utils/loras'
@@ -21,6 +21,8 @@ export type EditRunParams = {
   cfg: number
   loraName?: string
   loraStack?: import('~/utils/loras').LoraStackItem[]
+  scaleToTotalPixels?: boolean
+  scaleMegapixels?: number
 }
 
 export async function runEdit(job: Job, params: EditRunParams) {
@@ -93,6 +95,7 @@ export async function runEdit(job: Job, params: EditRunParams) {
           progress: 12
         })
         await ensureComfyLoraNames('image')
+        if (params.scaleToTotalPixels) await assertImageScaleToTotalPixelsNode()
         const graph = buildEditWorkflow({
           prompt,
           negative: params.negative,
@@ -102,7 +105,9 @@ export async function runEdit(job: Job, params: EditRunParams) {
           seed,
           cfg: params.cfg,
           filenamePrefix: `aigen_edit_${job.id.slice(0, 8)}_p${index + 1}`,
-          loraStack: params.loraStack || params.loraName
+          loraStack: params.loraStack || params.loraName,
+          scaleToTotalPixels: params.scaleToTotalPixels,
+          scaleMegapixels: params.scaleMegapixels
         })
         const queued = await queuePrompt(graph, job.clientId)
         job.promptId = queued.prompt_id

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',

server/utils/studioQueue.ts

@@ -47,6 +47,8 @@ export interface StudioJobPayload {
   passes?: { prompt: string }[]
   referenceStillId?: string
   referenceStillFilename?: string
+  scaleToTotalPixels?: boolean
+  scaleMegapixels?: number
 }
 
 export interface StudioJob {
@@ -656,6 +658,8 @@ async function startStudioEditJob(item: StudioJob) {
       steps: payload.steps,
       seed: job.library?.seed || payload.seed,
       cfg: payload.cfg,
+      scaleToTotalPixels: payload.scaleToTotalPixels === true,
+      scaleMegapixels: payload.scaleMegapixels,
       ...persistLoraFields(payload.loraStack || payload.loraName)
     }).catch((error) => {
       const message = error instanceof Error ? error.message : String(error)