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

Add local queued video upscale with Real-ESRGAN and RIFE

Browse Source
This commit is contained in:
Towsty
2026-09-08 17:55:02 -05:00
parent 4c97980646
commit 1e0cff9dea
24 changed files with 747 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/comfy-host-agent.mjs
+19 -5
View File
@@ -1,3 +1,4 @@
import { createUpscaleHost } from './upscale-host.mjs'
import { createGpuReservation } from './gpu-reservation.mjs'
import { createGpuProxy } from './gpu-proxy.mjs'
import { createYueGpHost } from './yuegp-host.mjs'
@@ -144,7 +145,7 @@ let proxyTarget = 0
function ensureProxyListening() {
if (proxyServer) return
proxyServer = createGpuProxy({ target: () => proxyTarget, reservation: gpuReservation, authorized, markWork, externalBusy: () => yueGp.busy() })
proxyServer = createGpuProxy({ target: () => proxyTarget, reservation: gpuReservation, authorized, markWork, externalBusy: () => yueGp.busy() || upscale.busy() })
proxyServer.on('error', (error) => {
console.log(JSON.stringify({ src: 'comfy-host-agent', event: 'proxy-error', error: String(error.message || error) }))
})
@@ -813,7 +814,7 @@ function purgeDesktopFiles(body) {
}
const gpuReservation = createGpuReservation({ idle: async () => {
if (yueGp.busy()) return false
if (yueGp.busy() || upscale.busy()) return false
if ((await trainingLock()).busy) return false
const healthy = await syncProxy()
if (healthy) {
@@ -839,9 +840,20 @@ const yueGp = createYueGpHost({
}
})
const upscale = createUpscaleHost({ leaseValid: token => gpuReservation.isOwner(token) })
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('/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'] || '')))
if (match && req.method === 'PUT' && match[2] === '/input') { await upscale.upload(match[1], req); return json(res, 200, { uploaded: true }) }
if (match && req.method === 'POST' && match[2] === '/cancel') return json(res, 200, await upscale.cancel(match[1]))
if (match && req.method === 'GET' && match[2] === '/video') { const path = upscale.output(match[1]); return path ? streamFile(res, path) : json(res, 404, { error: 'Video not ready' }) }
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' })
@@ -888,10 +900,12 @@ async function handleControl(req, res) {
proxyPort,
gpu: gpuReservation.availability(),
training: { busy: lastTraining.busy },
yuegp: { busy: yueGp.busy(), configured: yueGp.configured() }
yuegp: { busy: yueGp.busy(), configured: yueGp.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.' })
const training = await trainingLock()
if (training.busy) {
@@ -983,12 +997,12 @@ const server = http.createServer(async (req, res) => {
if (!authorized(req)) return json(res, 401, { ok: false, error: 'unauthorized' })
try {
const path = new URL(req.url || '/', 'http://localhost').pathname
if (req.method === 'POST' && !path.startsWith('/gpu/')) {
if (['POST', 'PUT'].includes(req.method) && !path.startsWith('/gpu/')) {
await gpuReservation.permit(String(req.headers['x-aigen-gpu-lease'] || ''), () => handleControl(req, res))
} else await handleControl(req, res)
} catch (error) {
req.resume()
if (String(req.url || '').startsWith('/yuegp/') && !res.headersSent) return json(res, error.statusCode || 400, { error: error.message || 'YuEGP request failed' })
if ((String(req.url || '').startsWith('/yuegp/') || String(req.url || '').startsWith('/upscale/')) && !res.headersSent) return json(res, error.statusCode || 400, { error: error.message || 'YuEGP 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.' })
}
})