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

Add image-to-text Describe on the bench via exclusive llama.cpp caption jobs.

Browse Source
This commit is contained in:
Towsty
2026-09-30 22:50:19 -05:00
parent 9481791ac5
commit 15ef891b26
15 changed files with 937 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/comfy-host-agent.mjs
+58 -19
View File
@@ -4,6 +4,7 @@ import { stableMemoryArgs } from './comfy-memory-policy.mjs'
import { createGpuReservation } from './gpu-reservation.mjs'
import { createGpuProxy } from './gpu-proxy.mjs'
import { createYue2Host } from './yue2-host.mjs'
import { createCaptionHost } from './caption-host.mjs'
import http from 'node:http'
import net from 'node:net'
import { execFile, spawn } from 'node:child_process'
@@ -147,7 +148,7 @@ let proxyTarget = 0
function ensureProxyListening() {
if (proxyServer) return
proxyServer = createGpuProxy({ target: () => proxyTarget, reservation: gpuReservation, authorized, markWork, externalBusy: () => yue2.busy() || upscale.busy() })
proxyServer = createGpuProxy({ target: () => proxyTarget, reservation: gpuReservation, authorized, markWork, externalBusy: () => yue2.busy() || upscale.busy() || caption.busy() })
proxyServer.on('error', (error) => {
console.log(JSON.stringify({ src: 'comfy-host-agent', event: 'proxy-error', error: String(error.message || error) }))
})
@@ -819,7 +820,7 @@ function purgeDesktopFiles(body) {
}
const gpuReservation = createGpuReservation({ idle: async () => {
if (yue2.busy() || upscale.busy()) return false
if (yue2.busy() || upscale.busy() || caption.busy()) return false
if ((await trainingLock()).busy) return false
const healthy = await syncProxy()
if (healthy) {
@@ -830,29 +831,65 @@ const gpuReservation = createGpuReservation({ idle: async () => {
return !(await processUp()) && !(await pythonMainUp().catch(() => true))
} })
const yue2 = createYue2Host({
leaseValid: lease => gpuReservation.isOwner(lease),
prepare: async () => {
if ((await trainingLock()).busy) throw new Error('GPU is busy with training.')
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.')
}
// Stop Comfy and refuse to launch while its Python still owns VRAM.
if (healthy || await processUp() || await pythonMainUp().catch(() => false)) {
await stopComfyProcesses()
markAsleep()
}
if (await pythonMainUp()) throw new Error('Comfy has not stopped; retry after the GPU is free.')
async function prepareExclusiveGpu(label) {
if ((await trainingLock()).busy) throw new Error('GPU is busy with training.')
if (label !== 'YuE2' && yue2.busy()) throw new Error('YuE2 is using the GPU.')
if (label !== 'caption' && caption.busy()) throw new Error('Caption is using the GPU.')
if (upscale.busy()) throw new Error('Local upscale 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; ${label} cannot start.`)
}
})
if (healthy || await processUp() || await pythonMainUp().catch(() => false)) {
await stopComfyProcesses()
markAsleep()
}
if (await pythonMainUp()) throw new Error('Comfy has not stopped; retry after the GPU is free.')
}
const upscale = createUpscaleHost({ leaseValid: token => gpuReservation.isOwner(token) })
const yue2 = createYue2Host({
leaseValid: lease => gpuReservation.isOwner(lease),
prepare: async () => prepareExclusiveGpu('YuE2')
})
const caption = createCaptionHost({
leaseValid: lease => gpuReservation.isOwner(lease),
prepare: async () => prepareExclusiveGpu('caption')
})
async function handleControl(req, res) {
if (!authorized(req)) return json(res, 401, { ok: false, error: 'unauthorized' })
const url = new URL(req.url || '/', 'http://localhost')
if (url.pathname.startsWith('/caption')) {
const match = url.pathname.match(/^\/caption\/jobs\/([a-zA-Z0-9-]{12,80})(\/input|\/cancel)?$/)
if (req.method === 'GET' && url.pathname === '/caption/status') {
return json(res, 200, { configured: caption.configured(), busy: caption.busy(), backend: 'llama.cpp', modelsDir: caption.modelsDir() })
}
if (req.method === 'POST' && url.pathname === '/caption') {
try {
const body = await readJson(req)
return json(res, 200, await caption.captionOnce(body, String(req.headers['x-aigen-gpu-lease'] || '')))
} catch (error) {
return json(res, error.statusCode || 400, { error: error.message || 'Caption failed' })
}
}
if (req.method === 'POST' && url.pathname === '/caption/jobs') {
return json(res, 200, await caption.start(await readJson(req), String(req.headers['x-aigen-gpu-lease'] || '')))
}
if (match && req.method === 'PUT' && match[2] === '/input') {
await caption.upload(match[1], req)
return json(res, 200, { uploaded: true })
}
if (match && req.method === 'POST' && match[2] === '/cancel') return json(res, 200, await caption.cancel(match[1]))
if (match && req.method === 'GET' && !match[2]) {
const job = caption.read(match[1])
return json(res, job ? 200 : 404, job || { error: 'Caption job not found' })
}
return json(res, 404, { error: 'Unknown caption endpoint' })
}
if (url.pathname.startsWith('/upscale/')) {
const match = url.pathname.match(/^\/upscale\/jobs\/([a-zA-Z0-9-]{12,80})(\/input|\/video|\/cancel)?$/)
if (req.method === 'POST' && url.pathname === '/upscale/jobs') return json(res, 200, upscale.start(await readJson(req), String(req.headers['x-aigen-gpu-lease'] || '')))
@@ -909,12 +946,14 @@ async function handleControl(req, res) {
gpu: gpuReservation.availability(),
training: { busy: lastTraining.busy },
yue2: { busy: yue2.busy(), configured: yue2.configured() },
caption: { busy: caption.busy(), configured: caption.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 (yue2.busy()) return json(res, 409, { message: 'YuE2 is using the GPU.' })
if (caption.busy()) return json(res, 409, { message: 'Caption is using the GPU.' })
const training = await trainingLock()
if (training.busy) {
return json(res, 409, {
@@ -1016,7 +1055,7 @@ const server = http.createServer(async (req, res) => {
} else await handleControl(req, res)
} catch (error) {
req.resume()
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 ((String(req.url || '').startsWith('/yue2/') || String(req.url || '').startsWith('/upscale/') || String(req.url || '').startsWith('/caption')) && !res.headersSent) return json(res, error.statusCode || 400, { error: error.message || '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.' })
}
})