From 15ef891b2651edb7460d308fa07521677806b017 Mon Sep 17 00:00:00 2001 From: Towsty Date: Wed, 30 Sep 2026 22:50:19 -0500 Subject: [PATCH] Add image-to-text Describe on the bench via exclusive llama.cpp caption jobs. --- components/studio2/Bench.vue | 78 +++++++- scripts/caption-host.mjs | 299 ++++++++++++++++++++++++++++ scripts/comfy-host-agent.mjs | 77 +++++-- scripts/setup-caption.ps1 | 54 +++++ scripts/start-comfy-host-agent.ps1 | 7 + server/api/generate/caption.post.ts | 50 +++++ server/api/health.get.ts | 7 +- server/api/studio-queue.get.ts | 3 +- server/plugins/00-resume-caption.ts | 3 + server/utils/caption.ts | 232 +++++++++++++++++++++ server/utils/jobs.ts | 5 +- server/utils/studioQueue.ts | 70 ++++++- shared/caption.mjs | 30 +++ shared/studio2/status.mjs | 5 +- tests/caption.test.mjs | 52 +++++ 15 files changed, 937 insertions(+), 35 deletions(-) create mode 100644 scripts/caption-host.mjs create mode 100644 scripts/setup-caption.ps1 create mode 100644 server/api/generate/caption.post.ts create mode 100644 server/plugins/00-resume-caption.ts create mode 100644 server/utils/caption.ts create mode 100644 shared/caption.mjs create mode 100644 tests/caption.test.mjs diff --git a/components/studio2/Bench.vue b/components/studio2/Bench.vue index 32406c6..bc1299a 100644 --- a/components/studio2/Bench.vue +++ b/components/studio2/Bench.vue @@ -234,6 +234,18 @@ > + +
+

{{ card.resultText }}

+
+ + +

No active jobs

