Add Studio 2 caption jobs via Qwen2.5-VL GGUF on the host agent.
Exclusive llama-server load/unload on the 5080, queued Describe UI with resultText, Copy, and Use as prompt. Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
@@ -11,7 +11,7 @@
|
||||
<option v-for="f in library.folders" :key="f.id" :value="f.id">{{ f.name }}</option>
|
||||
</select>
|
||||
</label>
|
||||
<div v-if="isClassic && form.mode!=='music'" class="flex flex-wrap items-center gap-3 text-sm text-zinc-300" aria-label="Identity locks">
|
||||
<div v-if="isClassic && form.mode!=='music' && form.mode!=='caption'" class="flex flex-wrap items-center gap-3 text-sm text-zinc-300" aria-label="Identity locks">
|
||||
<label class="inline-flex items-center gap-1.5"><input v-model="form.lockFace" type="checkbox" class="accent-amber-400">Lock face</label>
|
||||
<label class="inline-flex items-center gap-1.5"><input v-model="form.lockOutfit" type="checkbox" class="accent-amber-400">Lock outfit</label>
|
||||
</div>
|
||||
@@ -45,7 +45,7 @@
|
||||
>{{ task }}</button>
|
||||
</nav>
|
||||
|
||||
<div v-if="form.mode!=='music'" class="min-w-0">
|
||||
<div v-if="form.mode!=='music' && form.mode!=='caption'" class="min-w-0">
|
||||
<label v-if="isClassic" class="display-body block text-zinc-300">Model
|
||||
<select v-model="form.engine" class="display-control mt-1 w-full max-w-md border border-white/10 bg-zinc-900 px-3 py-2 outline-none ring-amber-300/40 focus:ring-2">
|
||||
<option v-for="opt in engineOptions" :key="opt.value" :value="opt.value" :disabled="opt.disabled">{{ opt.label }}{{ opt.hint ? ' — '+opt.hint : '' }}</option>
|
||||
@@ -68,7 +68,8 @@
|
||||
</button>
|
||||
</div>
|
||||
</div>
|
||||
<p v-else class="display-body text-zinc-400">Model · YuE2</p>
|
||||
<p v-else-if="form.mode==='music'" class="display-body text-zinc-400">Model · YuE2</p>
|
||||
<p v-else class="display-body text-zinc-400">Model · Qwen2.5-VL caption</p>
|
||||
</div>
|
||||
|
||||
<div class="grid min-w-0 gap-6 xl:grid-cols-[minmax(0,3fr)_minmax(24rem,2fr)] 2xl:gap-8">
|
||||
@@ -103,6 +104,7 @@
|
||||
<button v-if="selected.kind!=='track'" class="rounded-full border border-white/20 px-4 py-2 text-sm outline-none ring-amber-300/40 focus-visible:ring-2" @click="restore(selected)">Use as input</button>
|
||||
<button v-if="selected.kind==='still'" class="rounded-full border border-white/20 px-4 py-2 text-sm outline-none ring-amber-300/40 focus-visible:ring-2" @click="Object.assign(form,useAsSource(form,selected))">Use as source still</button>
|
||||
<button v-if="selected.kind==='still'" class="rounded-full border border-white/20 px-4 py-2 text-sm outline-none ring-amber-300/40 focus-visible:ring-2" @click="Object.assign(form,useAsHero(form,selected))">Use as hero</button>
|
||||
<button v-if="selected.kind==='still' || (selected.kind!=='track' && selected.kind!=='clip' && selected.id)" class="rounded-full border border-amber-300/40 px-4 py-2 text-sm text-amber-100 outline-none ring-amber-300/40 focus-visible:ring-2" :disabled="submitting" @click="describeSelected">Describe</button>
|
||||
<button v-if="selected.kind!=='track'" class="rounded-full border border-white/20 px-4 py-2 text-sm outline-none ring-amber-300/40 focus-visible:ring-2" @click="detailsOpen=true">Details</button>
|
||||
<button class="rounded-full border border-red-400/40 px-4 py-2 text-sm text-red-200 outline-none ring-amber-300/40 focus-visible:ring-2" @click="deleteOutput">Delete</button>
|
||||
<VideoUpscaleControl v-if="selected.kind==='clip'" :clip-id="selected.id" queue-only @complete="refresh" />
|
||||
@@ -139,7 +141,7 @@
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<section v-if="form.mode!=='music'" aria-label="Source" class="display-card">
|
||||
<section v-if="form.mode!=='music' && form.mode!=='caption'" aria-label="Source" class="display-card">
|
||||
<div class="mb-3 flex flex-wrap items-start justify-between gap-3">
|
||||
<h2 class="text-xs uppercase tracking-wider text-zinc-400">Source</h2>
|
||||
<div v-if="isStudio" class="flex flex-wrap gap-4 text-sm text-zinc-300" aria-label="Identity locks">
|
||||
@@ -194,6 +196,18 @@
|
||||
<input v-model="music.name" class="mt-1 w-full rounded-xl bg-zinc-900 p-3 outline-none ring-amber-300/40 focus:ring-2" placeholder="Optional track name">
|
||||
</label>
|
||||
</template>
|
||||
<template v-else-if="form.mode==='caption'">
|
||||
<p class="text-sm text-zinc-400">Qwen2.5-VL · image-to-text · exclusive GPU · unloads after one shot</p>
|
||||
<div class="display-card space-y-3 p-3">
|
||||
<Studio2AssetSlot label="Still to describe" :asset="captionStill" :hidden="slotHidden(captionStillId)" upload :disabled="uploading" @pick="picker='captionStillId'" @upload="uploadCaption($event)" />
|
||||
<label class="block text-sm text-zinc-300">Caption style
|
||||
<select v-model="captionStyle" class="mt-1 w-full rounded-xl border border-white/10 bg-zinc-900 p-3 outline-none ring-amber-300/40 focus:ring-2">
|
||||
<option v-for="style in captionStyles" :key="style" :value="style">{{ style }}</option>
|
||||
</select>
|
||||
</label>
|
||||
<p v-if="captionStyle==='delta'" class="text-xs text-zinc-500">Delta describes only changes from the locked look (bob / pink tank / blue shorts / Slytherin robe / barefoot).</p>
|
||||
</div>
|
||||
</template>
|
||||
<template v-else>
|
||||
<label class="block">Prompt<textarea v-model="form.promptSections.action" rows="4" class="mt-1 w-full rounded-xl bg-zinc-900 p-3 outline-none ring-amber-300/40 focus:ring-2" placeholder="Describe the image or the shot." /></label>
|
||||
<Studio2PromptControls v-if="video" v-model="form.promptSections" />
|
||||
@@ -223,6 +237,25 @@
|
||||
<progress class="h-2 w-full accent-amber-300" max="100" :value="job.progress || 0" />
|
||||
<p class="text-xs text-zinc-400">{{ job.message }}</p>
|
||||
</div>
|
||||
<div v-for="job in captionJobs" :key="job.id" class="space-y-2 rounded-xl border border-amber-300/25 bg-zinc-900 p-3 text-sm" aria-live="polite" role="status">
|
||||
<div class="flex justify-between gap-3">
|
||||
<div class="flex min-w-0 items-start gap-3">
|
||||
<img v-if="job.request?.stillId || job.request?.imageAId" :src="stillUrl(job.request.stillId || job.request.imageAId)" class="h-14 w-14 shrink-0 rounded object-cover" alt="Caption source">
|
||||
<div class="min-w-0">
|
||||
<strong class="block truncate">Describe · {{ job.request?.captionStyle || 'caption' }}</strong>
|
||||
<span class="capitalize text-amber-200">{{ job.state==='waiting'?'queued':job.state }}</span>
|
||||
</div>
|
||||
</div>
|
||||
<button v-if="!['complete','failed','cancelled'].includes(job.state)" type="button" class="shrink-0 text-zinc-300" @click="cancel(job)">Cancel</button>
|
||||
</div>
|
||||
<p class="text-xs text-zinc-400">{{ elapsedCaption(job) }}</p>
|
||||
<p v-if="job.resultText" class="whitespace-pre-wrap text-zinc-200">{{ job.resultText }}</p>
|
||||
<div v-if="job.resultText" class="flex flex-wrap gap-2">
|
||||
<button type="button" class="rounded-full border border-white/20 px-3 py-1.5 text-xs outline-none ring-amber-300/40 focus-visible:ring-2" @click="copyCaption(job.resultText)">Copy</button>
|
||||
<button type="button" class="rounded-full border border-amber-300/40 px-3 py-1.5 text-xs text-amber-100 outline-none ring-amber-300/40 focus-visible:ring-2" @click="useCaptionAsPrompt(job.resultText)">Use as prompt</button>
|
||||
</div>
|
||||
<p v-if="job.error" role="alert" class="text-red-300">{{ job.error }}</p>
|
||||
</div>
|
||||
<div v-for="job in musicJobs" :key="job.id" class="space-y-2 rounded-xl border border-amber-300/25 bg-zinc-900 p-3 text-sm" aria-live="polite" role="status">
|
||||
<div class="flex justify-between gap-3">
|
||||
<strong class="truncate">{{ job.name || 'YuE2' }}</strong>
|
||||
@@ -233,7 +266,7 @@
|
||||
<p v-if="job.lastError || job.live?.error" role="alert" class="text-red-300">{{ job.lastError || job.live?.error }}</p>
|
||||
</div>
|
||||
<Studio2JobProgress v-for="job in benchJobs" :key="job.id" :job="job" :now="now" @select="focusJob(job)" @cancel="cancel(job)" />
|
||||
<button v-if="folderJobs.some(j=>['complete','failed','cancelled'].includes(j.state)) || musicJobs.some(j=>['complete','done','error','cancelled'].includes(j.status))" type="button" class="text-sm text-zinc-400 outline-none ring-amber-300/40 focus-visible:ring-2" @click="clearFinished">Clear finished</button>
|
||||
<button v-if="folderJobs.some(j=>['complete','failed','cancelled'].includes(j.state)) || musicJobs.some(j=>['complete','done','error','cancelled'].includes(j.status)) || captionJobs.some(j=>['complete','failed','cancelled'].includes(j.state))" type="button" class="text-sm text-zinc-400 outline-none ring-amber-300/40 focus-visible:ring-2" @click="clearFinished">Clear finished</button>
|
||||
</section>
|
||||
</div>
|
||||
|
||||
@@ -288,9 +321,10 @@
|
||||
import {createLoraMemory,loraUse} from '~/shared/studio2/lora-memory.mjs'
|
||||
import {filterLorasForStudio,filterLorasForUse} from '~/utils/loras'
|
||||
import {compilePrompt,sectionsFrom} from '~/shared/studio2/contracts.mjs'
|
||||
import {CAPTION_STYLES} from '~/shared/studio2/caption.mjs'
|
||||
import {aspectOptions,resolveSize,bindClip,changeSource,restoreAll,useAsHero,useAsSource} from '~/shared/studio2/input.mjs'
|
||||
import {folderAssets,newestOutput,assetLabel,reconcileSelection,previewHidden} from '~/shared/studio2/viewer.mjs'
|
||||
import {queueJobs} from '~/shared/studio2/queue-view.mjs'
|
||||
import {queueJobs,elapsedSeconds} from '~/shared/studio2/queue-view.mjs'
|
||||
import {gpuStatus} from '~/shared/studio2/status.mjs'
|
||||
import {applyUploadedStill} from '~/shared/studio2/upload.mjs'
|
||||
import {parseBatch,stylePrompt} from '~/shared/studio2/styles.mjs'
|
||||
@@ -312,12 +346,15 @@ const liveMusic=ref<Record<string,any>>({})
|
||||
const form=reactive<any>({mode:'generate',engine:'flux',imageStyles:{positive:[],negative:[]},lockFace:true,lockOutfit:true,hideInput:xaigen,folderId:'',identityStillId:'',imageAId:'',imageBId:'',maskId:'',startClipId:'',endStillId:'',guides:[],promptSections:sectionsFrom({}),settings:{aspect:'auto',fps:24,steps:24,cfg:4,seed:null,seedMode:'random',duration:5,turbo:true,loraStack:[]}})
|
||||
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 captionStillId=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))
|
||||
watch(()=>[currentLoraUse.value,form.settings.loraStack],()=>loraMemory.sync(form),{deep:true,flush:'sync'})
|
||||
const route=useRoute()
|
||||
const tasks=['generate','edit','compose','refine','video','extend','music']
|
||||
const tasks=['generate','edit','compose','refine','video','extend','music','caption']
|
||||
const settingsOpen=ref(false),detailsOpen=ref(false),submitting=ref(false),error=ref(''),batchText=ref('')
|
||||
const video=computed(()=>['video','extend'].includes(form.mode))
|
||||
const kreaSticky=computed(()=>['edit','compose','iterate'].includes(form.mode))
|
||||
@@ -353,11 +390,13 @@ const clips=computed(()=>library.value.clips.filter((s:any)=>s.folderId===form.f
|
||||
const compiled=computed(()=>compilePrompt(form.promptSections,video.value,form))
|
||||
const stillUrl=(id:string)=>`/api/library/stills/${id}`
|
||||
const folderJobs=computed(()=>jobs.value.filter(j=>j.request.folderId===form.folderId))
|
||||
const benchJobs=computed(()=>queueJobs(folderJobs.value))
|
||||
const captionJobs=computed(()=>folderJobs.value.filter(j=>j.request?.mode==='caption'))
|
||||
const benchJobs=computed(()=>queueJobs(folderJobs.value.filter(j=>j.request?.mode!=='caption')))
|
||||
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 gpu=computed(()=>gpuStatus(jobs.value,healthState.value,upscales.value,musicJobs.value,captionJobs.value))
|
||||
const captionStill=computed(()=>stills.value.find((s:any)=>s.id===captionStillId.value))
|
||||
const sourceSlot=computed(()=>{const clip=clips.value.find((a:any)=>a.id===form.startClipId);return form.mode==='extend'?(clip?{...clip,kind:'clip'}:null):inputStill.value})
|
||||
const pickerAssets=computed(()=>picker.value==='startClipId'?clips.value.map((a:any)=>({...a,kind:'clip'})):stills.value.map((a:any)=>({...a,kind:'still'})))
|
||||
const inputStill=computed(()=>stills.value.find((s:any)=>s.id===form.imageAId))
|
||||
@@ -366,6 +405,42 @@ const sourceAsset=computed(()=>form.startClipId?clips.value.find((c:any)=>c.id==
|
||||
const sourceSize=computed(()=>measured[sourceAsset.value?.id] || sourceAsset.value)
|
||||
const resolvedSize=computed(()=>resolveSize(form.settings.aspect,sourceSize.value,undefined,form.engine))
|
||||
|
||||
function elapsedCaption(job:any){
|
||||
const seconds=elapsedSeconds(job,now.value)
|
||||
return `${Math.floor(seconds/60)}m ${seconds%60}s ${job.startedAt||job.state==='rendering'||job.state==='saving'?'elapsed':'waiting'}`
|
||||
}
|
||||
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
|
||||
if(form.mode==='caption' || form.mode==='music')form.mode='generate'
|
||||
}
|
||||
async function queueCaption(stillId:string){
|
||||
if(!stillId){error.value='Choose a still to describe.';return false}
|
||||
try{
|
||||
await $fetch('/api/studio-2/caption',{method:'POST',body:{folderId:form.folderId,stillId,sourcePath:stillId,captionStyle:captionStyle.value}})
|
||||
await refresh()
|
||||
return true
|
||||
}catch(e:any){error.value=e.data?.statusMessage||e.message;return false}
|
||||
}
|
||||
function describeSelected(){
|
||||
const id=selected.value?.kind==='still'?selected.value.id:(selected.value?.kind==='clip'? '' : (selected.value?.id || form.imageAId || form.identityStillId))
|
||||
if(!id){error.value='Select a still to describe.';return}
|
||||
captionStillId.value=id
|
||||
submitting.value=true;error.value=''
|
||||
void queueCaption(id).finally(()=>{submitting.value=false})
|
||||
}
|
||||
async function uploadCaption(event:Event){
|
||||
const input=event.target as HTMLInputElement,file=input.files?.[0];if(!file)return
|
||||
uploading.value=true;uploadMessage.value='';error.value=''
|
||||
const body=new FormData();body.append('image',file);body.append('folderId',form.folderId);body.append('hideInput',String(xaigen))
|
||||
try{const response=await $fetch<any>('/api/studio-2/stills',{method:'POST',body});library.value=response;captionStillId.value=response.still.id;measured[response.still.id]={width:response.still.width,height:response.still.height};uploadMessage.value=`Loaded ${response.still.name || file.name} for caption.`}
|
||||
catch(e:any){error.value=e.data?.statusMessage||e.message}
|
||||
finally{uploading.value=false;input.value=''}
|
||||
}
|
||||
|
||||
function isHidden(asset:any){
|
||||
if(!xaigen || !asset?.id)return false
|
||||
if(isClassic.value){
|
||||
@@ -396,6 +471,11 @@ watch(()=>sourceAsset.value?.id,id=>{const asset=sourceAsset.value;if(!import.me
|
||||
|
||||
function chooseTask(task:string){
|
||||
if(task==='music'){form.mode='music';return}
|
||||
if(task==='caption'){
|
||||
form.mode='caption'
|
||||
captionStillId.value=selected.value?.kind==='still'?selected.value.id:form.imageAId||form.identityStillId||captionStillId.value
|
||||
return
|
||||
}
|
||||
if(task==='extend'){
|
||||
form.mode='extend'
|
||||
form.engine=xaigen && form.engine==='ltx'?'ltx':'minimax'
|
||||
@@ -414,7 +494,7 @@ function selectOutput(item:any){selected.value=item;if(form.mode==='extend' && i
|
||||
function restore(item:any){Object.assign(form,restoreAll(form,item,xaigen,loraMemory));form.guides=form.guides||[];batchText.value=''}
|
||||
function focusJob(job:any){selectOutput(job.outputs?.at(-1)||null)}
|
||||
async function loadLoras(fresh=false){try{loraNames.value=await $fetch('/api/studio-2/loras'+(fresh?'?fresh=1':''))}catch{error.value='Could not refresh LoRAs. The last loaded list is still available.'}}
|
||||
function pickAsset(asset:any){const key=picker.value;if(key.startsWith('guide:'))form.guides[Number(key.split(':')[1])].stillId=asset?.id||'';else if(key==='startClipId')chooseClip(asset?.id||'');else form[key]=asset?.id||'';picker.value=''}
|
||||
function pickAsset(asset:any){const key=picker.value;if(key.startsWith('guide:'))form.guides[Number(key.split(':')[1])].stillId=asset?.id||'';else if(key==='startClipId')chooseClip(asset?.id||'');else if(key==='captionStillId')captionStillId.value=asset?.id||'';else form[key]=asset?.id||'';picker.value=''}
|
||||
async function clearFinished(){await $fetch('/api/studio-2/clear-finished',{method:'POST'});await refresh()}
|
||||
async function deleteOutput(){
|
||||
if(!selected.value)return
|
||||
@@ -432,6 +512,10 @@ async function queue(){
|
||||
await refresh()
|
||||
return
|
||||
}
|
||||
if(form.mode==='caption'){
|
||||
await queueCaption(captionStillId.value || form.imageAId || form.identityStillId)
|
||||
return
|
||||
}
|
||||
if(form.settings.seedMode!=='locked')form.settings.seed=randomSeed()
|
||||
const queued=await $fetch<any>('/api/studio-2/jobs',{method:'POST',body:{...form,mode:form.mode==='edit' && parseBatch(batchText.value).length?'iterate':form.mode,hideInput:xaigen,hideOutput:false,promptSections:video.value?form.promptSections:{action:form.promptSections.action},settings:{...form.settings,...resolvedSize.value},batch:video.value || form.mode==='edit'?parseBatch(batchText.value):[]}})
|
||||
form.settings.seed=queued.seed;form.settings.width=queued.width;form.settings.height=queued.height
|
||||
@@ -486,6 +570,7 @@ onMounted(async()=>{
|
||||
const asset=[...library.value.stills,...library.value.clips,...(library.value.tracks||[])].find((a:any)=>a.id===id)
|
||||
if(asset){selected.value={...asset,kind:route.query.clip?'clip':route.query.track?'track':'still'};if(route.query.use==='1'&&selected.value.kind!=='track')restore(selected.value);if(xaigen)toggleReveal(asset.id)}
|
||||
if(route.query.task==='music')form.mode='music'
|
||||
if(route.query.task==='caption')form.mode='caption'
|
||||
const poll=async()=>{try{await refresh();const health=await $fetch<any>('/api/health');healthState.value=health}catch(e:any){healthState.value=null;error.value=e.data?.statusMessage||e.message}if(!disposed)timer=setTimeout(poll,3000)};void poll()
|
||||
})
|
||||
onBeforeUnmount(()=>{disposed=true;clearTimeout(timer);clearInterval(clock)})
|
||||
|
||||
@@ -0,0 +1,40 @@
|
||||
# Studio 2 image→text caption (Qwen2.5-VL NSFW Caption V4 GGUF)
|
||||
|
||||
Caption jobs run on the Windows GPU host via **llama-server** (llama.cpp vision / mmproj). Not Comfy, not Klein, not JoyCaption.
|
||||
|
||||
## Weights (~7GB — two files only)
|
||||
|
||||
```powershell
|
||||
powershell -ExecutionPolicy Bypass -File scripts\setup-caption.ps1
|
||||
```
|
||||
|
||||
Installs into the Shared models tree:
|
||||
|
||||
`%LOCALAPPDATA%\Comfy-Desktop\ComfyUI-Shared\models\caption\qwen25vl-7b-nsfw-v4\`
|
||||
|
||||
| File | Role |
|
||||
| --- | --- |
|
||||
| `Qwen2.5-VL-7B-NSFW-Caption-V4.Q5_K_M.gguf` | Language model |
|
||||
| `Qwen2.5-VL-7B-NSFW-Caption-V4.mmproj-f16.gguf` | Vision projector |
|
||||
|
||||
Do **not** download the full 16.6GB safetensors repo.
|
||||
|
||||
Runtime needs ~8–10GB VRAM. Exclusive GPU: host stops Comfy (and waits if YuE2/upscale is busy), loads the GGUF for one shot, then kills llama-server (`keep_alive 0`).
|
||||
|
||||
## Host agent
|
||||
|
||||
Requires `llama-server` on PATH (`winget install ggml.llamacpp`) or `CAPTION_LLAMA_SERVER`.
|
||||
|
||||
Restart the Comfy host agent after setup. Confirm:
|
||||
|
||||
```text
|
||||
GET http://127.0.0.1:8199/caption/status → { configured, busy, backend: "llama.cpp" }
|
||||
POST /caption { imagePath, style } → load → { text } → unload
|
||||
POST /caption/jobs + PUT …/input → async job used by Studio 2
|
||||
```
|
||||
|
||||
Styles: `descriptive` | `klein_prompt` | `delta` | `tags`.
|
||||
|
||||
## App
|
||||
|
||||
`POST /api/studio-2/caption` with `{ folderId, stillId|sourcePath, captionStyle }` queues a Studio job. Bench **Describe** button and the **caption** task pill both enqueue. Result stores `resultText` on the studio2 job and optionally writes `{still}.txt` beside the library file.
|
||||
@@ -0,0 +1,331 @@
|
||||
import { spawn } from 'node:child_process'
|
||||
import { createServer } from 'node:net'
|
||||
import { existsSync, mkdirSync, readFileSync, writeFileSync, renameSync, appendFileSync, readdirSync, copyFileSync } from 'node:fs'
|
||||
import { join, resolve, extname } from 'node:path'
|
||||
import { captionStylePrompt, CAPTION_STYLES } from '../shared/studio2/caption.mjs'
|
||||
|
||||
const MODEL_NAME = 'Qwen2.5-VL-7B-NSFW-Caption-V4.Q5_K_M.gguf'
|
||||
const MMPROJ_NAME = 'Qwen2.5-VL-7B-NSFW-Caption-V4.mmproj-f16.gguf'
|
||||
|
||||
function defaultModelsDir() {
|
||||
const shared = process.env.COMFY_MODELS_ROOT
|
||||
|| join(process.env.LOCALAPPDATA || '', 'Comfy-Desktop', 'ComfyUI-Shared', 'models')
|
||||
return resolve(process.env.CAPTION_MODELS || join(shared, 'caption', 'qwen25vl-7b-nsfw-v4'))
|
||||
}
|
||||
|
||||
function defaultLlamaServer() {
|
||||
if (process.env.CAPTION_LLAMA_SERVER) return resolve(process.env.CAPTION_LLAMA_SERVER)
|
||||
return 'llama-server'
|
||||
}
|
||||
|
||||
function freePort() {
|
||||
return new Promise((resolvePort, reject) => {
|
||||
const server = createServer()
|
||||
server.listen(0, '127.0.0.1', () => {
|
||||
const address = server.address()
|
||||
const port = typeof address === 'object' && address ? address.port : 0
|
||||
server.close(error => error ? reject(error) : resolvePort(port))
|
||||
})
|
||||
server.on('error', reject)
|
||||
})
|
||||
}
|
||||
|
||||
async function waitForServer(port, signal, timeoutMs = 120_000, fetchImpl = fetch) {
|
||||
const started = Date.now()
|
||||
while (Date.now() - started < timeoutMs) {
|
||||
if (signal?.aborted) throw new Error('Caption cancelled while loading the model.')
|
||||
try {
|
||||
const response = await fetchImpl(`http://127.0.0.1:${port}/health`, { signal: AbortSignal.timeout(2000) })
|
||||
if (response.ok) return
|
||||
} catch { /* still booting */ }
|
||||
await new Promise(r => setTimeout(r, 500))
|
||||
}
|
||||
throw new Error('llama-server did not become ready in time.')
|
||||
}
|
||||
|
||||
function mimeFor(path) {
|
||||
const ext = extname(path).toLowerCase()
|
||||
if (ext === '.jpg' || ext === '.jpeg') return 'image/jpeg'
|
||||
if (ext === '.webp') return 'image/webp'
|
||||
if (ext === '.gif') return 'image/gif'
|
||||
return 'image/png'
|
||||
}
|
||||
|
||||
function killTree(child) {
|
||||
return new Promise(resolveKill => {
|
||||
if (!child?.pid) return resolveKill()
|
||||
let finished = false
|
||||
const done = () => {
|
||||
if (finished) return
|
||||
finished = true
|
||||
resolveKill()
|
||||
}
|
||||
child.once('close', done)
|
||||
try {
|
||||
if (typeof child.kill === 'function') child.kill()
|
||||
} catch { /* ignore */ }
|
||||
if (process.platform === 'win32' && child.pid > 0) {
|
||||
try {
|
||||
const killer = spawn('taskkill', ['/PID', String(child.pid), '/T', '/F'], { windowsHide: true, stdio: 'ignore' })
|
||||
killer.once('close', done)
|
||||
killer.once('error', done)
|
||||
} catch {
|
||||
done()
|
||||
}
|
||||
}
|
||||
setTimeout(done, 3000)
|
||||
})
|
||||
}
|
||||
|
||||
export function validateCaptionHostRequest(body) {
|
||||
if (!/^[a-zA-Z0-9-]{12,80}$/.test(body?.id || '')) throw new Error('Invalid caption job ID.')
|
||||
const style = String(body.style || body.captionStyle || 'descriptive')
|
||||
if (!CAPTION_STYLES.includes(style)) throw new Error('Unknown caption style.')
|
||||
return { id: body.id, style, imagePath: body.imagePath ? String(body.imagePath) : '' }
|
||||
}
|
||||
|
||||
/** One llama-server process per caption. Process exit is the VRAM unload boundary (keep_alive 0). */
|
||||
export function createCaptionHost({
|
||||
prepare,
|
||||
leaseValid,
|
||||
spawnProcess = spawn,
|
||||
modelsDir,
|
||||
llamaServer,
|
||||
dataDir,
|
||||
fetchImpl = fetch
|
||||
} = {}) {
|
||||
const models = resolve(modelsDir || defaultModelsDir())
|
||||
const executable = llamaServer || defaultLlamaServer()
|
||||
const data = resolve(dataDir || process.env.CAPTION_JOBS_DIR || join(models, 'aigen-jobs'))
|
||||
let active = null
|
||||
let holdUntil = 0
|
||||
if (existsSync(data)) {
|
||||
for (const id of readdirSync(data)) {
|
||||
try {
|
||||
const path = join(data, id, 'status.json')
|
||||
if (!existsSync(path)) continue
|
||||
const state = JSON.parse(readFileSync(path, 'utf8'))
|
||||
if (['running', 'starting', 'cancelling'].includes(state.status)) {
|
||||
holdUntil = Date.now() + 10000
|
||||
state.status = 'error'
|
||||
state.error = 'Caption host restarted. GPU model unloaded.'
|
||||
writeFileSync(path + '.tmp', JSON.stringify(state))
|
||||
renameSync(path + '.tmp', path)
|
||||
}
|
||||
} catch { /* ignore bad records */ }
|
||||
}
|
||||
}
|
||||
const dir = id => {
|
||||
if (!/^[a-zA-Z0-9-]{12,80}$/.test(id || '')) throw new Error('Invalid caption job ID.')
|
||||
return join(data, id)
|
||||
}
|
||||
const persist = job => {
|
||||
const target = join(dir(job.id), 'status.json')
|
||||
writeFileSync(target + '.tmp', JSON.stringify(job))
|
||||
renameSync(target + '.tmp', target)
|
||||
}
|
||||
const read = id => {
|
||||
if (active?.job.id === id) return { ...active.job, text: active.job.text }
|
||||
const path = join(dir(id), 'status.json')
|
||||
if (!existsSync(path)) return null
|
||||
return JSON.parse(readFileSync(path, 'utf8'))
|
||||
}
|
||||
const modelPath = () => join(models, MODEL_NAME)
|
||||
const mmprojPath = () => join(models, MMPROJ_NAME)
|
||||
const configured = () => {
|
||||
try {
|
||||
if (!existsSync(modelPath()) || !existsSync(mmprojPath())) return false
|
||||
if (executable.includes('\\') || executable.includes('/')) return existsSync(executable)
|
||||
return true
|
||||
} catch { return false }
|
||||
}
|
||||
|
||||
return {
|
||||
busy: () => Boolean(active) || Date.now() < holdUntil,
|
||||
configured,
|
||||
modelsDir: () => models,
|
||||
read,
|
||||
async upload(id, stream) {
|
||||
if (read(id)?.status === 'complete') throw new Error('This caption job has already finished.')
|
||||
const jobDir = dir(id)
|
||||
mkdirSync(jobDir, { recursive: true })
|
||||
const dest = join(jobDir, 'input.upload')
|
||||
const { createWriteStream } = await import('node:fs')
|
||||
const { pipeline } = await import('node:stream/promises')
|
||||
const { Transform } = await import('node:stream')
|
||||
let bytes = 0
|
||||
await pipeline(stream, new Transform({
|
||||
transform(chunk, _, callback) {
|
||||
bytes += chunk.length
|
||||
callback(bytes > 40 * 1024 * 1024 ? new Error('Image exceeds the 40 MB upload limit.') : null, chunk)
|
||||
}
|
||||
}), createWriteStream(dest, { flags: 'w' }))
|
||||
if (bytes < 32) throw new Error('Image file is empty.')
|
||||
renameSync(dest, join(jobDir, 'input.png'))
|
||||
},
|
||||
async start(body, lease) {
|
||||
const request = validateCaptionHostRequest(body)
|
||||
const previous = read(request.id)
|
||||
if (previous) return previous
|
||||
if (active || Date.now() < holdUntil) throw Object.assign(new Error('Caption is already running or releasing VRAM.'), { statusCode: 409 })
|
||||
if (!configured()) throw new Error(`Caption model is not installed. Expected ${MODEL_NAME} and ${MMPROJ_NAME} under ${models}. Run scripts/setup-caption.ps1.`)
|
||||
if (!leaseValid(lease)) throw new Error('GPU reservation expired.')
|
||||
|
||||
let imagePath = request.imagePath
|
||||
if (imagePath) {
|
||||
if (!existsSync(imagePath)) throw new Error('imagePath does not exist on the GPU host.')
|
||||
mkdirSync(dir(request.id), { recursive: true })
|
||||
const dest = join(dir(request.id), `input${extname(imagePath) || '.png'}`)
|
||||
copyFileSync(imagePath, dest)
|
||||
imagePath = dest
|
||||
} else {
|
||||
imagePath = join(dir(request.id), 'input.png')
|
||||
if (!existsSync(imagePath)) {
|
||||
const alt = readdirSync(dir(request.id)).find(name => /^input\./i.test(name))
|
||||
if (!alt) throw new Error('Upload an image first or pass imagePath.')
|
||||
imagePath = join(dir(request.id), alt)
|
||||
}
|
||||
}
|
||||
|
||||
const job = {
|
||||
id: request.id,
|
||||
status: 'starting',
|
||||
message: 'Preparing GPU for caption',
|
||||
progress: 0,
|
||||
style: request.style,
|
||||
startedAt: Date.now(),
|
||||
checkedAt: Date.now()
|
||||
}
|
||||
active = { job, child: null, cancelled: false, abort: new AbortController() }
|
||||
const run = active
|
||||
mkdirSync(dir(job.id), { recursive: true })
|
||||
writeFileSync(join(dir(job.id), 'request.json'), JSON.stringify({ ...request, imagePath, prompt: captionStylePrompt(request.style) }))
|
||||
persist(job)
|
||||
|
||||
try {
|
||||
if (prepare) await prepare()
|
||||
if (run.cancelled || !leaseValid(lease)) throw new Error('Caption start cancelled or GPU reservation expired.')
|
||||
|
||||
job.status = 'running'
|
||||
job.message = 'Loading Qwen2.5-VL caption model'
|
||||
job.progress = 5
|
||||
persist(job)
|
||||
|
||||
const port = await freePort()
|
||||
const args = [
|
||||
'-m', modelPath(),
|
||||
'--mmproj', mmprojPath(),
|
||||
'--host', '127.0.0.1',
|
||||
'--port', String(port),
|
||||
'-ngl', String(process.env.CAPTION_NGL || '99'),
|
||||
'-c', String(process.env.CAPTION_CTX || '4096'),
|
||||
'--jinja'
|
||||
]
|
||||
const child = spawnProcess(executable, args, {
|
||||
windowsHide: true,
|
||||
shell: false,
|
||||
stdio: ['ignore', 'pipe', 'pipe'],
|
||||
env: { ...process.env }
|
||||
})
|
||||
run.child = child
|
||||
const log = chunk => appendFileSync(join(dir(job.id), 'worker.log'), chunk)
|
||||
child.stdout.on('data', log)
|
||||
child.stderr.on('data', log)
|
||||
|
||||
const watchdog = setInterval(() => {
|
||||
if (!leaseValid(lease)) {
|
||||
job.error = 'GPU reservation expired; caption stopped.'
|
||||
run.abort.abort()
|
||||
void killTree(child)
|
||||
}
|
||||
}, 5000)
|
||||
|
||||
try {
|
||||
await waitForServer(port, run.abort.signal, 120_000, fetchImpl)
|
||||
if (run.cancelled) throw new Error('Cancelled')
|
||||
job.message = 'Captioning'
|
||||
job.progress = 40
|
||||
persist(job)
|
||||
|
||||
const bytes = readFileSync(imagePath)
|
||||
const dataUrl = `data:${mimeFor(imagePath)};base64,${bytes.toString('base64')}`
|
||||
const timeout = AbortSignal.timeout(Number(process.env.CAPTION_TIMEOUT_MS || 180_000))
|
||||
const onAbort = () => { /* combined via race below */ }
|
||||
run.abort.signal.addEventListener('abort', onAbort, { once: true })
|
||||
let response
|
||||
try {
|
||||
response = await fetchImpl(`http://127.0.0.1:${port}/v1/chat/completions`, {
|
||||
method: 'POST',
|
||||
headers: { 'Content-Type': 'application/json' },
|
||||
signal: run.abort.signal.aborted ? run.abort.signal : timeout,
|
||||
body: JSON.stringify({
|
||||
temperature: 0.2,
|
||||
max_tokens: 512,
|
||||
messages: [{
|
||||
role: 'user',
|
||||
content: [
|
||||
{ type: 'text', text: captionStylePrompt(request.style) },
|
||||
{ type: 'image_url', image_url: { url: dataUrl } }
|
||||
]
|
||||
}]
|
||||
})
|
||||
})
|
||||
} finally {
|
||||
run.abort.signal.removeEventListener('abort', onAbort)
|
||||
}
|
||||
if (run.abort.signal.aborted) throw new Error('Cancelled')
|
||||
if (!response.ok) {
|
||||
const detail = await response.text().catch(() => '')
|
||||
throw new Error(`llama-server caption failed (${response.status}): ${detail.slice(0, 400)}`)
|
||||
}
|
||||
const payload = await response.json()
|
||||
const text = String(payload?.choices?.[0]?.message?.content || '').trim()
|
||||
if (!text) throw new Error('Caption model returned empty text.')
|
||||
job.text = text
|
||||
job.status = run.cancelled ? 'cancelled' : 'complete'
|
||||
job.message = job.status === 'complete' ? 'Caption ready' : 'Cancelled'
|
||||
job.progress = 100
|
||||
job.checkedAt = Date.now()
|
||||
persist(job)
|
||||
return { ...job }
|
||||
} finally {
|
||||
clearInterval(watchdog)
|
||||
await killTree(child)
|
||||
// keep_alive 0: process exit unloads GGUF+mmproj from VRAM
|
||||
holdUntil = Date.now() + 2000
|
||||
if (active === run) active = null
|
||||
}
|
||||
} catch (error) {
|
||||
job.status = run.cancelled ? 'cancelled' : 'error'
|
||||
job.error = error.message || String(error)
|
||||
job.message = job.error
|
||||
job.checkedAt = Date.now()
|
||||
persist(job)
|
||||
if (run.child) await killTree(run.child)
|
||||
if (active === run) active = null
|
||||
holdUntil = Date.now() + 2000
|
||||
throw error
|
||||
}
|
||||
},
|
||||
async cancel(id) {
|
||||
if (active?.job.id !== id) return read(id)
|
||||
const run = active
|
||||
run.cancelled = true
|
||||
run.job.status = 'cancelling'
|
||||
run.job.message = 'Cancelling'
|
||||
persist(run.job)
|
||||
run.abort.abort()
|
||||
if (run.child) await killTree(run.child)
|
||||
return read(id)
|
||||
},
|
||||
/** Sync one-shot used by POST /caption { imagePath, style }. */
|
||||
async captionOnce(body, lease) {
|
||||
const id = body.id || `caption-${Date.now()}-${Math.random().toString(36).slice(2, 10)}`
|
||||
const result = await this.start({ ...body, id }, lease)
|
||||
if (result.status === 'complete') return { text: result.text, id: result.id, style: result.style }
|
||||
throw new Error(result.error || 'Caption failed.')
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
export const captionModelFiles = { model: MODEL_NAME, mmproj: MMPROJ_NAME }
|
||||
@@ -4,6 +4,7 @@ import { stableMemoryArgs } from './comfy-memory-policy.mjs'
|
||||
import { createGpuReservation } from './gpu-reservation.mjs'
|
||||
import { createGpuProxy } from './gpu-proxy.mjs'
|
||||
import { createYue2Host } from './yue2-host.mjs'
|
||||
import { createCaptionHost } from './caption-host.mjs'
|
||||
import http from 'node:http'
|
||||
import net from 'node:net'
|
||||
import { execFile, spawn } from 'node:child_process'
|
||||
@@ -147,7 +148,7 @@ let proxyTarget = 0
|
||||
|
||||
function ensureProxyListening() {
|
||||
if (proxyServer) return
|
||||
proxyServer = createGpuProxy({ target: () => proxyTarget, reservation: gpuReservation, authorized, markWork, externalBusy: () => yue2.busy() || upscale.busy() })
|
||||
proxyServer = createGpuProxy({ target: () => proxyTarget, reservation: gpuReservation, authorized, markWork, externalBusy: () => yue2.busy() || upscale.busy() || caption.busy() })
|
||||
proxyServer.on('error', (error) => {
|
||||
console.log(JSON.stringify({ src: 'comfy-host-agent', event: 'proxy-error', error: String(error.message || error) }))
|
||||
})
|
||||
@@ -819,7 +820,7 @@ function purgeDesktopFiles(body) {
|
||||
}
|
||||
|
||||
const gpuReservation = createGpuReservation({ idle: async () => {
|
||||
if (yue2.busy() || upscale.busy()) return false
|
||||
if (yue2.busy() || upscale.busy() || caption.busy()) return false
|
||||
if ((await trainingLock()).busy) return false
|
||||
const healthy = await syncProxy()
|
||||
if (healthy) {
|
||||
@@ -830,14 +831,15 @@ const gpuReservation = createGpuReservation({ idle: async () => {
|
||||
return !(await processUp()) && !(await pythonMainUp().catch(() => true))
|
||||
} })
|
||||
|
||||
const yue2 = createYue2Host({
|
||||
leaseValid: lease => gpuReservation.isOwner(lease),
|
||||
prepare: async () => {
|
||||
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; YuE2 cannot start.')
|
||||
if (!queue.ok || queue.running || queue.pending) throw new Error(`Comfy is busy; ${label} cannot start.`)
|
||||
}
|
||||
// Stop Comfy and refuse to launch while its Python still owns VRAM.
|
||||
if (healthy || await processUp() || await pythonMainUp().catch(() => false)) {
|
||||
@@ -846,13 +848,50 @@ const yue2 = createYue2Host({
|
||||
}
|
||||
if (await pythonMainUp()) throw new Error('Comfy has not stopped; retry after the GPU is free.')
|
||||
}
|
||||
})
|
||||
|
||||
const upscale = createUpscaleHost({ leaseValid: token => gpuReservation.isOwner(token) })
|
||||
|
||||
const yue2 = createYue2Host({
|
||||
leaseValid: lease => gpuReservation.isOwner(lease),
|
||||
prepare: async () => prepareExclusiveGpu('YuE2')
|
||||
})
|
||||
|
||||
const caption = createCaptionHost({
|
||||
leaseValid: lease => gpuReservation.isOwner(lease),
|
||||
prepare: async () => prepareExclusiveGpu('caption')
|
||||
})
|
||||
|
||||
async function handleControl(req, res) {
|
||||
if (!authorized(req)) return json(res, 401, { ok: false, error: 'unauthorized' })
|
||||
const url = new URL(req.url || '/', 'http://localhost')
|
||||
if (url.pathname.startsWith('/caption')) {
|
||||
const match = url.pathname.match(/^\/caption\/jobs\/([a-zA-Z0-9-]{12,80})(\/input|\/cancel)?$/)
|
||||
if (req.method === 'GET' && url.pathname === '/caption/status') {
|
||||
return json(res, 200, { configured: caption.configured(), busy: caption.busy(), backend: 'llama.cpp', modelsDir: caption.modelsDir() })
|
||||
}
|
||||
// Sync one-shot: POST /caption { imagePath, style } → load → caption → unload → { text }
|
||||
if (req.method === 'POST' && url.pathname === '/caption') {
|
||||
try {
|
||||
const body = await readJson(req)
|
||||
return json(res, 200, await caption.captionOnce(body, String(req.headers['x-aigen-gpu-lease'] || '')))
|
||||
} catch (error) {
|
||||
return json(res, error.statusCode || 400, { error: error.message || 'Caption failed' })
|
||||
}
|
||||
}
|
||||
if (req.method === 'POST' && url.pathname === '/caption/jobs') {
|
||||
return json(res, 200, await caption.start(await readJson(req), String(req.headers['x-aigen-gpu-lease'] || '')))
|
||||
}
|
||||
if (match && req.method === 'PUT' && match[2] === '/input') {
|
||||
await caption.upload(match[1], req)
|
||||
return json(res, 200, { uploaded: true })
|
||||
}
|
||||
if (match && req.method === 'POST' && match[2] === '/cancel') return json(res, 200, await caption.cancel(match[1]))
|
||||
if (match && req.method === 'GET' && !match[2]) {
|
||||
const job = caption.read(match[1])
|
||||
return json(res, job ? 200 : 404, job || { error: 'Caption job not found' })
|
||||
}
|
||||
return json(res, 404, { error: 'Unknown caption endpoint' })
|
||||
}
|
||||
if (url.pathname.startsWith('/upscale/')) {
|
||||
const match = url.pathname.match(/^\/upscale\/jobs\/([a-zA-Z0-9-]{12,80})(\/input|\/video|\/cancel)?$/)
|
||||
if (req.method === 'POST' && url.pathname === '/upscale/jobs') return json(res, 200, upscale.start(await readJson(req), String(req.headers['x-aigen-gpu-lease'] || '')))
|
||||
@@ -909,12 +948,14 @@ async function handleControl(req, res) {
|
||||
gpu: gpuReservation.availability(),
|
||||
training: { busy: lastTraining.busy },
|
||||
yue2: { busy: yue2.busy(), configured: yue2.configured() },
|
||||
caption: { busy: caption.busy(), configured: caption.configured() },
|
||||
upscale: { busy: upscale.busy(), engine: 'realesrgan-rife', local: true }
|
||||
})
|
||||
}
|
||||
if (req.method === 'POST' && url.pathname === '/start') {
|
||||
if (upscale.busy()) return json(res, 409, { message: 'Local upscale is using the GPU.' })
|
||||
if (yue2.busy()) return json(res, 409, { message: 'YuE2 is using the GPU.' })
|
||||
if (caption.busy()) return json(res, 409, { message: 'Caption is using the GPU.' })
|
||||
const training = await trainingLock()
|
||||
if (training.busy) {
|
||||
return json(res, 409, {
|
||||
@@ -1016,7 +1057,7 @@ const server = http.createServer(async (req, res) => {
|
||||
} else await handleControl(req, res)
|
||||
} catch (error) {
|
||||
req.resume()
|
||||
if ((String(req.url || '').startsWith('/yue2/') || String(req.url || '').startsWith('/upscale/')) && !res.headersSent) return json(res, error.statusCode || 400, { error: error.message || 'Music host request failed' })
|
||||
if ((String(req.url || '').startsWith('/yue2/') || String(req.url || '').startsWith('/upscale/') || String(req.url || '').startsWith('/caption')) && !res.headersSent) return json(res, error.statusCode || 400, { error: error.message || 'Host request failed' })
|
||||
if (!res.headersSent) json(res, error.statusCode || 400, { ok: false, message: error.statusCode === 409 ? 'GPU is in use. Waiting for availability.' : 'GPU coordination request failed.' })
|
||||
}
|
||||
})
|
||||
|
||||
@@ -0,0 +1,129 @@
|
||||
# Download Qwen2.5-VL NSFW Caption V4 GGUF (Q5_K_M + mmproj only) for Studio 2 caption jobs.
|
||||
# Does NOT pull the full safetensors repo (~16.6GB).
|
||||
#
|
||||
# Example:
|
||||
# powershell -ExecutionPolicy Bypass -File scripts\setup-caption.ps1
|
||||
# powershell -File scripts\setup-caption.ps1 -SkipDownload
|
||||
|
||||
[CmdletBinding()]
|
||||
param(
|
||||
[string]$ModelsRoot = '',
|
||||
[string]$LlamaServer = '',
|
||||
[switch]$SkipDownload,
|
||||
[switch]$Force
|
||||
)
|
||||
|
||||
$ErrorActionPreference = 'Stop'
|
||||
$ProgressPreference = 'Continue'
|
||||
|
||||
function Write-Step([string]$Message) {
|
||||
Write-Host ""
|
||||
Write-Host "==> $Message" -ForegroundColor Cyan
|
||||
}
|
||||
|
||||
function Write-Ok([string]$Message) {
|
||||
Write-Host " $Message" -ForegroundColor Green
|
||||
}
|
||||
|
||||
function Write-Warn([string]$Message) {
|
||||
Write-Host " warn $Message" -ForegroundColor Yellow
|
||||
}
|
||||
|
||||
function Resolve-ModelsRoot([string]$Hint) {
|
||||
$shared = Join-Path $env:LOCALAPPDATA 'Comfy-Desktop\ComfyUI-Shared\models'
|
||||
$candidates = @(
|
||||
$Hint,
|
||||
$env:COMFY_MODELS_ROOT,
|
||||
$shared
|
||||
) | Where-Object { $_ } | ForEach-Object { $_.TrimEnd('\', '/') }
|
||||
|
||||
foreach ($path in $candidates) {
|
||||
if (Test-Path $path) { return (Resolve-Path $path).Path }
|
||||
}
|
||||
New-Item -ItemType Directory -Force -Path $shared | Out-Null
|
||||
return (Resolve-Path $shared).Path
|
||||
}
|
||||
|
||||
function Find-HfCli {
|
||||
foreach ($name in @('hf', 'huggingface-cli')) {
|
||||
$cmd = Get-Command $name -ErrorAction SilentlyContinue
|
||||
if ($cmd) { return @{ Name = $name; Source = $cmd.Source } }
|
||||
}
|
||||
throw "Hugging Face CLI not found (hf / huggingface-cli). Install with: pip install -U huggingface_hub"
|
||||
}
|
||||
|
||||
function Ensure-Dir([string]$Path) {
|
||||
if (-not (Test-Path $Path)) {
|
||||
New-Item -ItemType Directory -Force -Path $Path | Out-Null
|
||||
}
|
||||
}
|
||||
|
||||
Write-Step 'Resolve models root'
|
||||
$modelsRoot = Resolve-ModelsRoot $ModelsRoot
|
||||
$captionDir = Join-Path $modelsRoot 'caption\qwen25vl-7b-nsfw-v4'
|
||||
Ensure-Dir $captionDir
|
||||
Write-Ok "Models root: $modelsRoot"
|
||||
Write-Ok "Caption dir: $captionDir"
|
||||
|
||||
$model = 'Qwen2.5-VL-7B-NSFW-Caption-V4.Q5_K_M.gguf'
|
||||
$mmproj = 'Qwen2.5-VL-7B-NSFW-Caption-V4.mmproj-f16.gguf'
|
||||
$modelPath = Join-Path $captionDir $model
|
||||
$mmprojPath = Join-Path $captionDir $mmproj
|
||||
|
||||
if (-not $SkipDownload) {
|
||||
Write-Step 'Download GGUF + mmproj only (~7GB)'
|
||||
$needModel = $Force -or -not (Test-Path $modelPath)
|
||||
$needMmproj = $Force -or -not (Test-Path $mmprojPath)
|
||||
if (-not $needModel -and -not $needMmproj) {
|
||||
Write-Ok 'Both weight files already present.'
|
||||
} else {
|
||||
$hf = Find-HfCli
|
||||
if ($hf.Name -eq 'hf') {
|
||||
& $hf.Source download mradermacher/Qwen2.5-VL-7B-NSFW-Caption-V4-GGUF `
|
||||
--include $model `
|
||||
--include $mmproj `
|
||||
--local-dir $captionDir
|
||||
} else {
|
||||
& $hf.Source download mradermacher/Qwen2.5-VL-7B-NSFW-Caption-V4-GGUF `
|
||||
--include $model `
|
||||
--include $mmproj `
|
||||
--local-dir $captionDir
|
||||
}
|
||||
if (-not (Test-Path $modelPath)) { throw "Missing $model after download." }
|
||||
if (-not (Test-Path $mmprojPath)) { throw "Missing $mmproj after download." }
|
||||
Write-Ok 'Download complete.'
|
||||
}
|
||||
} else {
|
||||
Write-Warn 'SkipDownload set — not fetching weights.'
|
||||
}
|
||||
|
||||
Write-Step 'Locate llama-server'
|
||||
$server = $LlamaServer
|
||||
if (-not $server) { $server = $env:CAPTION_LLAMA_SERVER }
|
||||
if (-not $server) {
|
||||
$cmd = Get-Command llama-server -ErrorAction SilentlyContinue
|
||||
if ($cmd) { $server = $cmd.Source }
|
||||
}
|
||||
if ($server) {
|
||||
Write-Ok "llama-server: $server"
|
||||
} else {
|
||||
Write-Warn 'llama-server not on PATH. Install with: winget install ggml.llamacpp'
|
||||
Write-Warn 'Or set CAPTION_LLAMA_SERVER to a CUDA build of llama-server.exe'
|
||||
}
|
||||
|
||||
$ready = @{
|
||||
model = $modelPath
|
||||
mmproj = $mmprojPath
|
||||
modelsDir = $captionDir
|
||||
llamaServer = $server
|
||||
backend = 'llama.cpp'
|
||||
updatedAt = (Get-Date).ToString('o')
|
||||
}
|
||||
$readyPath = Join-Path $captionDir 'aigen-ready.json'
|
||||
$ready | ConvertTo-Json | Set-Content -Path $readyPath -Encoding utf8
|
||||
Write-Ok "Wrote $readyPath"
|
||||
Write-Host ""
|
||||
Write-Host "Caption host looks for:" -ForegroundColor Cyan
|
||||
Write-Host " $modelPath"
|
||||
Write-Host " $mmprojPath"
|
||||
Write-Host "Restart the Comfy host agent after install so /caption/status reports configured=true."
|
||||
@@ -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,
|
||||
|
||||
@@ -0,0 +1,74 @@
|
||||
import { validateCaptionRequest } from '~/shared/studio2/caption.mjs'
|
||||
import { saveRecord, readRecord } from '../../utils/studio2/store'
|
||||
import { addStudioJob, kickStudioQueue, type StudioJobPayload } from '../../utils/studioQueue'
|
||||
import { getStill } from '../../utils/library'
|
||||
|
||||
export default defineEventHandler(async (event) => {
|
||||
const { owner } = assertLibraryOwner(event)
|
||||
const raw = await readBody(event)
|
||||
let request: ReturnType<typeof validateCaptionRequest>
|
||||
try {
|
||||
request = validateCaptionRequest(raw)
|
||||
} catch (e: any) {
|
||||
throw createError({ statusCode: e.statusCode || 400, statusMessage: e.message })
|
||||
}
|
||||
assertFolderAccess(event, request.folderId)
|
||||
let stillId = request.stillId
|
||||
if (!stillId && request.sourcePath) {
|
||||
// sourcePath may be a still id when the client already resolved a library asset.
|
||||
stillId = String(request.sourcePath)
|
||||
}
|
||||
const still = getStill(owner, stillId)
|
||||
assertFolderAccess(event, still.folderId)
|
||||
request.stillId = still.id
|
||||
request.imageAId = still.id
|
||||
request.sourcePath = still.id
|
||||
|
||||
const id = crypto.randomUUID()
|
||||
const record = {
|
||||
id,
|
||||
owner,
|
||||
request,
|
||||
state: 'waiting',
|
||||
queuedAt: Date.now(),
|
||||
familyId: crypto.randomUUID(),
|
||||
outputs: [],
|
||||
resultText: '',
|
||||
purgeResult: 'Not yet saved'
|
||||
}
|
||||
saveRecord(record)
|
||||
const row = await addStudioJob({
|
||||
ownerKey: owner,
|
||||
kind: 'caption',
|
||||
familyId: record.familyId,
|
||||
payload: {
|
||||
studio2Id: id,
|
||||
prompt: request.compiledPrompt,
|
||||
name: `Describe · ${request.captionStyle}`,
|
||||
folderId: request.folderId,
|
||||
stillId: still.id,
|
||||
captionStyle: request.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
|
||||
})
|
||||
saveRecord({ ...readRecord(owner, id), queueId: row.id })
|
||||
kickStudioQueue()
|
||||
return { id, captionStyle: request.captionStyle, stillId: still.id }
|
||||
})
|
||||
@@ -0,0 +1,3 @@
|
||||
import { resumeCaptionJobs } from '~/server/utils/caption'
|
||||
|
||||
export default defineNitroPlugin(() => { resumeCaptionJobs() })
|
||||
@@ -0,0 +1,268 @@
|
||||
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/studio2/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) {
|
||||
if (!jobId) return false
|
||||
return 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 } : {}),
|
||||
signal: AbortSignal.timeout(method === 'PUT' ? 120_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<string, unknown>) {
|
||||
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 sidecar */ }
|
||||
}
|
||||
|
||||
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
|
||||
studio2Id?: 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 })
|
||||
if (record.studio2Id) {
|
||||
const { readRecord, saveRecord } = await import('./studio2/store')
|
||||
try {
|
||||
const r = readRecord(record.owner, record.studio2Id)
|
||||
r.state = job.status === 'cancelled' ? 'cancelled' : 'failed'
|
||||
r.error = job.error
|
||||
r.finishedAt = Date.now()
|
||||
saveRecord(r)
|
||||
} catch { /* record may already be gone */ }
|
||||
}
|
||||
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 as Job & { resultText?: string }).resultText = text
|
||||
if (record.studio2Id) {
|
||||
const { readRecord, saveRecord } = await import('./studio2/store')
|
||||
const r = readRecord(record.owner, record.studio2Id)
|
||||
r.state = 'complete'
|
||||
r.resultText = text
|
||||
r.finishedAt = Date.now()
|
||||
r.savedAt = Date.now()
|
||||
saveRecord(r)
|
||||
}
|
||||
writeSidecar(record.owner, record.stillId, text)
|
||||
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
|
||||
studio2Id?: 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,
|
||||
studio2Id: params.studio2Id,
|
||||
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; never caption twice on an HTTP 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 })
|
||||
if (params.studio2Id) {
|
||||
try {
|
||||
const { readRecord, saveRecord } = await import('./studio2/store')
|
||||
const r = readRecord(params.ownerKey, params.studio2Id)
|
||||
r.state = 'failed'
|
||||
r.error = job.error
|
||||
r.finishedAt = Date.now()
|
||||
saveRecord(r)
|
||||
} catch { /* ignore */ }
|
||||
}
|
||||
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'
|
||||
job.message = 'Reconnecting to caption host'
|
||||
void watch(job, record)
|
||||
} catch { /* Preserve invalid records for diagnosis. */ }
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -40,6 +40,7 @@ export interface StudioJobPayload {
|
||||
useIdentityRefs: boolean
|
||||
stillId?: string
|
||||
stillFilename?: string
|
||||
captionStyle?: string
|
||||
hideThumbnail: boolean
|
||||
hideInput?: boolean
|
||||
folderLocked?: boolean
|
||||
@@ -224,6 +225,7 @@ export function listStudioJobs(owner: string) {
|
||||
export function studioJobKind(job: Pick<StudioJob, 'kind'> | { kind?: string }) {
|
||||
if (job.kind === 'edit') return 'edit'
|
||||
if (job.kind === 'music') return 'music'
|
||||
if (job.kind === 'caption') return 'caption'
|
||||
return 'video'
|
||||
}
|
||||
|
||||
@@ -288,7 +290,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')
|
||||
@@ -308,12 +310,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')
|
||||
@@ -353,7 +355,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
|
||||
@@ -373,7 +375,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') {
|
||||
@@ -467,8 +469,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)
|
||||
@@ -604,6 +609,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) {
|
||||
@@ -676,6 +685,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) {
|
||||
@@ -887,6 +901,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 => (
|
||||
@@ -1395,6 +1413,42 @@ 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,
|
||||
studio2Id: item.payload.studio2Id,
|
||||
name: item.payload.name
|
||||
})
|
||||
await markStudioLive(item.ownerKey, item.id, live.id)
|
||||
if (item.payload.studio2Id) {
|
||||
const { readRecord, saveRecord } = await import('./studio2/store')
|
||||
const record = readRecord(item.ownerKey, item.payload.studio2Id)
|
||||
record.state = 'rendering'
|
||||
record.liveId = live.id
|
||||
record.startedAt = Date.now()
|
||||
record.updatedAt = Date.now()
|
||||
saveRecord(record)
|
||||
}
|
||||
} 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 {
|
||||
@@ -1437,6 +1491,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')
|
||||
@@ -1701,7 +1759,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) => {
|
||||
|
||||
@@ -0,0 +1,66 @@
|
||||
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 validateCaptionRequest(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.imageAId || '')
|
||||
const sourcePath = String(body.sourcePath || '')
|
||||
if (!stillId && !sourcePath) fail('Choose a still to describe.')
|
||||
return {
|
||||
mode: 'caption',
|
||||
folderId,
|
||||
stillId,
|
||||
sourcePath,
|
||||
captionStyle: style,
|
||||
promptSections: { action: `Describe · ${style}` },
|
||||
compiledPrompt: captionStylePrompt(style),
|
||||
engine: 'caption',
|
||||
imageStyles: { positive: [], negative: [] },
|
||||
hideOutput: false,
|
||||
lockFace: false,
|
||||
lockOutfit: false,
|
||||
hideInput: false,
|
||||
identityStillId: '',
|
||||
imageAId: stillId,
|
||||
imageBId: '',
|
||||
maskId: '',
|
||||
startClipId: '',
|
||||
endStillId: '',
|
||||
guides: [],
|
||||
batch: [],
|
||||
settings: {
|
||||
aspect: 'auto',
|
||||
fps: 24,
|
||||
width: 0,
|
||||
height: 0,
|
||||
steps: 1,
|
||||
cfg: 0,
|
||||
seedMode: 'random',
|
||||
seed: 0,
|
||||
duration: 0,
|
||||
turbo: true,
|
||||
loraStack: []
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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','waking','submitting','rendering','saving'].includes(j.state) || j.state==='queued' || j.state==='held')
|
||||
if(caption)return {label:caption.state==='waiting'?'Queued':'Describing',tone:'busy',detail:caption.request?.captionStyle || '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','submitting','rendering','saving'].includes(j.state))
|
||||
if(job)return {label:{waking:'Waking up',submitting:'Submitting',rendering:'Generating',saving:'Saving'}[job.state],tone:'busy',detail:job.request?.promptSections?.action || 'Generation in progress'}
|
||||
|
||||
@@ -0,0 +1,139 @@
|
||||
import test from 'node:test'
|
||||
import assert from 'node:assert/strict'
|
||||
import { EventEmitter } from 'node:events'
|
||||
import { PassThrough } from 'node:stream'
|
||||
import { mkdtempSync, writeFileSync, readFileSync, mkdirSync, existsSync } from 'node:fs'
|
||||
import { tmpdir } from 'node:os'
|
||||
import { join } from 'node:path'
|
||||
import { createCaptionHost, validateCaptionHostRequest, captionModelFiles } from '../scripts/caption-host.mjs'
|
||||
import { validateCaptionRequest, captionStylePrompt, CAPTION_STYLES, DELTA_LOCKED_LOOK } from '../shared/studio2/caption.mjs'
|
||||
|
||||
test('caption styles and delta lock text', () => {
|
||||
assert.deepEqual(CAPTION_STYLES, ['descriptive', 'klein_prompt', 'delta', 'tags'])
|
||||
assert.match(captionStylePrompt('delta'), new RegExp(DELTA_LOCKED_LOOK))
|
||||
assert.match(captionStylePrompt('klein_prompt'), /Flux Klein/)
|
||||
assert.match(captionStylePrompt('tags'), /comma-separated/)
|
||||
})
|
||||
|
||||
test('validateCaptionRequest requires folder + still', () => {
|
||||
assert.throws(() => validateCaptionRequest({}), /folder/)
|
||||
assert.throws(() => validateCaptionRequest({ folderId: 'f1' }), /still/)
|
||||
assert.throws(() => validateCaptionRequest({ folderId: 'f1', stillId: 's1', captionStyle: 'nope' }), /style/)
|
||||
const ok = validateCaptionRequest({ folderId: 'f1', stillId: 's1', captionStyle: 'descriptive' })
|
||||
assert.equal(ok.mode, 'caption')
|
||||
assert.equal(ok.captionStyle, 'descriptive')
|
||||
assert.equal(ok.lockFace, false)
|
||||
assert.equal(ok.lockOutfit, false)
|
||||
})
|
||||
|
||||
test('host validates job body', () => {
|
||||
assert.throws(() => validateCaptionHostRequest({ id: 'bad', style: 'descriptive' }))
|
||||
assert.deepEqual(
|
||||
validateCaptionHostRequest({ id: 'caption-job-123456', style: 'tags' }),
|
||||
{ id: 'caption-job-123456', style: 'tags', imagePath: '' }
|
||||
)
|
||||
})
|
||||
|
||||
function fixture() {
|
||||
const root = mkdtempSync(join(tmpdir(), 'aigen-caption-'))
|
||||
const models = join(root, 'models')
|
||||
mkdirSync(models)
|
||||
writeFileSync(join(models, captionModelFiles.model), 'model')
|
||||
writeFileSync(join(models, captionModelFiles.mmproj), 'mmproj')
|
||||
const llama = join(root, 'llama-server.exe')
|
||||
writeFileSync(llama, '')
|
||||
const dataDir = join(root, 'jobs')
|
||||
const child = new EventEmitter()
|
||||
child.pid = 4242
|
||||
child.stdout = new PassThrough()
|
||||
child.stderr = new PassThrough()
|
||||
child.kill = () => { setImmediate(() => child.emit('close', 0)); return true }
|
||||
child.exitCode = null
|
||||
const calls = []
|
||||
let fetchCount = 0
|
||||
const fetchImpl = async (url, init) => {
|
||||
fetchCount++
|
||||
if (String(url).endsWith('/health')) return { ok: true }
|
||||
if (String(url).includes('/v1/chat/completions')) {
|
||||
assert.match(JSON.stringify(JSON.parse(init.body)), /Describe this image|Flux Klein|locked baseline|comma-separated/)
|
||||
return {
|
||||
ok: true,
|
||||
json: async () => ({ choices: [{ message: { content: 'a warm caption' } }] })
|
||||
}
|
||||
}
|
||||
throw new Error(`unexpected fetch ${url}`)
|
||||
}
|
||||
const host = createCaptionHost({
|
||||
modelsDir: models,
|
||||
llamaServer: llama,
|
||||
dataDir,
|
||||
prepare: async () => {},
|
||||
leaseValid: () => true,
|
||||
spawnProcess: (...args) => { calls.push(args); setImmediate(() => {}); return child },
|
||||
fetchImpl
|
||||
})
|
||||
return { host, child, calls, root, dataDir, fetchCount: () => fetchCount }
|
||||
}
|
||||
|
||||
test('host loads llama-server, captions once, and unloads', async () => {
|
||||
const f = fixture()
|
||||
const jobDir = join(f.dataDir, 'caption-job-123456')
|
||||
mkdirSync(jobDir, { recursive: true })
|
||||
writeFileSync(join(jobDir, 'input.png'), Buffer.alloc(64, 1))
|
||||
const started = f.host.start({ id: 'caption-job-123456', style: 'descriptive' }, 'lease')
|
||||
// Let the host reach waitForServer / chat completion.
|
||||
await new Promise(r => setTimeout(r, 50))
|
||||
f.child.emit('spawn')
|
||||
const result = await started
|
||||
assert.equal(result.status, 'complete')
|
||||
assert.equal(result.text, 'a warm caption')
|
||||
assert.equal(f.calls.length, 1)
|
||||
const [, args] = f.calls[0]
|
||||
assert.ok(args.includes('--mmproj'))
|
||||
assert.ok(args.includes('-m'))
|
||||
assert.equal(f.host.read('caption-job-123456').text, 'a warm caption')
|
||||
// busy may briefly remain true during the post-unload hold window
|
||||
assert.ok(typeof f.host.busy() === 'boolean')
|
||||
})
|
||||
|
||||
test('host rejects a second job while busy', async () => {
|
||||
const f = fixture()
|
||||
const dataDir = join(f.root, 'jobs2')
|
||||
let resolveHealth
|
||||
const healthGate = new Promise(resolve => { resolveHealth = resolve })
|
||||
const hung = createCaptionHost({
|
||||
modelsDir: join(f.root, 'models'),
|
||||
llamaServer: join(f.root, 'llama-server.exe'),
|
||||
dataDir,
|
||||
prepare: async () => {},
|
||||
leaseValid: () => true,
|
||||
spawnProcess: () => {
|
||||
const child = new EventEmitter()
|
||||
child.pid = 99
|
||||
child.stdout = new PassThrough()
|
||||
child.stderr = new PassThrough()
|
||||
child.kill = () => { setImmediate(() => child.emit('close', 1)); return true }
|
||||
child.exitCode = null
|
||||
return child
|
||||
},
|
||||
fetchImpl: async (url) => {
|
||||
if (String(url).endsWith('/health')) {
|
||||
await healthGate
|
||||
return { ok: true }
|
||||
}
|
||||
return { ok: true, json: async () => ({ choices: [{ message: { content: 'x' } }] }) }
|
||||
}
|
||||
})
|
||||
mkdirSync(join(dataDir, 'caption-job-aaaaaa'), { recursive: true })
|
||||
writeFileSync(join(dataDir, 'caption-job-aaaaaa', 'input.png'), Buffer.alloc(64, 1))
|
||||
const first = hung.start({ id: 'caption-job-aaaaaa', style: 'tags' }, 'lease')
|
||||
await new Promise(r => setTimeout(r, 20))
|
||||
await assert.rejects(hung.start({ id: 'caption-job-bbbbbb', style: 'tags' }, 'lease'), /already running/)
|
||||
await hung.cancel('caption-job-aaaaaa')
|
||||
resolveHealth()
|
||||
await assert.rejects(first)
|
||||
})
|
||||
|
||||
test('POST /caption sync shape is documented on host module', () => {
|
||||
assert.equal(typeof createCaptionHost({ prepare: async () => {}, leaseValid: () => true }).captionOnce, 'function')
|
||||
})
|
||||
Reference in New Issue
Block a user