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.elapsed }}
+ {{ 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')