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

Remove YuEGP handlers from the Comfy host agent.

Browse Source 
Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Towsty
2026-09-15 18:17:59 -05:00
co-authored by Cursor
parent cce1faa9b1
commit a216afc4ce
2 changed files with 6 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/comfy-host-agent.mjs
+4 -40
View File
@@ -3,7 +3,6 @@ import { purgeStudio2Files } from './studio2-purge.mjs'
import { stableMemoryArgs } from './comfy-memory-policy.mjs'
import { createGpuReservation } from './gpu-reservation.mjs'
import { createGpuProxy } from './gpu-proxy.mjs'
import { createYueGpHost } from './yuegp-host.mjs'
import { createYue2Host } from './yue2-host.mjs'
import http from 'node:http'
import net from 'node:net'
@@ -148,7 +147,7 @@ let proxyTarget = 0
function ensureProxyListening() {
if (proxyServer) return
proxyServer = createGpuProxy({ target: () => proxyTarget, reservation: gpuReservation, authorized, markWork, externalBusy: () => yueGp.busy() || yue2.busy() || upscale.busy() })
proxyServer = createGpuProxy({ target: () => proxyTarget, reservation: gpuReservation, authorized, markWork, externalBusy: () => yue2.busy() || upscale.busy() })
proxyServer.on('error', (error) => {
console.log(JSON.stringify({ src: 'comfy-host-agent', event: 'proxy-error', error: String(error.message || error) }))
})
@@ -820,7 +819,7 @@ function purgeDesktopFiles(body) {
}
const gpuReservation = createGpuReservation({ idle: async () => {
if (yueGp.busy() || yue2.busy() || upscale.busy()) return false
if (yue2.busy() || upscale.busy()) return false
if ((await trainingLock()).busy) return false
const healthy = await syncProxy()
if (healthy) {
@@ -831,33 +830,16 @@ const gpuReservation = createGpuReservation({ idle: async () => {
return !(await processUp()) && !(await pythonMainUp().catch(() => true))
} })
const yueGp = createYueGpHost({
leaseValid: lease => gpuReservation.isOwner(lease),
prepare: async () => {
if ((await trainingLock()).busy) throw new Error('GPU is busy with training.')
if (yue2.busy()) throw new Error('YuE2 is using the GPU.')
const healthy = await syncProxy()
if (healthy) {
const queue = await fetchLocalQueue(healthy)
if (!queue.ok || queue.running || queue.pending) throw new Error('Comfy is busy; YuEGP cannot start.')
await stopComfyProcesses()
markAsleep()
}
if (await pythonMainUp()) throw new Error('Comfy has not stopped; retry after the GPU is free.')
}
})
const yue2 = createYue2Host({
leaseValid: lease => gpuReservation.isOwner(lease),
prepare: async () => {
if ((await trainingLock()).busy) throw new Error('GPU is busy with training.')
if (yueGp.busy()) throw new Error('YuEGP is using the GPU.')
const healthy = await syncProxy()
if (healthy) {
const queue = await fetchLocalQueue(healthy)
if (!queue.ok || queue.running || queue.pending) throw new Error('Comfy is busy; YuE2 cannot start.')
}
// Same gate as YuEGP: stop Comfy and refuse to launch while its Python still owns VRAM.
// Stop Comfy and refuse to launch while its Python still owns VRAM.
if (healthy || await processUp() || await pythonMainUp().catch(() => false)) {
await stopComfyProcesses()
markAsleep()
@@ -880,23 +862,7 @@ async function handleControl(req, res) {
if (match && req.method === 'GET' && !match[2]) { const job = upscale.read(match[1]); return json(res, job ? 200 : 404, job || { error: 'Upscale job not found' }) }
return json(res, 404, { error: 'Unknown upscale endpoint' })
}
if (url.pathname.startsWith('/yuegp/')) {
const match = url.pathname.match(/^\/yuegp\/jobs\/([a-zA-Z0-9-]{12,80})(\/audio|\/cancel)?$/)
if (req.method === 'GET' && url.pathname === '/yuegp/status') return json(res, 200, { configured: yueGp.configured(), busy: yueGp.busy(), backend: 'yuegp' })
if (req.method === 'POST' && url.pathname === '/yuegp/jobs') return json(res, 200, await yueGp.start(await readJson(req), String(req.headers['x-aigen-gpu-lease'] || '')))
if (match && req.method === 'POST' && match[2] === '/cancel') return json(res, 200, await yueGp.cancel(match[1]))
if (match && req.method === 'GET' && match[2] === '/audio') {
const path = yueGp.audio(match[1])
return path ? streamFile(res, path) : json(res, 404, { error: 'Audio not ready' })
}
if (match && req.method === 'GET' && !match[2]) {
const job = yueGp.read(match[1])
return json(res, job ? 200 : 404, job || { error: 'YuEGP job not found' })
}
return json(res, 404, { error: 'Unknown YuEGP endpoint' })
}
if (url.pathname.startsWith('/yue2/')) {
if (yueGp.busy()) return json(res, 409, { message: 'YuEGP is using the GPU.' })
const match = url.pathname.match(/^\/yue2\/jobs\/([a-zA-Z0-9-]{12,80})(\/audio|\/cancel)?$/)
if (req.method === 'GET' && url.pathname === '/yue2/status') return json(res, 200, { configured: yue2.configured(), busy: yue2.busy(), backend: 'yue2' })
if (req.method === 'POST' && url.pathname === '/yue2/jobs') return json(res, 200, await yue2.start(await readJson(req), String(req.headers['x-aigen-gpu-lease'] || '')))
@@ -942,14 +908,12 @@ async function handleControl(req, res) {
proxyPort,
gpu: gpuReservation.availability(),
training: { busy: lastTraining.busy },
yuegp: { busy: yueGp.busy(), configured: yueGp.configured() },
yue2: { busy: yue2.busy(), configured: yue2.configured() },
upscale: { busy: upscale.busy(), engine: 'realesrgan-rife', local: true }
})
}
if (req.method === 'POST' && url.pathname === '/start') {
if (upscale.busy()) return json(res, 409, { message: 'Local upscale is using the GPU.' })
if (yueGp.busy()) return json(res, 409, { message: 'YuEGP is using the GPU.' })
if (yue2.busy()) return json(res, 409, { message: 'YuE2 is using the GPU.' })
const training = await trainingLock()
if (training.busy) {
@@ -1052,7 +1016,7 @@ const server = http.createServer(async (req, res) => {
} else await handleControl(req, res)
} catch (error) {
req.resume()
if ((String(req.url || '').startsWith('/yuegp/') || String(req.url || '').startsWith('/yue2/') || String(req.url || '').startsWith('/upscale/')) && !res.headersSent) return json(res, error.statusCode || 400, { error: error.message || 'Music host request failed' })
if ((String(req.url || '').startsWith('/yue2/') || String(req.url || '').startsWith('/upscale/')) && !res.headersSent) return json(res, error.statusCode || 400, { error: error.message || 'Music host request failed' })
if (!res.headersSent) json(res, error.statusCode || 400, { ok: false, message: error.statusCode === 409 ? 'GPU is in use. Waiting for availability.' : 'GPU coordination request failed.' })
}
})