aigen/server/utils/comfy.ts

function comfyBase() {
  const config = useRuntimeConfig()
  const host = (config.comfyHost || process.env.COMFY_HOST || '').replace(/\/$/, '')
  if (!host) {
    throw createError({ statusCode: 500, statusMessage: 'COMFY_HOST is not configured' })
  }
  return host
}

export function getComfyHost() {
  return comfyBase()
}

export function comfyWsUrl(clientId: string) {
  return `${comfyBase().replace(/^http/, 'ws')}/ws?clientId=${encodeURIComponent(clientId)}`
}

export async function comfyFetch(path: string, init?: RequestInit) {
  const url = `${comfyBase()}${path}`
  try {
    return await fetch(url, init)
  } catch (error) {
    throw createError({
      statusCode: 502,
      statusMessage: `ComfyUI host unreachable (${comfyBase()})`,
      data: { cause: error instanceof Error ? error.message : String(error) }
    })
  }
}

export async function uploadImage(file: { filename: string; data: Buffer; type?: string }) {
  const body = new FormData()
  const blob = new Blob([new Uint8Array(file.data)], { type: file.type || 'application/octet-stream' })
  body.append('image', blob, file.filename)
  body.append('overwrite', 'true')
  body.append('type', 'input')
  const res = await comfyFetch('/upload/image', { method: 'POST', body })
  if (!res.ok) {
    throw createError({ statusCode: 502, statusMessage: `Image upload failed (${res.status})` })
  }
  return (await res.json()) as { name: string; subfolder?: string; type?: string }
}

export async function queuePrompt(graph: unknown, clientId: string) {
  const res = await comfyFetch('/prompt', {
    method: 'POST',
    headers: { 'Content-Type': 'application/json' },
    body: JSON.stringify({ prompt: graph, client_id: clientId })
  })
  const payload = await res.json().catch(() => ({}))
  if (!res.ok) {
    const message = (payload as { error?: { message?: string } }).error?.message
      || (payload as { node_errors?: unknown }).node_errors
      || `Queue failed (${res.status})`
    throw createError({ statusCode: 502, statusMessage: String(message), data: payload })
  }
  return payload as { prompt_id: string; number?: number }
}

export async function interruptComfy() {
  const res = await comfyFetch('/interrupt', { method: 'POST' })
  return res.ok
}

export async function fetchHistory(promptId: string) {
  const res = await comfyFetch(`/history/${encodeURIComponent(promptId)}`)
  if (!res.ok) return null
  return (await res.json()) as Record<string, unknown>
}

function isVideoFile(item: { filename?: string; format?: string } | null | undefined) {
  if (!item) return false
  const name = String(item.filename || '').toLowerCase()
  const format = String(item.format || '').toLowerCase()
  return name.endsWith('.mp4') || name.endsWith('.webm') || format.includes('video') || format.includes('mp4')
}

function normalizeVideo(item: { filename?: string; subfolder?: string; type?: string }) {
  return {
    filename: String(item.filename),
    subfolder: String(item.subfolder || ''),
    type: String(item.type || 'output')
  }
}

function findVideo(value: unknown, depth = 0): { filename: string; subfolder: string; type: string } | null {
  if (!value || typeof value !== 'object' || depth > 8) return null
  if (Array.isArray(value)) {
    const files = value.filter((item): item is { filename?: string; format?: string; subfolder?: string; type?: string } => Boolean(item && typeof item === 'object'))
    const match = files.find(item => String(item.filename || '').toLowerCase().endsWith('.mp4'))
      || files.find(isVideoFile)
    if (match?.filename) return normalizeVideo(match)
    for (const item of value) {
      const nested = findVideo(item, depth + 1)
      if (nested) return nested
    }
    return null
  }
  const record = value as { filename?: string; format?: string; subfolder?: string; type?: string }
  if (record.filename && isVideoFile(record)) return normalizeVideo(record)
  for (const nested of Object.values(value as Record<string, unknown>)) {
    const found = findVideo(nested, depth + 1)
    if (found) return found
  }
  return null
}

export function extractVideo(history: Record<string, unknown> | null, promptId: string) {
  if (!history) return null
  const entry = (history[promptId] || Object.values(history)[0]) as { outputs?: Record<string, unknown> } | undefined
  return findVideo(entry?.outputs || {})
}

export function inspectHistory(history: Record<string, unknown> | null, promptId: string) {
  const video = extractVideo(history, promptId)
  if (!history) return { video: null, completed: false, error: null as string | null }
  const entry = (history[promptId] || Object.values(history)[0]) as {
    status?: { status_str?: string; completed?: boolean; messages?: unknown[] }
  } | undefined
  const status = entry?.status
  const error = status?.status_str === 'error'
    ? 'ComfyUI reported an execution error'
    : null
  const completed = Boolean(video) || status?.completed === true || status?.status_str === 'success'
  return { video, completed, error }
}

export async function probeComfy() {
  try {
    const res = await fetch(`${comfyBase()}/system_stats`, { signal: AbortSignal.timeout(2500) })
    if (!res.ok) return { ok: false, host: comfyBase() }
    const stats = await res.json()
    return { ok: true, host: comfyBase(), stats }
  } catch {
    return { ok: false, host: comfyBase() }
  }
}