Add Studio 2 caption jobs via Qwen2.5-VL GGUF on the host agent.
Exclusive llama-server load/unload on the 5080, queued Describe UI with resultText, Copy, and Use as prompt. Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
@@ -0,0 +1,331 @@
|
||||
import { spawn } from 'node:child_process'
|
||||
import { createServer } from 'node:net'
|
||||
import { existsSync, mkdirSync, readFileSync, writeFileSync, renameSync, appendFileSync, readdirSync, copyFileSync } from 'node:fs'
|
||||
import { join, resolve, extname } from 'node:path'
|
||||
import { captionStylePrompt, CAPTION_STYLES } from '../shared/studio2/caption.mjs'
|
||||
|
||||
const MODEL_NAME = 'Qwen2.5-VL-7B-NSFW-Caption-V4.Q5_K_M.gguf'
|
||||
const MMPROJ_NAME = 'Qwen2.5-VL-7B-NSFW-Caption-V4.mmproj-f16.gguf'
|
||||
|
||||
function defaultModelsDir() {
|
||||
const shared = process.env.COMFY_MODELS_ROOT
|
||||
|| join(process.env.LOCALAPPDATA || '', 'Comfy-Desktop', 'ComfyUI-Shared', 'models')
|
||||
return resolve(process.env.CAPTION_MODELS || join(shared, 'caption', 'qwen25vl-7b-nsfw-v4'))
|
||||
}
|
||||
|
||||
function defaultLlamaServer() {
|
||||
if (process.env.CAPTION_LLAMA_SERVER) return resolve(process.env.CAPTION_LLAMA_SERVER)
|
||||
return 'llama-server'
|
||||
}
|
||||
|
||||
function freePort() {
|
||||
return new Promise((resolvePort, reject) => {
|
||||
const server = createServer()
|
||||
server.listen(0, '127.0.0.1', () => {
|
||||
const address = server.address()
|
||||
const port = typeof address === 'object' && address ? address.port : 0
|
||||
server.close(error => error ? reject(error) : resolvePort(port))
|
||||
})
|
||||
server.on('error', reject)
|
||||
})
|
||||
}
|
||||
|
||||
async function waitForServer(port, signal, timeoutMs = 120_000, fetchImpl = fetch) {
|
||||
const started = Date.now()
|
||||
while (Date.now() - started < timeoutMs) {
|
||||
if (signal?.aborted) throw new Error('Caption cancelled while loading the model.')
|
||||
try {
|
||||
const response = await fetchImpl(`http://127.0.0.1:${port}/health`, { signal: AbortSignal.timeout(2000) })
|
||||
if (response.ok) return
|
||||
} catch { /* still booting */ }
|
||||
await new Promise(r => setTimeout(r, 500))
|
||||
}
|
||||
throw new Error('llama-server did not become ready in time.')
|
||||
}
|
||||
|
||||
function mimeFor(path) {
|
||||
const ext = extname(path).toLowerCase()
|
||||
if (ext === '.jpg' || ext === '.jpeg') return 'image/jpeg'
|
||||
if (ext === '.webp') return 'image/webp'
|
||||
if (ext === '.gif') return 'image/gif'
|
||||
return 'image/png'
|
||||
}
|
||||
|
||||
function killTree(child) {
|
||||
return new Promise(resolveKill => {
|
||||
if (!child?.pid) return resolveKill()
|
||||
let finished = false
|
||||
const done = () => {
|
||||
if (finished) return
|
||||
finished = true
|
||||
resolveKill()
|
||||
}
|
||||
child.once('close', done)
|
||||
try {
|
||||
if (typeof child.kill === 'function') child.kill()
|
||||
} catch { /* ignore */ }
|
||||
if (process.platform === 'win32' && child.pid > 0) {
|
||||
try {
|
||||
const killer = spawn('taskkill', ['/PID', String(child.pid), '/T', '/F'], { windowsHide: true, stdio: 'ignore' })
|
||||
killer.once('close', done)
|
||||
killer.once('error', done)
|
||||
} catch {
|
||||
done()
|
||||
}
|
||||
}
|
||||
setTimeout(done, 3000)
|
||||
})
|
||||
}
|
||||
|
||||
export function validateCaptionHostRequest(body) {
|
||||
if (!/^[a-zA-Z0-9-]{12,80}$/.test(body?.id || '')) throw new Error('Invalid caption job ID.')
|
||||
const style = String(body.style || body.captionStyle || 'descriptive')
|
||||
if (!CAPTION_STYLES.includes(style)) throw new Error('Unknown caption style.')
|
||||
return { id: body.id, style, imagePath: body.imagePath ? String(body.imagePath) : '' }
|
||||
}
|
||||
|
||||
/** One llama-server process per caption. Process exit is the VRAM unload boundary (keep_alive 0). */
|
||||
export function createCaptionHost({
|
||||
prepare,
|
||||
leaseValid,
|
||||
spawnProcess = spawn,
|
||||
modelsDir,
|
||||
llamaServer,
|
||||
dataDir,
|
||||
fetchImpl = fetch
|
||||
} = {}) {
|
||||
const models = resolve(modelsDir || defaultModelsDir())
|
||||
const executable = llamaServer || defaultLlamaServer()
|
||||
const data = resolve(dataDir || process.env.CAPTION_JOBS_DIR || join(models, 'aigen-jobs'))
|
||||
let active = null
|
||||
let holdUntil = 0
|
||||
if (existsSync(data)) {
|
||||
for (const id of readdirSync(data)) {
|
||||
try {
|
||||
const path = join(data, id, 'status.json')
|
||||
if (!existsSync(path)) continue
|
||||
const state = JSON.parse(readFileSync(path, 'utf8'))
|
||||
if (['running', 'starting', 'cancelling'].includes(state.status)) {
|
||||
holdUntil = Date.now() + 10000
|
||||
state.status = 'error'
|
||||
state.error = 'Caption host restarted. GPU model unloaded.'
|
||||
writeFileSync(path + '.tmp', JSON.stringify(state))
|
||||
renameSync(path + '.tmp', path)
|
||||
}
|
||||
} catch { /* ignore bad records */ }
|
||||
}
|
||||
}
|
||||
const dir = id => {
|
||||
if (!/^[a-zA-Z0-9-]{12,80}$/.test(id || '')) throw new Error('Invalid caption job ID.')
|
||||
return join(data, id)
|
||||
}
|
||||
const persist = job => {
|
||||
const target = join(dir(job.id), 'status.json')
|
||||
writeFileSync(target + '.tmp', JSON.stringify(job))
|
||||
renameSync(target + '.tmp', target)
|
||||
}
|
||||
const read = id => {
|
||||
if (active?.job.id === id) return { ...active.job, text: active.job.text }
|
||||
const path = join(dir(id), 'status.json')
|
||||
if (!existsSync(path)) return null
|
||||
return JSON.parse(readFileSync(path, 'utf8'))
|
||||
}
|
||||
const modelPath = () => join(models, MODEL_NAME)
|
||||
const mmprojPath = () => join(models, MMPROJ_NAME)
|
||||
const configured = () => {
|
||||
try {
|
||||
if (!existsSync(modelPath()) || !existsSync(mmprojPath())) return false
|
||||
if (executable.includes('\\') || executable.includes('/')) return existsSync(executable)
|
||||
return true
|
||||
} catch { return false }
|
||||
}
|
||||
|
||||
return {
|
||||
busy: () => Boolean(active) || Date.now() < holdUntil,
|
||||
configured,
|
||||
modelsDir: () => models,
|
||||
read,
|
||||
async upload(id, stream) {
|
||||
if (read(id)?.status === 'complete') throw new Error('This caption job has already finished.')
|
||||
const jobDir = dir(id)
|
||||
mkdirSync(jobDir, { recursive: true })
|
||||
const dest = join(jobDir, 'input.upload')
|
||||
const { createWriteStream } = await import('node:fs')
|
||||
const { pipeline } = await import('node:stream/promises')
|
||||
const { Transform } = await import('node:stream')
|
||||
let bytes = 0
|
||||
await pipeline(stream, new Transform({
|
||||
transform(chunk, _, callback) {
|
||||
bytes += chunk.length
|
||||
callback(bytes > 40 * 1024 * 1024 ? new Error('Image exceeds the 40 MB upload limit.') : null, chunk)
|
||||
}
|
||||
}), createWriteStream(dest, { flags: 'w' }))
|
||||
if (bytes < 32) throw new Error('Image file is empty.')
|
||||
renameSync(dest, join(jobDir, 'input.png'))
|
||||
},
|
||||
async start(body, lease) {
|
||||
const request = validateCaptionHostRequest(body)
|
||||
const previous = read(request.id)
|
||||
if (previous) return previous
|
||||
if (active || Date.now() < holdUntil) throw Object.assign(new Error('Caption is already running or releasing VRAM.'), { statusCode: 409 })
|
||||
if (!configured()) throw new Error(`Caption model is not installed. Expected ${MODEL_NAME} and ${MMPROJ_NAME} under ${models}. Run scripts/setup-caption.ps1.`)
|
||||
if (!leaseValid(lease)) throw new Error('GPU reservation expired.')
|
||||
|
||||
let imagePath = request.imagePath
|
||||
if (imagePath) {
|
||||
if (!existsSync(imagePath)) throw new Error('imagePath does not exist on the GPU host.')
|
||||
mkdirSync(dir(request.id), { recursive: true })
|
||||
const dest = join(dir(request.id), `input${extname(imagePath) || '.png'}`)
|
||||
copyFileSync(imagePath, dest)
|
||||
imagePath = dest
|
||||
} else {
|
||||
imagePath = join(dir(request.id), 'input.png')
|
||||
if (!existsSync(imagePath)) {
|
||||
const alt = readdirSync(dir(request.id)).find(name => /^input\./i.test(name))
|
||||
if (!alt) throw new Error('Upload an image first or pass imagePath.')
|
||||
imagePath = join(dir(request.id), alt)
|
||||
}
|
||||
}
|
||||
|
||||
const job = {
|
||||
id: request.id,
|
||||
status: 'starting',
|
||||
message: 'Preparing GPU for caption',
|
||||
progress: 0,
|
||||
style: request.style,
|
||||
startedAt: Date.now(),
|
||||
checkedAt: Date.now()
|
||||
}
|
||||
active = { job, child: null, cancelled: false, abort: new AbortController() }
|
||||
const run = active
|
||||
mkdirSync(dir(job.id), { recursive: true })
|
||||
writeFileSync(join(dir(job.id), 'request.json'), JSON.stringify({ ...request, imagePath, prompt: captionStylePrompt(request.style) }))
|
||||
persist(job)
|
||||
|
||||
try {
|
||||
if (prepare) await prepare()
|
||||
if (run.cancelled || !leaseValid(lease)) throw new Error('Caption start cancelled or GPU reservation expired.')
|
||||
|
||||
job.status = 'running'
|
||||
job.message = 'Loading Qwen2.5-VL caption model'
|
||||
job.progress = 5
|
||||
persist(job)
|
||||
|
||||
const port = await freePort()
|
||||
const args = [
|
||||
'-m', modelPath(),
|
||||
'--mmproj', mmprojPath(),
|
||||
'--host', '127.0.0.1',
|
||||
'--port', String(port),
|
||||
'-ngl', String(process.env.CAPTION_NGL || '99'),
|
||||
'-c', String(process.env.CAPTION_CTX || '4096'),
|
||||
'--jinja'
|
||||
]
|
||||
const child = spawnProcess(executable, args, {
|
||||
windowsHide: true,
|
||||
shell: false,
|
||||
stdio: ['ignore', 'pipe', 'pipe'],
|
||||
env: { ...process.env }
|
||||
})
|
||||
run.child = child
|
||||
const log = chunk => appendFileSync(join(dir(job.id), 'worker.log'), chunk)
|
||||
child.stdout.on('data', log)
|
||||
child.stderr.on('data', log)
|
||||
|
||||
const watchdog = setInterval(() => {
|
||||
if (!leaseValid(lease)) {
|
||||
job.error = 'GPU reservation expired; caption stopped.'
|
||||
run.abort.abort()
|
||||
void killTree(child)
|
||||
}
|
||||
}, 5000)
|
||||
|
||||
try {
|
||||
await waitForServer(port, run.abort.signal, 120_000, fetchImpl)
|
||||
if (run.cancelled) throw new Error('Cancelled')
|
||||
job.message = 'Captioning'
|
||||
job.progress = 40
|
||||
persist(job)
|
||||
|
||||
const bytes = readFileSync(imagePath)
|
||||
const dataUrl = `data:${mimeFor(imagePath)};base64,${bytes.toString('base64')}`
|
||||
const timeout = AbortSignal.timeout(Number(process.env.CAPTION_TIMEOUT_MS || 180_000))
|
||||
const onAbort = () => { /* combined via race below */ }
|
||||
run.abort.signal.addEventListener('abort', onAbort, { once: true })
|
||||
let response
|
||||
try {
|
||||
response = await fetchImpl(`http://127.0.0.1:${port}/v1/chat/completions`, {
|
||||
method: 'POST',
|
||||
headers: { 'Content-Type': 'application/json' },
|
||||
signal: run.abort.signal.aborted ? run.abort.signal : timeout,
|
||||
body: JSON.stringify({
|
||||
temperature: 0.2,
|
||||
max_tokens: 512,
|
||||
messages: [{
|
||||
role: 'user',
|
||||
content: [
|
||||
{ type: 'text', text: captionStylePrompt(request.style) },
|
||||
{ type: 'image_url', image_url: { url: dataUrl } }
|
||||
]
|
||||
}]
|
||||
})
|
||||
})
|
||||
} finally {
|
||||
run.abort.signal.removeEventListener('abort', onAbort)
|
||||
}
|
||||
if (run.abort.signal.aborted) throw new Error('Cancelled')
|
||||
if (!response.ok) {
|
||||
const detail = await response.text().catch(() => '')
|
||||
throw new Error(`llama-server caption failed (${response.status}): ${detail.slice(0, 400)}`)
|
||||
}
|
||||
const payload = await response.json()
|
||||
const text = String(payload?.choices?.[0]?.message?.content || '').trim()
|
||||
if (!text) throw new Error('Caption model returned empty text.')
|
||||
job.text = text
|
||||
job.status = run.cancelled ? 'cancelled' : 'complete'
|
||||
job.message = job.status === 'complete' ? 'Caption ready' : 'Cancelled'
|
||||
job.progress = 100
|
||||
job.checkedAt = Date.now()
|
||||
persist(job)
|
||||
return { ...job }
|
||||
} finally {
|
||||
clearInterval(watchdog)
|
||||
await killTree(child)
|
||||
// keep_alive 0: process exit unloads GGUF+mmproj from VRAM
|
||||
holdUntil = Date.now() + 2000
|
||||
if (active === run) active = null
|
||||
}
|
||||
} catch (error) {
|
||||
job.status = run.cancelled ? 'cancelled' : 'error'
|
||||
job.error = error.message || String(error)
|
||||
job.message = job.error
|
||||
job.checkedAt = Date.now()
|
||||
persist(job)
|
||||
if (run.child) await killTree(run.child)
|
||||
if (active === run) active = null
|
||||
holdUntil = Date.now() + 2000
|
||||
throw error
|
||||
}
|
||||
},
|
||||
async cancel(id) {
|
||||
if (active?.job.id !== id) return read(id)
|
||||
const run = active
|
||||
run.cancelled = true
|
||||
run.job.status = 'cancelling'
|
||||
run.job.message = 'Cancelling'
|
||||
persist(run.job)
|
||||
run.abort.abort()
|
||||
if (run.child) await killTree(run.child)
|
||||
return read(id)
|
||||
},
|
||||
/** Sync one-shot used by POST /caption { imagePath, style }. */
|
||||
async captionOnce(body, lease) {
|
||||
const id = body.id || `caption-${Date.now()}-${Math.random().toString(36).slice(2, 10)}`
|
||||
const result = await this.start({ ...body, id }, lease)
|
||||
if (result.status === 'complete') return { text: result.text, id: result.id, style: result.style }
|
||||
throw new Error(result.error || 'Caption failed.')
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
export const captionModelFiles = { model: MODEL_NAME, mmproj: MMPROJ_NAME }
|
||||
@@ -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,67 @@ 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.`)
|
||||
}
|
||||
})
|
||||
// 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.')
|
||||
}
|
||||
|
||||
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() })
|
||||
}
|
||||
// Sync one-shot: POST /caption { imagePath, style } → load → caption → unload → { text }
|
||||
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 +948,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 +1057,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.' })
|
||||
}
|
||||
})
|
||||
|
||||
@@ -0,0 +1,129 @@
|
||||
# Download Qwen2.5-VL NSFW Caption V4 GGUF (Q5_K_M + mmproj only) for Studio 2 caption jobs.
|
||||
# Does NOT pull the full safetensors repo (~16.6GB).
|
||||
#
|
||||
# Example:
|
||||
# powershell -ExecutionPolicy Bypass -File scripts\setup-caption.ps1
|
||||
# powershell -File scripts\setup-caption.ps1 -SkipDownload
|
||||
|
||||
[CmdletBinding()]
|
||||
param(
|
||||
[string]$ModelsRoot = '',
|
||||
[string]$LlamaServer = '',
|
||||
[switch]$SkipDownload,
|
||||
[switch]$Force
|
||||
)
|
||||
|
||||
$ErrorActionPreference = 'Stop'
|
||||
$ProgressPreference = 'Continue'
|
||||
|
||||
function Write-Step([string]$Message) {
|
||||
Write-Host ""
|
||||
Write-Host "==> $Message" -ForegroundColor Cyan
|
||||
}
|
||||
|
||||
function Write-Ok([string]$Message) {
|
||||
Write-Host " $Message" -ForegroundColor Green
|
||||
}
|
||||
|
||||
function Write-Warn([string]$Message) {
|
||||
Write-Host " warn $Message" -ForegroundColor Yellow
|
||||
}
|
||||
|
||||
function Resolve-ModelsRoot([string]$Hint) {
|
||||
$shared = Join-Path $env:LOCALAPPDATA 'Comfy-Desktop\ComfyUI-Shared\models'
|
||||
$candidates = @(
|
||||
$Hint,
|
||||
$env:COMFY_MODELS_ROOT,
|
||||
$shared
|
||||
) | Where-Object { $_ } | ForEach-Object { $_.TrimEnd('\', '/') }
|
||||
|
||||
foreach ($path in $candidates) {
|
||||
if (Test-Path $path) { return (Resolve-Path $path).Path }
|
||||
}
|
||||
New-Item -ItemType Directory -Force -Path $shared | Out-Null
|
||||
return (Resolve-Path $shared).Path
|
||||
}
|
||||
|
||||
function Find-HfCli {
|
||||
foreach ($name in @('hf', 'huggingface-cli')) {
|
||||
$cmd = Get-Command $name -ErrorAction SilentlyContinue
|
||||
if ($cmd) { return @{ Name = $name; Source = $cmd.Source } }
|
||||
}
|
||||
throw "Hugging Face CLI not found (hf / huggingface-cli). Install with: pip install -U huggingface_hub"
|
||||
}
|
||||
|
||||
function Ensure-Dir([string]$Path) {
|
||||
if (-not (Test-Path $Path)) {
|
||||
New-Item -ItemType Directory -Force -Path $Path | Out-Null
|
||||
}
|
||||
}
|
||||
|
||||
Write-Step 'Resolve models root'
|
||||
$modelsRoot = Resolve-ModelsRoot $ModelsRoot
|
||||
$captionDir = Join-Path $modelsRoot 'caption\qwen25vl-7b-nsfw-v4'
|
||||
Ensure-Dir $captionDir
|
||||
Write-Ok "Models root: $modelsRoot"
|
||||
Write-Ok "Caption dir: $captionDir"
|
||||
|
||||
$model = 'Qwen2.5-VL-7B-NSFW-Caption-V4.Q5_K_M.gguf'
|
||||
$mmproj = 'Qwen2.5-VL-7B-NSFW-Caption-V4.mmproj-f16.gguf'
|
||||
$modelPath = Join-Path $captionDir $model
|
||||
$mmprojPath = Join-Path $captionDir $mmproj
|
||||
|
||||
if (-not $SkipDownload) {
|
||||
Write-Step 'Download GGUF + mmproj only (~7GB)'
|
||||
$needModel = $Force -or -not (Test-Path $modelPath)
|
||||
$needMmproj = $Force -or -not (Test-Path $mmprojPath)
|
||||
if (-not $needModel -and -not $needMmproj) {
|
||||
Write-Ok 'Both weight files already present.'
|
||||
} else {
|
||||
$hf = Find-HfCli
|
||||
if ($hf.Name -eq 'hf') {
|
||||
& $hf.Source download mradermacher/Qwen2.5-VL-7B-NSFW-Caption-V4-GGUF `
|
||||
--include $model `
|
||||
--include $mmproj `
|
||||
--local-dir $captionDir
|
||||
} else {
|
||||
& $hf.Source download mradermacher/Qwen2.5-VL-7B-NSFW-Caption-V4-GGUF `
|
||||
--include $model `
|
||||
--include $mmproj `
|
||||
--local-dir $captionDir
|
||||
}
|
||||
if (-not (Test-Path $modelPath)) { throw "Missing $model after download." }
|
||||
if (-not (Test-Path $mmprojPath)) { throw "Missing $mmproj after download." }
|
||||
Write-Ok 'Download complete.'
|
||||
}
|
||||
} else {
|
||||
Write-Warn 'SkipDownload set — not fetching weights.'
|
||||
}
|
||||
|
||||
Write-Step 'Locate llama-server'
|
||||
$server = $LlamaServer
|
||||
if (-not $server) { $server = $env:CAPTION_LLAMA_SERVER }
|
||||
if (-not $server) {
|
||||
$cmd = Get-Command llama-server -ErrorAction SilentlyContinue
|
||||
if ($cmd) { $server = $cmd.Source }
|
||||
}
|
||||
if ($server) {
|
||||
Write-Ok "llama-server: $server"
|
||||
} else {
|
||||
Write-Warn 'llama-server not on PATH. Install with: winget install ggml.llamacpp'
|
||||
Write-Warn 'Or set CAPTION_LLAMA_SERVER to a CUDA build of llama-server.exe'
|
||||
}
|
||||
|
||||
$ready = @{
|
||||
model = $modelPath
|
||||
mmproj = $mmprojPath
|
||||
modelsDir = $captionDir
|
||||
llamaServer = $server
|
||||
backend = 'llama.cpp'
|
||||
updatedAt = (Get-Date).ToString('o')
|
||||
}
|
||||
$readyPath = Join-Path $captionDir 'aigen-ready.json'
|
||||
$ready | ConvertTo-Json | Set-Content -Path $readyPath -Encoding utf8
|
||||
Write-Ok "Wrote $readyPath"
|
||||
Write-Host ""
|
||||
Write-Host "Caption host looks for:" -ForegroundColor Cyan
|
||||
Write-Host " $modelPath"
|
||||
Write-Host " $mmprojPath"
|
||||
Write-Host "Restart the Comfy host agent after install so /caption/status reports configured=true."
|
||||
Reference in New Issue
Block a user