@@ -460,6 +477,7 @@ import {aspectOptions,resolveSize,bindClip,restoreAll,useAsSource,hydrateStudio2 import {folderAssets,newestOutput,assetLabel,reconcileSelection,previewHidden} from '~/shared/studio2/viewer.mjs' import {queueJobs,jobStageLabel,jobProgressView,elapsedSeconds} from '~/shared/studio2/queue-view.mjs' import {gpuStatus} from '~/shared/studio2/status.mjs' +import {CAPTION_STYLES} from '~/shared/caption.mjs' import {applyUploadedStill} from '~/shared/studio2/upload.mjs' import {defaultNegatives,parseBatch} from '~/shared/studio2/styles.mjs' import {randomSeed} from '~/shared/studio2/seed.mjs' @@ -614,6 +632,9 @@ const promptArea=ref(null) const promptCaret=ref({start:0,end:0}) const music=reactive({tags:DEFAULT_MUSIC_TAGS,lyrics:DEFAULT_MUSIC_LYRICS,duration:60,seed:'',name:''}) const tagsPlaceholder=DEFAULT_MUSIC_TAGS,lyricsPlaceholder=DEFAULT_MUSIC_LYRICS +const captionStyles=CAPTION_STYLES +const captionStyle=ref('descriptive') +const liveCaption=ref>({}) let loraMemory=createLoraMemory(xaigen) const currentLoraUse=computed(()=>loraUse(form)) const pickerLoras=computed(()=>!xaigen && currentLoraUse.value==='ltx'?[]:filterLorasForUse(filterLorasForStudio([...loraNames.value.image,...loraNames.value.video],xaigen),currentLoraUse.value)) @@ -725,7 +746,10 @@ const benchJobs=computed(()=>queueJobs(folderJobs.value)) const musicJobs=computed(()=>studioQueue.value .filter((j:any)=>j.kind==='music' && j.folderId===form.folderId && !['complete','done'].includes(j.status)) .map((j:any)=>({...j,live:j.liveJobId?liveMusic.value[j.liveJobId]:null}))) -const gpu=computed(()=>gpuStatus(jobs.value,healthState.value,upscales.value,musicJobs.value)) +const captionJobs=computed(()=>studioQueue.value + .filter((j:any)=>j.kind==='caption' && j.folderId===form.folderId) + .map((j:any)=>({...j,live:j.liveJobId?liveCaption.value[j.liveJobId]:null}))) +const gpu=computed(()=>gpuStatus(jobs.value,healthState.value,upscales.value,musicJobs.value,captionJobs.value)) const gpuChipLabel=computed(()=>{ if(!healthState.value) return 'GPU ?' const raw=String(healthState.value?.comfy?.gpu || healthState.value?.gpuName || healthState.value?.comfy?.device || '').trim() @@ -792,6 +816,22 @@ const queueCards=computed(()=>{ cancelable:!['complete','done','error','cancelled'].includes(job.status) }) } + for(const job of captionJobs.value){ + const live=job.live + const text=live?.resultText || job.resultText || '' + const state=live?.status || job.status + cards.push({ + id:job.id, + job, + code:'#'+String(job.id).replace(/\D/g,'').slice(-3).padStart(3,'0'), + title:job.name || `Describe · ${job.captionStyle || job.payload?.captionStyle || 'caption'}`, + meta:`Caption · ${job.captionStyle || job.payload?.captionStyle || 'descriptive'}`, + state, + statusLabel:live?.message || state, + resultText:text, + cancelable:!['complete','done','error','cancelled'].includes(job.status) && !['complete','error','cancelled'].includes(String(state)) + }) + } for(const job of upscales.value.filter((j:any)=>j.folderId===form.folderId && ['queued','running'].includes(j.status))){ cards.push({ id:job.id, @@ -807,7 +847,7 @@ const queueCards=computed(()=>{ return cards }) const queueCount=computed(()=>queueCards.value.length) -const canClearFinished=computed(()=>folderJobs.value.some(j=>['complete','failed','cancelled'].includes(j.state)) || musicJobs.value.some(j=>['complete','done','error','cancelled'].includes(j.status))) +const canClearFinished=computed(()=>folderJobs.value.some(j=>['complete','failed','cancelled'].includes(j.state)) || musicJobs.value.some(j=>['complete','done','error','cancelled'].includes(j.status)) || captionJobs.value.some(j=>['complete','done','error','cancelled'].includes(j.status) || ['complete','error','cancelled'].includes(j.live?.status))) const transportLabel=computed(()=>{ if(selected.value?.kind==='clip') return `${formatTime(videoTime.value)} / ${formatTime(videoDuration.value)}` if(selected.value?.kind==='track') return 'Audio' @@ -1000,6 +1040,23 @@ function applyUseImage(item?:any){ uploadMessage.value=`Added ${assetLabel(asset)} to Photos.` saveBenchDraft(true) } +async function copyCaption(text:string){ + try{await navigator.clipboard.writeText(text)}catch{error.value='Could not copy caption.'} +} +function useCaptionAsPrompt(text:string){ + const current=String(form.promptSections.action||'').trim() + form.promptSections.action=current?`${current}\n\n${text}`:text +} +async function describeSelected(){ + const id=selected.value?.kind==='still'?selected.value.id:'' + if(!id){error.value='Select a still to describe.';return} + submitting.value=true;error.value='' + try{ + await $fetch('/api/generate/caption',{method:'POST',body:{folderId:form.folderId,stillId:id,sourcePath:id,captionStyle:captionStyle.value}}) + await refreshJobs() + }catch(e:any){error.value=e.data?.statusMessage||e.message} + finally{submitting.value=false} +} function applyUsePrompt(item?:any){ const asset=resolveActionAsset(item) if(!asset){ @@ -1369,6 +1426,14 @@ async function refreshJobs(){ const nextLive:Record={} for(const [id,snap] of liveEntries){if(snap)nextLive[id]=snap} liveMusic.value=nextLive + const captionRows=studioQueue.value.filter((j:any)=>j.kind==='caption' && j.liveJobId) + const captionEntries=await Promise.all(captionRows.map(async(j:any)=>{ + try{return [j.liveJobId,await $fetch(`/api/generate/${j.liveJobId}`)] as const} + catch{return [j.liveJobId,null] as const} + })) + const nextCaption:Record={} + for(const [id,snap] of captionEntries){if(snap)nextCaption[id]=snap} + liveCaption.value=nextCaption ingestNewOutputs(nextJobs) } /** Merge freshly completed outputs into the local catalog — never refetch the whole library. */ @@ -1559,7 +1624,14 @@ async function queue(){ logTiming('done') } } -async function cancel(job:any){await $fetch('/api/studio-2/cancel',{method:'POST',body:{id:job.id}});await refreshJobs()} +async function cancel(job:any){ + if(job?.kind==='music' || job?.kind==='caption'){ + await $fetch(`/api/studio-queue/${job.id}`,{method:'DELETE'}) + }else{ + await $fetch('/api/studio-2/cancel',{method:'POST',body:{id:job.id}}) + } + await refreshJobs() +} async function loadLibraryOnce(){ // Current folder page only — never the yearly Inbox JSON on / diff --git a/scripts/caption-host.mjs b/scripts/caption-host.mjs new file mode 100644 index 0000000..986d0d4 --- /dev/null +++ b/scripts/caption-host.mjs @@ -0,0 +1,299 @@ +import { spawn } from 'node:child_process' +import { createServer } from 'node:net' +import { existsSync, mkdirSync, readFileSync, writeFileSync, renameSync, appendFileSync, readdirSync, copyFileSync, createWriteStream } from 'node:fs' +import { join, resolve, extname } from 'node:path' +import { pipeline } from 'node:stream/promises' +import { Transform } from 'node:stream' +import { captionStylePrompt, CAPTION_STYLES } from '../shared/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')) +} + +/** Prefer winget ggml.llamacpp; never download. Process exit after one caption = keep_alive 0. */ +function resolveLlamaServer(explicit) { + if (explicit) return resolve(explicit) + if (process.env.CAPTION_LLAMA_SERVER) return resolve(process.env.CAPTION_LLAMA_SERVER) + const local = process.env.LOCALAPPDATA || '' + const candidates = [ + join(local, 'Microsoft', 'WinGet', 'Packages', 'ggml.llamacpp_Microsoft.Winget.Source_8wekyb3d8bbwe', 'llama-server.exe'), + join(local, 'llama.cpp', 'llama-server.exe'), + join('C:\\', 'llama.cpp', 'llama-server.exe') + ] + for (const path of candidates) { + if (existsSync(path)) return path + } + 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 { /* 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 unloads VRAM (keep_alive 0). */ +export function createCaptionHost({ + prepare, + leaseValid, + spawnProcess = spawn, + modelsDir, + llamaServer, + dataDir, + fetchImpl = fetch, + serverWaitMs = 120_000 +} = {}) { + const models = resolve(modelsDir || defaultModelsDir()) + const executable = resolveLlamaServer(llamaServer) + 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 */ } + } + } + 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 } + const path = join(dir(id), 'status.json') + return existsSync(path) ? JSON.parse(readFileSync(path, 'utf8')) : null + } + 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') + 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 missing 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 = existsSync(dir(request.id)) && 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() + // One caption per process. killTree after response unloads VRAM (keep_alive 0). + const child = spawnProcess(executable, [ + '-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' + ], { 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, serverWaitMs, 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 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 : AbortSignal.timeout(Number(process.env.CAPTION_TIMEOUT_MS || 180_000)), + 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 } } + ] + }] + }) + }) + 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) + 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) + }, + 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 } diff --git a/scripts/comfy-host-agent.mjs b/scripts/comfy-host-agent.mjs index 5258981..0327020 100644 --- a/scripts/comfy-host-agent.mjs +++ b/scripts/comfy-host-agent.mjs @@ -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.' }) } }) diff --git a/scripts/setup-caption.ps1 b/scripts/setup-caption.ps1 new file mode 100644 index 0000000..383fada --- /dev/null +++ b/scripts/setup-caption.ps1 @@ -0,0 +1,54 @@ +# Caption (Qwen2.5-VL NSFW Caption V4 GGUF) — verify host files only. +# Models are already on the 5080 host. Do not download. Do not search Hugging Face. + +$ErrorActionPreference = 'Stop' + +function Resolve-ModelsRoot { + if ($env:COMFY_MODELS_ROOT) { return $env:COMFY_MODELS_ROOT } + $shared = Join-Path $env:LOCALAPPDATA 'Comfy-Desktop\ComfyUI-Shared\models' + if (Test-Path $shared) { return $shared } + throw 'Set COMFY_MODELS_ROOT or install Comfy Desktop Shared models.' +} + +$root = Join-Path (Resolve-ModelsRoot) 'caption\qwen25vl-7b-nsfw-v4' +if (-not (Test-Path $root)) { + throw "Caption models folder missing: $root" +} + +$files = @( + 'Qwen2.5-VL-7B-NSFW-Caption-V4.Q5_K_M.gguf', + 'Qwen2.5-VL-7B-NSFW-Caption-V4.mmproj-f16.gguf' +) + +foreach ($name in $files) { + $path = Join-Path $root $name + if (-not (Test-Path $path)) { + throw "Missing $name under $root. Do not download; use the files already on this host." + } + Write-Host "OK $name" +} + +$llama = $null +if ($env:CAPTION_LLAMA_SERVER -and (Test-Path $env:CAPTION_LLAMA_SERVER)) { + $llama = Get-Item $env:CAPTION_LLAMA_SERVER +} else { + $winget = Join-Path $env:LOCALAPPDATA 'Microsoft\WinGet\Packages\ggml.llamacpp_Microsoft.Winget.Source_8wekyb3d8bbwe\llama-server.exe' + if (Test-Path $winget) { $llama = Get-Item $winget } + else { $llama = Get-Command llama-server -ErrorAction SilentlyContinue } +} + +if (-not $llama) { + throw 'llama-server not found. Install with: winget install ggml.llamacpp' +} +$llamaPath = if ($llama.PSObject.Properties['FullName'] -and $llama.FullName) { $llama.FullName } elseif ($llama.PSObject.Properties['Source'] -and $llama.Source) { $llama.Source } else { [string]$llama } +Write-Host "llama-server: $llamaPath" + +$ready = @{ + model = 'Qwen2.5-VL-7B-NSFW-Caption-V4.Q5_K_M.gguf' + mmproj = 'Qwen2.5-VL-7B-NSFW-Caption-V4.mmproj-f16.gguf' + modelsDir = $root + llamaServer = $llamaPath + readyAt = (Get-Date).ToUniversalTime().ToString('o') +} | ConvertTo-Json +Set-Content -Path (Join-Path $root 'aigen-ready.json') -Value $ready -Encoding utf8 +Write-Host "Caption ready: $root" diff --git a/scripts/start-comfy-host-agent.ps1 b/scripts/start-comfy-host-agent.ps1 index 7f03178..ccaad88 100644 --- a/scripts/start-comfy-host-agent.ps1 +++ b/scripts/start-comfy-host-agent.ps1 @@ -18,6 +18,12 @@ Get-NetTCPConnection -LocalPort 8199 -State Listen -ErrorAction SilentlyContinue Start-Sleep -Milliseconds 400 +# Caption VLM — point at existing host GGUFs (no download). +$captionModels = Join-Path $env:LOCALAPPDATA 'Comfy-Desktop\ComfyUI-Shared\models\caption\qwen25vl-7b-nsfw-v4' +$llamaServer = Join-Path $env:LOCALAPPDATA 'Microsoft\WinGet\Packages\ggml.llamacpp_Microsoft.Winget.Source_8wekyb3d8bbwe\llama-server.exe' +if (Test-Path $captionModels) { $env:CAPTION_MODELS = $captionModels } +if (Test-Path $llamaServer) { $env:CAPTION_LLAMA_SERVER = $llamaServer } + if ($Wait) { Set-Location $root & $node $agent @@ -31,4 +37,5 @@ $psi.WorkingDirectory = $root $psi.UseShellExecute = $false $psi.CreateNoWindow = $true $psi.WindowStyle = [System.Diagnostics.ProcessWindowStyle]::Hidden +# CAPTION_* already set on this process; child inherits when UseShellExecute is false [void][System.Diagnostics.Process]::Start($psi) diff --git a/server/api/generate/caption.post.ts b/server/api/generate/caption.post.ts new file mode 100644 index 0000000..968a95f --- /dev/null +++ b/server/api/generate/caption.post.ts @@ -0,0 +1,50 @@ +import { validateCaptionBody } from '~/shared/caption.mjs' +import { addStudioJob, kickStudioQueue, type StudioJobPayload } from '~/server/utils/studioQueue' +import { getStill } from '~/server/utils/library' + +export default defineEventHandler(async (event) => { + const { owner } = assertLibraryOwner(event) + const raw = await readBody(event) + let body: ReturnType + try { + body = validateCaptionBody(raw) + } catch (e: any) { + throw createError({ statusCode: e.statusCode || 400, statusMessage: e.message }) + } + assertFolderAccess(event, body.folderId) + const still = getStill(owner, body.stillId) + assertFolderAccess(event, still.folderId) + + const row = await addStudioJob({ + ownerKey: owner, + kind: 'caption', + familyId: crypto.randomUUID(), + payload: { + prompt: `Describe · ${body.captionStyle}`, + name: `Describe · ${body.captionStyle}`, + folderId: body.folderId, + stillId: still.id, + captionStyle: body.captionStyle, + aspect: 'auto', + width: 0, + height: 0, + steps: 1, + turbo: true, + seed: 0, + cfg: 0, + fps: 24, + samplerName: 'euler', + scheduler: 'simple', + duration: 0, + sound: false, + workflow: 'minimax', + useIdentityRefs: false, + hideThumbnail: false, + referenceStillIds: [], + extensions: [], + queueAutoRun: true + } as StudioJobPayload + }) + kickStudioQueue() + return { id: row.id, liveJobId: row.liveJobId, captionStyle: body.captionStyle, stillId: still.id } +}) diff --git a/server/api/health.get.ts b/server/api/health.get.ts index f82db21..d177cca 100644 --- a/server/api/health.get.ts +++ b/server/api/health.get.ts @@ -48,7 +48,7 @@ export default defineEventHandler(async () => { host: comfy.host || '', queueRunning: queue.running, queuePending: queue.pending, - asleep: Boolean(control?.asleep) && !comfyReachable && !control?.yue2?.busy, + asleep: Boolean(control?.asleep) && !comfyReachable && !control?.yue2?.busy && !control?.caption?.busy, processUp: Boolean(control?.process ?? control?.processUp), lastActivityAt: control?.lastActivityAt || null, stats: comfy.ok ? comfy.stats : undefined @@ -58,6 +58,11 @@ export default defineEventHandler(async () => { busy: Boolean(control?.yue2?.busy), backend: control?.yue2?.backend || 'yue2' }, + caption: { + configured: Boolean(control?.caption?.configured), + busy: Boolean(control?.caption?.busy), + backend: control?.caption?.backend || 'llama.cpp' + }, imageComfy: { configured: imageConfigured, ok: imageOk, diff --git a/server/api/studio-queue.get.ts b/server/api/studio-queue.get.ts index 03e30fd..ff60bde 100644 --- a/server/api/studio-queue.get.ts +++ b/server/api/studio-queue.get.ts @@ -10,13 +10,14 @@ export default defineEventHandler((event) => { || job.status === 'running' || job.status === 'held' || (job.status === 'error' && Date.now() - job.updatedAt < 60 * 60 * 1000) + || (job.kind === 'caption' && job.status === 'complete' && Date.now() - job.updatedAt < 60 * 60 * 1000) )) const claimed = new Set(active.map(job => job.shotQueueId).filter(Boolean) as string[]) const queues = listShotQueues(owner) const byId = new Map(queues.map(queue => [queue.id, queue])) const jobs = active.map((job) => { const shots = job.shotQueueId ? byId.get(job.shotQueueId) || null : null - const kind = job.kind === 'edit' ? 'edit' : job.kind === 'music' ? 'music' : 'video' + const kind = job.kind === 'edit' ? 'edit' : job.kind === 'music' ? 'music' : job.kind === 'caption' ? 'caption' : 'video' const shotSummary = shots ? summarizeQueue(shots) : null const plannedShots = shotSummary ? undefined diff --git a/server/plugins/00-resume-caption.ts b/server/plugins/00-resume-caption.ts new file mode 100644 index 0000000..fa70bfd --- /dev/null +++ b/server/plugins/00-resume-caption.ts @@ -0,0 +1,3 @@ +import { resumeCaptionJobs } from '~/server/utils/caption' + +export default defineNitroPlugin(() => { resumeCaptionJobs() }) diff --git a/server/utils/caption.ts b/server/utils/caption.ts new file mode 100644 index 0000000..7b79f1f --- /dev/null +++ b/server/utils/caption.ts @@ -0,0 +1,232 @@ +import { createReadStream, existsSync, mkdirSync, readdirSync, readFileSync, writeFileSync, renameSync, unlinkSync } from 'node:fs' +import { join } from 'node:path' +import { createJob, emitJob, getJob, restoreJob, type Job } from './jobs' +import { getStill, stillPath } from './library' +import { sharedGpuHeaders } from './sharedGpu' +import { CAPTION_STYLES } from '~/shared/caption.mjs' + +function settings() { + const config = useRuntimeConfig() + return { + url: String(config.comfyControlUrl || process.env.COMFY_CONTROL_URL || '').replace(/\/$/, ''), + token: String(config.comfyControlToken || process.env.COMFY_CONTROL_TOKEN || '') + } +} + +export function captionConfigured() { return Boolean(settings().url) } + +function pendingRoot() { + return join(String(useRuntimeConfig().libraryDir || process.env.LIBRARY_DIR || '/data/library'), 'caption-pending') +} + +export function captionPendingAlive(jobId: string) { + return Boolean(jobId) && existsSync(join(pendingRoot(), `${jobId}.json`)) +} + +async function request(path: string, method = 'GET', body?: unknown) { + const { url, token } = settings() + if (!url) throw new Error('Caption host is not configured. Set COMFY_CONTROL_URL.') + const response = await fetch(`${url}/caption/${path}`, { + method, + headers: { + ...(token ? { Authorization: `Bearer ${token}` } : {}), + ...(method === 'GET' ? {} : sharedGpuHeaders()), + ...(method === 'POST' ? { 'Content-Type': 'application/json' } : {}) + }, + body: method === 'POST' ? JSON.stringify(body) : body as BodyInit | undefined, + ...(method === 'PUT' ? { duplex: 'half' as const } : {}), + // Caption loads a 7B VLM; allow the host POST to finish or fall through to watch + signal: AbortSignal.timeout(method === 'PUT' ? 120_000 : method === 'POST' ? 300_000 : 60_000) + } as RequestInit) + if (!response.ok) { + const detail = await response.json().catch(() => ({})) as { message?: string; error?: string } + throw Object.assign(new Error(detail.message || detail.error || `Caption host returned ${response.status}`), { statusCode: response.status }) + } + return response +} + +function persist(record: Record) { + mkdirSync(pendingRoot(), { recursive: true }) + const path = join(pendingRoot(), `${record.id}.json`) + writeFileSync(path + '.tmp', JSON.stringify(record)) + renameSync(path + '.tmp', path) +} + +function clearPending(id: string) { + const path = join(pendingRoot(), `${id}.json`) + if (existsSync(path)) unlinkSync(path) +} + +function writeSidecar(owner: string, stillId: string, text: string) { + try { + const image = stillPath(owner, stillId) + if (!existsSync(image)) return + writeFileSync(`${image}.txt`, text, 'utf8') + } catch { /* optional */ } +} + +export async function cancelCaptionJob(job: Job) { + await request(`jobs/${job.id}/cancel`, 'POST', {}) + job.status = 'cancelled' + emitJob(job, { type: 'error', error: 'Cancelled', message: 'Cancelled' }) +} + +async function settle(job: Job) { + const { onLiveVideoSettled } = await import('./studioQueue') + await onLiveVideoSettled(job) + clearPending(job.id) +} + +async function watch(job: Job, record: { + id: string + owner: string + stillId: string + folderId: string + captionStyle: string + liveId: string +}) { + let failures = 0 + while (job.status !== 'cancelled') { + try { + const state = await (await request(`jobs/${job.id}`)).json() as { + status: string; message?: string; error?: string; progress?: number; text?: string + } + if (job.status === 'cancelled') { await settle(job); return } + if (state.status === 'error' || state.status === 'cancelled') { + job.status = state.status === 'cancelled' ? 'cancelled' : 'error' + job.error = state.error || state.message || 'Caption failed' + emitJob(job, { type: 'error', error: job.error, message: job.error }) + await settle(job) + return + } + if (state.status === 'complete') { + const text = String(state.text || '').trim() + job.message = 'Caption ready' + job.progress = 100 + job.status = 'complete' + job.resultText = text + writeSidecar(record.owner, record.stillId, text) + persist({ ...record, resultText: text, status: 'complete' }) + emitJob(job, { type: 'complete', message: 'Caption ready', progress: 100 }) + await settle(job) + return + } + failures = 0 + job.status = 'running' + emitJob(job, { + type: 'progress', + message: state.message || 'Captioning', + progress: Math.min(95, Math.max(1, Number(state.progress || 10))) + }) + } catch (error) { + failures++ + emitJob(job, { type: 'status', message: `Caption host check failed; retrying: ${error instanceof Error ? error.message : String(error)}` }) + if (failures >= 10) { + job.status = 'error' + job.error = 'Caption host unreachable.' + emitJob(job, { type: 'error', error: job.error, message: job.error }) + await settle(job) + return + } + } + await new Promise(resolve => setTimeout(resolve, 1500)) + } + await settle(job) +} + +export async function startCaptionJob(params: { + ownerKey: string + folderId: string + stillId: string + captionStyle: string + name?: string +}) { + if (!CAPTION_STYLES.includes(params.captionStyle)) throw new Error('Unknown caption style.') + getStill(params.ownerKey, params.stillId) + const job = createJob('caption') + job.caption = true + job.library = { + ownerKey: params.ownerKey, + folderId: params.folderId, + hideThumbnail: false, + name: params.name || `Describe · ${params.captionStyle}`, + prompt: params.captionStyle, + aspect: 'image', + width: 0, + height: 0, + steps: 1, + turbo: true, + seed: 0, + stillId: params.stillId, + engine: 'caption' + } + const record = { + id: job.id, + liveId: job.id, + owner: params.ownerKey, + stillId: params.stillId, + folderId: params.folderId, + captionStyle: params.captionStyle, + createdAt: Date.now(), + status: 'running' + } + persist(record) + void (async () => { + try { + if (job.status === 'cancelled') { await settle(job); return } + await request(`jobs/${job.id}/input`, 'PUT', createReadStream(stillPath(params.ownerKey, params.stillId))) + if (job.status === 'cancelled') { await settle(job); return } + try { + await request('jobs', 'POST', { id: job.id, style: params.captionStyle }) + } catch { /* confirm by ID on timeout */ } + await watch(job, record) + } catch (error) { + const statusCode = (error as { statusCode?: number })?.statusCode + if (statusCode && statusCode >= 400 && statusCode < 500) { + job.status = 'error' + job.error = error instanceof Error ? error.message : String(error) + emitJob(job, { type: 'error', error: job.error, message: job.error }) + await settle(job) + return + } + emitJob(job, { type: 'status', message: `Checking caption submission: ${error instanceof Error ? error.message : String(error)}` }) + await watch(job, record) + } + })() + return job +} + +export function resumeCaptionJobs() { + if (!existsSync(pendingRoot())) return + for (const file of readdirSync(pendingRoot()).filter(file => /^[a-zA-Z0-9-]+\.json$/.test(file))) { + try { + const record = JSON.parse(readFileSync(join(pendingRoot(), file), 'utf8')) + if (getJob(record.id)) continue + const job = restoreJob({ + id: record.liveId || record.id, + clientId: record.clientId || crypto.randomUUID(), + promptId: '', + startedAt: record.createdAt || Date.now(), + library: { + ownerKey: record.owner, + folderId: record.folderId, + hideThumbnail: false, + prompt: record.captionStyle || 'descriptive', + aspect: 'image', + width: 0, + height: 0, + steps: 1, + turbo: true, + seed: 0, + stillId: record.stillId, + engine: 'caption' + } + }) + job.caption = true + job.kind = 'caption' + if (record.resultText) job.resultText = record.resultText + job.message = 'Reconnecting to caption host' + void watch(job, record) + } catch { /* keep bad records */ } + } +} diff --git a/server/utils/jobs.ts b/server/utils/jobs.ts index bd25f78..0b5932d 100644 --- a/server/utils/jobs.ts +++ b/server/utils/jobs.ts @@ -36,9 +36,11 @@ export interface Job { upscale?: boolean yueGp?: boolean yue2?: boolean + caption?: boolean + resultText?: string musicActivity?: { checkedAt: number; running: boolean } id: string - kind?: 'video' | 'edit' | 'music' + kind?: 'video' | 'edit' | 'music' | 'caption' promptId?: string clientId: string status: JobStatus @@ -222,6 +224,7 @@ export function jobSnapshot(job: Job) { clipId: job.clipId, stillId: job.stillId, trackId: job.trackId, + resultText: job.resultText, hideThumbnail: job.hideThumbnail, error: job.error, folderLocked: job.library?.folderLocked, diff --git a/server/utils/studioQueue.ts b/server/utils/studioQueue.ts index ba8214a..90a5817 100644 --- a/server/utils/studioQueue.ts +++ b/server/utils/studioQueue.ts @@ -13,7 +13,7 @@ import { imageV2StackSpecials } from '~/utils/imageV2' import { allowIdentityRefs, type PermanenceRef } from '~/utils/globalLocks' export type StudioJobStatus = 'waiting' | 'running' | 'held' | 'complete' | 'error' | 'cancelled' -export type StudioJobKind = 'video' | 'edit' | 'music' +export type StudioJobKind = 'video' | 'edit' | 'music' | 'caption' export interface StudioJobPayload { studio2Id?: string @@ -42,6 +42,7 @@ export interface StudioJobPayload { useIdentityRefs: boolean stillId?: string stillFilename?: string + captionStyle?: string hideThumbnail: boolean hideInput?: boolean folderLocked?: boolean @@ -105,6 +106,7 @@ export interface StudioJob { resumeAutoRun?: boolean lastError?: string waitReason?: string + resultText?: string } type StudioQueueStore = { @@ -226,6 +228,7 @@ export function listStudioJobs(owner: string) { export function studioJobKind(job: Pick | { kind?: string }) { if (job.kind === 'edit') return 'edit' if (job.kind === 'music') return 'music' + if (job.kind === 'caption') return 'caption' return 'video' } @@ -245,6 +248,8 @@ export function summarizeStudioJob(job: StudioJob) { folderId: job.payload.folderId, musicEngine: job.payload.musicEngine, stillId: job.payload.stillId, + captionStyle: job.payload.captionStyle, + resultText: job.resultText, workflow: job.payload.workflow, imagePipeline: job.payload.imagePipeline || 'v1', duration: job.payload.duration, @@ -290,7 +295,7 @@ function failZombieLiveJob(job: Job, error: string) { function sweepStaleLiveJobs() { const now = Date.now() for (const job of listJobs()) { - if (job.studio2 || job.yueGp || job.yue2 || job.upscale) continue + if (job.studio2 || job.yueGp || job.yue2 || job.upscale || job.caption) continue if (job.saving) continue if (job.status === 'queued' && !job.promptId && now - job.startedAt >= QUEUED_GRACE_MS) { failZombieLiveJob(job, 'Job never started') @@ -310,12 +315,12 @@ async function reapZombieLiveJobs() { const { fetchHistory } = await import('~/server/utils/comfy') for (const job of listJobs()) { // Never interrupt download/stitch/library save — Comfy is idle then by design. - if (job.studio2 || job.yueGp || job.yue2 || job.upscale) continue + if (job.studio2 || job.yueGp || job.yue2 || job.upscale || job.caption) continue if (job.saving) continue if (job.library?.chainContinuing) continue if (job.status !== 'running' && job.status !== 'uploading' && job.status !== 'queued') continue if (jobIsLocallySubmitting(job)) continue - const zombieMs = (job.kind === 'edit' || job.kind === 'music') ? 45_000 : ZOMBIE_EMPTY_COMFY_MS + const zombieMs = (job.kind === 'edit' || job.kind === 'music' || job.kind === 'caption') ? 45_000 : ZOMBIE_EMPTY_COMFY_MS if (jobAgeMs(job) < zombieMs) continue if (!job.promptId) { if (jobAgeMs(job) >= QUEUED_GRACE_MS) failZombieLiveJob(job, 'Job never started on ComfyUI') @@ -355,7 +360,7 @@ async function reapZombieLiveJobs() { } function liveJobOwnsGpu(job: Job) { - if ((job.studio2 || job.yueGp || job.yue2 || job.upscale) && ['running', 'queued', 'uploading'].includes(job.status)) return true + if ((job.studio2 || job.yueGp || job.yue2 || job.upscale || job.caption) && ['running', 'queued', 'uploading'].includes(job.status)) return true if (job.library?.stopAfterCurrent) return false if (job.library?.chainContinuing) return true if (job.saving) return true @@ -375,7 +380,7 @@ function liveJobOwnsGpu(job: Job) { */ function clearDeadGpuClaimsForForceStart() { for (const live of listJobs()) { - if (live.studio2 || live.yueGp || live.yue2 || live.upscale) continue + if (live.studio2 || live.yueGp || live.yue2 || live.upscale || live.caption) continue if (live.saving) continue if (jobIsLocallySubmitting(live)) continue if (live.status !== 'running' && live.status !== 'queued' && live.status !== 'uploading') { @@ -469,8 +474,11 @@ export async function addStudioJob(params: { kind?: StudioJobKind }) { const now = Date.now() - const kind = params.kind === 'edit' ? 'edit' : params.kind === 'music' ? 'music' : 'video' - const shotCount = kind === 'music' + const kind = params.kind === 'edit' ? 'edit' + : params.kind === 'music' ? 'music' + : params.kind === 'caption' ? 'caption' + : 'video' + const shotCount = kind === 'music' || kind === 'caption' ? 1 : kind === 'edit' ? 1 + (params.payload.passes?.length || 0) @@ -606,6 +614,10 @@ export async function clearStuckStudioWork(owner: string) { const { cancelYue2Job } = await import('./yue2') await cancelYue2Job(job) } + if (job.caption) { + const { cancelCaptionJob } = await import('./caption') + await cancelCaptionJob(job) + } job.status = 'cancelled' job.error = 'Cleared by force reset' if (job.library) { @@ -678,6 +690,11 @@ async function stopLiveGeneration(liveJobId?: string, shotQueueId?: string) { await cancelYue2Job(live) return } + if (live?.caption) { + const { cancelCaptionJob } = await import('./caption') + await cancelCaptionJob(live) + return + } if (live) { live.status = 'cancelled' if (live.library) { @@ -889,6 +906,10 @@ function pendingAlive(job: StudioJob) { const root = join(String(useRuntimeConfig().libraryDir || process.env.LIBRARY_DIR || '/data/library'), 'yue2-pending', `${job.liveJobId}.json`) if (existsSync(root)) return true } + if (job.kind === 'caption') { + const root = join(String(useRuntimeConfig().libraryDir || process.env.LIBRARY_DIR || '/data/library'), 'caption-pending', `${job.liveJobId}.json`) + if (existsSync(root)) return true + } } if (!job.shotQueueId) return false return listPendingJobs().some(pending => ( @@ -1397,6 +1418,32 @@ async function startStudioExtendJob(item: StudioJob) { } } +async function startStudioCaptionJob(item: StudioJob) { + let live: import('~/server/utils/jobs').Job | undefined + try { + const { startCaptionJob } = await import('./caption') + const style = String(item.payload.captionStyle || 'descriptive') + const stillId = String(item.payload.stillId || '') + if (!stillId) throw new Error('Caption requires a still.') + live = await startCaptionJob({ + ownerKey: item.ownerKey, + folderId: item.payload.folderId, + stillId, + captionStyle: style, + name: item.payload.name + }) + await markStudioLive(item.ownerKey, item.id, live.id) + } catch (error) { + const message = error instanceof Error ? error.message : String(error) + if (live && (live.status === 'queued' || live.status === 'uploading' || live.status === 'running')) { + live.status = 'error' + live.error = message + } + await parkStudioOnStartFailure(item.ownerKey, item.id, message) + kickStudioQueue() + } +} + async function startStudioMusicJob(item: StudioJob) { let live: import('~/server/utils/jobs').Job | undefined try { @@ -1439,6 +1486,10 @@ export async function startStudioJob(item: StudioJob) { } async function startStudioJobReserved(item: StudioJob) { + if (studioJobKind(item) === 'caption') { + await startStudioCaptionJob(item) + return + } if (item.payload.studio2Id) { try { const { startStudio2Job } = await import('./studio2/runner') @@ -1703,7 +1754,7 @@ export async function onLiveVideoSettled(job: Job) { return } const remaining = remainingStudioShots(job) - const wakeFail = !job.studio2 && !job.yueGp && !job.yue2 && !job.upscale && job.status === 'error' && remaining > 0 && isTransientComfyError(job.error) + const wakeFail = !job.studio2 && !job.yueGp && !job.yue2 && !job.upscale && !job.caption && job.status === 'error' && remaining > 0 && isTransientComfyError(job.error) const failed = (job.status === 'error' || job.status === 'cancelled') && !wakeFail await mutateStore(owner, (store) => { @@ -1777,6 +1828,7 @@ export async function onLiveVideoSettled(job: Job) { : failed ? 'error' : 'complete' + if (job.caption && job.resultText) row.resultText = job.resultText clearStudioRowSlot(row, status, failed ? job.error : undefined) syncPausedFlag(store) }) diff --git a/shared/caption.mjs b/shared/caption.mjs new file mode 100644 index 0000000..3933009 --- /dev/null +++ b/shared/caption.mjs @@ -0,0 +1,30 @@ +export const CAPTION_STYLES = ['descriptive', 'klein_prompt', 'delta', 'tags'] + +export const DELTA_LOCKED_LOOK = + 'short light-brown bob, pink tank, light blue shorts, open green Slytherin robe with red-gold trim, barefoot' + +export function captionStylePrompt(style) { + switch (style) { + case 'klein_prompt': + return 'Write one self-contained image generation prompt for Flux Klein. Include subject, pose, clothing, setting, lighting, and camera. No preamble or quotes.' + case 'delta': + return `The character's locked baseline look is: ${DELTA_LOCKED_LOOK}. Describe ONLY what changed from that baseline in this image. Do not restate the locked look.` + case 'tags': + return 'List comma-separated descriptive tags for this image. Tags only, no sentences.' + case 'descriptive': + default: + return 'Describe this image in clear, detailed prose. Cover subject, appearance, clothing, pose, setting, and lighting.' + } +} + +export function validateCaptionBody(raw) { + const fail = message => { throw Object.assign(new Error(message), { statusCode: 400 }) } + const body = raw || {} + const style = String(body.captionStyle || body.style || 'descriptive') + if (!CAPTION_STYLES.includes(style)) fail('Unknown caption style.') + const folderId = String(body.folderId || '') + if (!folderId) fail('Choose a folder.') + const stillId = String(body.stillId || body.sourcePath || '') + if (!stillId) fail('Choose a still to describe.') + return { folderId, stillId, captionStyle: style } +} diff --git a/shared/studio2/status.mjs b/shared/studio2/status.mjs index dada742..ff11029 100644 --- a/shared/studio2/status.mjs +++ b/shared/studio2/status.mjs @@ -1,4 +1,4 @@ -export function gpuStatus(jobs, health, upscales=[], musicJobs=[]) { +export function gpuStatus(jobs, health, upscales=[], musicJobs=[], captionJobs=[]) { if (!health) return {label:'Status unavailable',tone:'unknown',detail:'Checking GPU connection…'} const upscale=upscales.find(j=>j.status==='running') if(upscale)return {label:'Upscaling',tone:'busy',detail:upscale.name || 'Video upscale'} @@ -7,6 +7,9 @@ export function gpuStatus(jobs, health, upscales=[], musicJobs=[]) { const stage=music.live?.message || music.live?.stage || music.waitReason || 'YuE2' return {label:music.status==='waiting'?'Queued':'Generating',tone:'busy',detail:String(stage)} } + const caption=captionJobs.find(j=>['waiting','running','held'].includes(j.status) || ['queued','running','uploading'].includes(j.live?.status)) + if(caption)return {label:caption.status==='waiting'?'Queued':'Describing',tone:'busy',detail:caption.payload?.captionStyle || caption.live?.message || 'Caption'} + if(health.caption?.busy)return {label:'Describing',tone:'busy',detail:'Caption model is using the GPU'} if(health.yue2?.busy)return {label:'Generating',tone:'busy',detail:'YuE2 is using the GPU'} const job=jobs.find(j=>['waking','enhancing','submitting','rendering','saving'].includes(j.state)) if(job)return {label:{waking:'Waking up',enhancing:'Enhancing prompt…',submitting:'Submitting',rendering:'Generating…',saving:'Saving'}[job.state],tone:'busy',detail:job.request?.promptSections?.action || 'Generation in progress'} diff --git a/tests/caption.test.mjs b/tests/caption.test.mjs new file mode 100644 index 0000000..6e38522 --- /dev/null +++ b/tests/caption.test.mjs @@ -0,0 +1,52 @@ +import assert from 'node:assert/strict' +import { createCaptionHost, validateCaptionHostRequest } from '../scripts/caption-host.mjs' +import { CAPTION_STYLES, captionStylePrompt, DELTA_LOCKED_LOOK, validateCaptionBody } from '../shared/caption.mjs' +import { mkdtempSync, writeFileSync, mkdirSync, existsSync } from 'node:fs' +import { join } from 'node:path' +import { tmpdir } from 'node:os' +import { EventEmitter } from 'node:events' + +assert.deepEqual(CAPTION_STYLES, ['descriptive', 'klein_prompt', 'delta', 'tags']) +assert.match(captionStylePrompt('delta'), new RegExp(DELTA_LOCKED_LOOK.replace(/[.*+?^${}()|[\]\\]/g, '\\$&'))) +assert.equal(validateCaptionBody({ folderId: 'f1', stillId: 's1', captionStyle: 'tags' }).captionStyle, 'tags') +assert.throws(() => validateCaptionBody({ folderId: 'f1' }), /still/) +assert.throws(() => validateCaptionHostRequest({ id: 'short', style: 'descriptive' }), /Invalid/) + +const root = mkdtempSync(join(tmpdir(), 'caption-test-')) +const models = join(root, 'models') +mkdirSync(models) +writeFileSync(join(models, 'Qwen2.5-VL-7B-NSFW-Caption-V4.Q5_K_M.gguf'), 'x') +writeFileSync(join(models, 'Qwen2.5-VL-7B-NSFW-Caption-V4.mmproj-f16.gguf'), 'x') +writeFileSync(join(root, 'llama-server.exe'), 'x') + +let prepared = 0 +const host = createCaptionHost({ + modelsDir: models, + llamaServer: join(root, 'llama-server.exe'), + dataDir: join(root, 'jobs'), + leaseValid: () => true, + prepare: async () => { prepared += 1 }, + serverWaitMs: 800, + spawnProcess: () => { + const child = new EventEmitter() + child.pid = 4242 + child.stdout = new EventEmitter() + child.stderr = new EventEmitter() + child.kill = () => child.emit('close', 0) + setTimeout(() => child.emit('close', 0), 50) + return child + }, + fetchImpl: async () => ({ ok: false, status: 503, text: async () => 'down' }) +}) + +assert.equal(host.configured(), true) +assert.equal(host.busy(), false) + +const id = 'caption-test-abcdefghijkl' +mkdirSync(join(root, 'jobs', id), { recursive: true }) +writeFileSync(join(root, 'jobs', id, 'input.png'), Buffer.alloc(64, 1)) + +await assert.rejects(() => host.start({ id, style: 'descriptive' }, 'lease'), /llama-server|Caption/) +assert.equal(prepared, 1) +assert.ok(existsSync(join(root, 'jobs', id, 'status.json'))) +console.log('caption tests ok')