aigen/scripts/caption-host.mjs

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 }