towsty/aigen
1
0
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases Wiki Activity

Run image edit on Beast with Flux.2 Klein and stop routing jobs to Sidecar.

Browse Source 
Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Towsty
2026-08-27 09:06:50 -05:00
co-authored by Cursor
parent 6fec10febd
commit 13ecec949c
8 changed files with 899 additions and 212 deletions
Download Patch File Download Diff File Expand all files Collapse all files
server/utils/imageComfy.ts
+6 -6
View File
@@ -55,7 +55,7 @@ export function getImageComfyHost() {
}
export function imageComfyConfigured() {
return Boolean(getBeastImageHost() || getSidecarImageHost())
return Boolean(getBeastImageHost())
}
export function sameImageHost(a: string, b: string) {
@@ -246,17 +246,17 @@ export async function waitForImageEdit(opts: {
const max = Math.max(1, Number(payload.data.max || 4))
const value = Number(payload.data.value || 0)
opts.onProgress?.({
message: 'Sampling edit...',
message: 'Sampling Flux.2 Klein...',
progress: Math.min(90, 20 + Math.round((value / max) * 60)),
step: value,
maxStep: max,
node: payload.data.node || '8'
node: payload.data.node || null
})
}
if (payload.type === 'executing' && payload.data?.node) {
opts.onProgress?.({
message: payload.data.node === '8' ? 'Sampling edit...' : 'Running Qwen Edit...',
progress: payload.data.node === '10' ? 92 : 30,
message: 'Running Flux.2 Klein...',
progress: payload.data.node === '9' || payload.data.node === '94' ? 92 : 30,
node: payload.data.node
})
}
@@ -291,7 +291,7 @@ export async function waitForImageEdit(opts: {
if (result) return result
throw createError({
statusCode: 502,
statusMessage: lastError || 'Image edit timed out waiting for Qwen Edit to finish'
statusMessage: lastError || 'Image edit timed out waiting for Flux.2 Klein to finish'
})
}
server/utils/imageComfyLifecycle.ts
+27 -54
View File
@@ -1,4 +1,4 @@
import { getImageComfyHost, setImageComfyHostOverride, imageComfyConfigured, getBeastImageHost, getSidecarImageHost, sameImageHost, withImageComfyHost, type ImageEditBox } from '~/server/utils/imageComfy'
import { getImageComfyHost, setImageComfyHostOverride, imageComfyConfigured, getBeastImageHost, withImageComfyHost, type ImageEditBox } from '~/server/utils/imageComfy'
export type ImageComfyState = 'online' | 'booting' | 'busy' | 'starting' | 'offline'
@@ -231,67 +231,40 @@ export function ensureImageComfyReady(onStatus: StatusFn = () => undefined) {
export async function pickImageEditHost(onStatus: StatusFn = () => undefined): Promise<ImageEditTarget> {
const beast = getBeastImageHost()
const sidecar = getSidecarImageHost()
if (!beast && !sidecar) {
if (!beast) {
throw createError({
statusCode: 503,
statusMessage: 'No image ComfyUI host is configured. Set COMFY_HOST (Beast) and/or IMAGE_COMFY_HOST (Sidecar).'
statusMessage: 'Beast ComfyUI is not configured. Set COMFY_HOST.'
})
}
if (beast) {
const health = await checkImageComfyHttp(2500, beast)
if (health.ok) {
const queue = await fetchQueueAt(beast)
if (queue.running === 0) {
onStatus({
state: 'online',
message: 'Beast is free. Editing on this machine.',
host: beast,
box: 'beast',
queueRunning: 0,
queuePending: queue.pending
})
return { host: beast, box: 'beast', label: 'Beast' }
}
const canOverflow = Boolean(sidecar && !sameImageHost(beast, sidecar))
onStatus({
state: 'busy',
message: canOverflow
? 'Beast is busy. Rolling over to Sidecar...'
: `Beast is busy (${queue.running} running). Waiting for a free slot...`,
host: beast,
box: 'beast',
queueRunning: queue.running,
queuePending: queue.pending
})
if (!canOverflow) {
await withImageComfyHost(beast, () => waitWhileBusy(onStatus))
return { host: beast, box: 'beast', label: 'Beast' }
}
} else if (sidecar && !sameImageHost(beast, sidecar)) {
onStatus({
state: 'offline',
message: 'Beast is offline. Trying Sidecar...',
host: beast,
box: 'beast'
})
}
const health = await checkImageComfyHttp(2500, beast)
if (!health.ok) {
throw createError({
statusCode: 503,
statusMessage: 'Beast is offline. Image edit runs on the 5080 only.'
})
}
if (sidecar && (!beast || !sameImageHost(beast, sidecar))) {
const queue = await fetchQueueAt(beast)
if (queue.running === 0) {
onStatus({
state: 'starting',
message: 'Checking Sidecar for this edit...',
host: sidecar,
box: 'sidecar'
state: 'online',
message: 'Beast is free. Editing on this machine.',
host: beast,
box: 'beast',
queueRunning: 0,
queuePending: queue.pending
})
await withImageComfyHost(sidecar, () => ensureImageComfyReady(onStatus))
return { host: sidecar, box: 'sidecar', label: 'Sidecar' }
return { host: beast, box: 'beast', label: 'Beast' }
}
throw createError({
statusCode: 503,
statusMessage: 'Beast is unavailable and Sidecar is not configured.'
onStatus({
state: 'busy',
message: `Beast is busy (${queue.running} running). Waiting for a free slot...`,
host: beast,
box: 'beast',
queueRunning: queue.running,
queuePending: queue.pending
})
await withImageComfyHost(beast, () => waitWhileBusy(onStatus))
return { host: beast, box: 'beast', label: 'Beast' }
}
server/utils/imageWorkflow.ts
+75 -37
View File
@@ -1,61 +1,99 @@
import template from '../assets/workflow_qwen_edit.json'
import template from '../assets/workflow_flux2_klein_edit.json'
export const DEFAULT_EDIT_NEGATIVE = 'blurry, low quality, deformed, bad anatomy, watermark, text, extra fingers'
const BYPASS = 4
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'
type WorkflowNode = { class_type: string; inputs: Record<string, unknown>; _meta?: { title?: string } }
type WorkflowNode = { class_type: string; inputs: Record<string, unknown>; mode?: number; _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
denoise: number
samplerName?: string
scheduler?: string
filenamePrefix?: string
}
export function buildEditWorkflow(params: EditParams) {
const graph = structuredClone(template) as WorkflowGraph
const loadImage = nodeByClass(graph, 'LoadImage') || graph['4']
const qwenPrompt = nodeByClass(graph, 'TextEncodeQwenImageEdit') || graph['17']
const sampler = nodeByClass(graph, 'KSampler') || graph['8']
const save = nodeByClass(graph, 'SaveImage') || graph['10']
if (loadImage) loadImage.inputs.image = params.imageName
if (qwenPrompt) {
if ('prompt' in qwenPrompt.inputs) qwenPrompt.inputs.prompt = params.prompt
if ('text' in qwenPrompt.inputs) qwenPrompt.inputs.text = params.prompt
const dual = Boolean(params.referenceImageName)
const prompt = String(params.prompt || '').trim()
const negative = String(params.negative || '')
const prefix = params.filenamePrefix || 'aigen-edit'
setImage(graph, IMAGE1, params.imageName)
setImage(graph, IMAGE2, params.referenceImageName || '')
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)
for (const [id, node] of Object.entries(graph)) {
if (!node) continue
if (id === IMAGE1) {
delete node.mode
continue
}
const dualNode = id === IMAGE2 || id === DUAL_SAVE || id.startsWith('92:')
const singleNode = id === SINGLE_SAVE || id.startsWith('75:')
if (dual ? singleNode : dualNode) node.mode = BYPASS
else delete node.mode
}
if (sampler) {
sampler.inputs.seed = params.seed
sampler.inputs.steps = params.steps
sampler.inputs.cfg = params.cfg
sampler.inputs.denoise = params.denoise
sampler.inputs.sampler_name = params.samplerName || 'euler'
sampler.inputs.scheduler = params.scheduler || 'simple'
}
if (save && params.filenamePrefix) save.inputs.filename_prefix = params.filenamePrefix
return graph
}
function nodeByClass(graph: WorkflowGraph, classType: string, index = 0) {
const matches = Object.values(graph).filter(node => node?.class_type === classType)
return matches[index]
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> = {
'1': 'Loading Qwen Edit UNET',
'11': 'Applying Lightning LoRA',
'2': 'Loading Qwen CLIP',
'3': 'Loading Qwen VAE',
'4': 'Loading input still',
'17': 'Encoding Qwen edit prompt',
'19': 'Preparing negative',
'7': 'Encoding still',
'8': 'Sampling edit',
'9': 'Decoding still',
'10': 'Saving still'
'76': 'Loading image 1',
'81': 'Loading 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'
}